diff --git a/lib/python3.12/site-packages/wandb/agents/__init__.py b/lib/python3.12/site-packages/wandb/agents/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/wandb/agents/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/agents/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e6d33588b53908afb09201c90439407d586d7fce Binary files /dev/null and b/lib/python3.12/site-packages/wandb/agents/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/agents/__pycache__/pyagent.cpython-312.pyc b/lib/python3.12/site-packages/wandb/agents/__pycache__/pyagent.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9cabc4c0600b0ce836aee0e60b96c67b5e52d4be Binary files /dev/null and b/lib/python3.12/site-packages/wandb/agents/__pycache__/pyagent.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/agents/pyagent.py b/lib/python3.12/site-packages/wandb/agents/pyagent.py new file mode 100644 index 0000000000000000000000000000000000000000..d43b1b8e63b29a0b2b565bd98c0c0bef55d8830b --- /dev/null +++ b/lib/python3.12/site-packages/wandb/agents/pyagent.py @@ -0,0 +1,357 @@ +"""Agent - Agent object. + +Manage wandb agent. + +""" + +import ctypes +import logging +import os +import queue +import socket +import threading +import time +import traceback + +import wandb +from wandb.apis import InternalApi +from wandb.sdk.launch.sweeps import utils as sweep_utils + +logger = logging.getLogger(__name__) + + +def _terminate_thread(thread): + if not thread.is_alive(): + return + if hasattr(thread, "_terminated"): + return + thread._terminated = True + tid = getattr(thread, "_thread_id", None) + if tid is None: + for k, v in threading._active.items(): + if v is thread: + tid = k + if tid is None: + # This should never happen + return + logger.debug(f"Terminating thread: {tid}") + res = ctypes.pythonapi.PyThreadState_SetAsyncExc( + ctypes.c_long(tid), ctypes.py_object(Exception) + ) + if res == 0: + # This should never happen + return + elif res != 1: + # Revert + logger.debug(f"Termination failed for thread {tid}") + ctypes.pythonapi.PyThreadState_SetAsyncExc(ctypes.c_long(tid), None) + + +class Job: + def __init__(self, command): + self.command = command + job_type = command.get("type") + self.type = job_type + self.run_id = command.get("run_id") + self.config = command.get("args") + + def __repr__(self): + if self.type == "run": + return f"Job({self.run_id},{self.config})" + elif self.type == "stop": + return f"stop({self.run_id})" + else: + return "exit" + + +class RunStatus: + QUEUED = "QUEUED" + RUNNING = "RUNNING" + STOPPED = "STOPPED" + ERRORED = "ERRORED" + DONE = "DONE" + + +class Agent: + FLAPPING_MAX_SECONDS = 60 + FLAPPING_MAX_FAILURES = 3 + MAX_INITIAL_FAILURES = 5 + + def __init__( + self, sweep_id=None, project=None, entity=None, function=None, count=None + ): + self._sweep_path = sweep_id + self._sweep_id = None + self._project = project + self._entity = entity + self._function = function + self._count = count + # glob_config = os.path.expanduser('~/.config/wandb/settings') + # loc_config = 'wandb/settings' + # files = (glob_config, loc_config) + self._api = InternalApi() + self._agent_id = None + self._max_initial_failures = wandb.env.get_agent_max_initial_failures( + self.MAX_INITIAL_FAILURES + ) + # if the directory to log to is not set, set it + if os.environ.get(wandb.env.DIR) is None: + os.environ[wandb.env.DIR] = os.path.abspath(os.getcwd()) + + def _init(self): + # These are not in constructor so that Agent instance can be rerun + self._run_threads = {} + self._run_status = {} + self._queue = queue.Queue() + self._exit_flag = False + self._exceptions = {} + self._start_time = time.time() + + def _register(self): + logger.debug("Agent._register()") + agent = self._api.register_agent(socket.gethostname(), sweep_id=self._sweep_id) + self._agent_id = agent["id"] + logger.debug(f"agent_id = {self._agent_id}") + + def _setup(self): + logger.debug("Agent._setup()") + self._init() + parts = dict(entity=self._entity, project=self._project, name=self._sweep_path) + err = sweep_utils.parse_sweep_id(parts) + if err: + wandb.termerror(err) + return + entity = parts.get("entity") or self._entity + project = parts.get("project") or self._project + sweep_id = parts.get("name") or self._sweep_id + if sweep_id: + os.environ[wandb.env.SWEEP_ID] = sweep_id + if entity: + wandb.env.set_entity(entity) + if project: + wandb.env.set_project(project) + if sweep_id: + self._sweep_id = sweep_id + self._register() + + def _stop_run(self, run_id): + logger.debug(f"Stopping run {run_id}.") + self._run_status[run_id] = RunStatus.STOPPED + thread = self._run_threads.get(run_id) + if thread: + _terminate_thread(thread) + + def _stop_all_runs(self): + logger.debug("Stopping all runs.") + for run in list(self._run_threads.keys()): + self._stop_run(run) + + def _exit(self): + self._stop_all_runs() + self._exit_flag = True + # _terminate_thread(self._main_thread) + + def _heartbeat(self): + while True: + if self._exit_flag: + return + # if not self._main_thread.is_alive(): + # return + run_status = { + run: True + for run, status in self._run_status.items() + if status in (RunStatus.QUEUED, RunStatus.RUNNING) + } + commands = self._api.agent_heartbeat(self._agent_id, {}, run_status) + if commands: + job = Job(commands[0]) + logger.debug(f"Job received: {job}") + if job.type in ["run", "resume"]: + self._queue.put(job) + self._run_status[job.run_id] = RunStatus.QUEUED + elif job.type == "stop": + self._stop_run(job.run_id) + elif job.type == "exit": + self._exit() + return + time.sleep(5) + + def _run_jobs_from_queue(self): + global _INSTANCES + _INSTANCES += 1 + try: + waiting = False + count = 0 + while True: + if self._exit_flag: + return + try: + try: + job = self._queue.get(timeout=5) + if self._exit_flag: + logger.debug("Exiting main loop due to exit flag.") + wandb.termlog("Sweep Agent: Exiting.") + return + except queue.Empty: + if not waiting: + logger.debug("Paused.") + wandb.termlog("Sweep Agent: Waiting for job.") + waiting = True + time.sleep(5) + if self._exit_flag: + logger.debug("Exiting main loop due to exit flag.") + wandb.termlog("Sweep Agent: Exiting.") + return + continue + if waiting: + logger.debug("Resumed.") + wandb.termlog("Job received.") + waiting = False + count += 1 + run_id = job.run_id + if self._run_status[run_id] == RunStatus.STOPPED: + continue + logger.debug(f"Spawning new thread for run {run_id}.") + thread = threading.Thread(target=self._run_job, args=(job,)) + self._run_threads[run_id] = thread + thread.start() + self._run_status[run_id] = RunStatus.RUNNING + thread.join() + logger.debug(f"Thread joined for run {run_id}.") + if self._run_status[run_id] == RunStatus.RUNNING: + self._run_status[run_id] = RunStatus.DONE + elif self._run_status[run_id] == RunStatus.ERRORED: + exc = self._exceptions[run_id] + exc_type, exc_value, exc_traceback = ( + exc.__class__, + exc, + exc.__traceback__, + ) + exc_traceback_formatted = traceback.format_exception( + exc_type, exc_value, exc_traceback + ) + exc_repr = "".join(exc_traceback_formatted) + logger.error(f"Run {run_id} errored:\n{exc_repr}") + wandb.termerror(f"Run {run_id} errored:\n{exc_repr}") + if os.getenv(wandb.env.AGENT_DISABLE_FLAPPING) == "true": + self._exit_flag = True + return + elif ( + time.time() - self._start_time < self.FLAPPING_MAX_SECONDS + ) and (len(self._exceptions) >= self.FLAPPING_MAX_FAILURES): + msg = f"Detected {self.FLAPPING_MAX_FAILURES} failed runs in the first {self.FLAPPING_MAX_SECONDS} seconds, killing sweep." + logger.error(msg) + wandb.termerror(msg) + wandb.termlog( + "To disable this check set WANDB_AGENT_DISABLE_FLAPPING=true" + ) + self._exit_flag = True + return + if ( + self._max_initial_failures < len(self._exceptions) + and len(self._exceptions) >= count + ): + msg = f"Detected {self._max_initial_failures} failed runs in a row at start, killing sweep." + logger.error(msg) + wandb.termerror(msg) + wandb.termlog( + "To change this value set WANDB_AGENT_MAX_INITIAL_FAILURES=val" + ) + self._exit_flag = True + return + if self._count and self._count == count: + logger.debug("Exiting main loop because max count reached.") + self._exit_flag = True + return + except KeyboardInterrupt: + logger.debug("Ctrl + C detected. Stopping sweep.") + wandb.termlog("Ctrl + C detected. Stopping sweep.") + self._exit() + return + except Exception: + if self._exit_flag: + logger.debug("Exiting main loop due to exit flag.") + wandb.termlog("Sweep Agent: Killed.") + return + else: + raise + finally: + _INSTANCES -= 1 + + def _run_job(self, job): + try: + run_id = job.run_id + + config_file = os.path.join( + "wandb", "sweep-" + self._sweep_id, "config-" + run_id + ".yaml" + ) + os.environ[wandb.env.RUN_ID] = run_id + base_dir = os.environ.get(wandb.env.DIR, "") + sweep_param_path = os.path.join(base_dir, config_file) + os.environ[wandb.env.SWEEP_PARAM_PATH] = sweep_param_path + wandb.wandb_lib.config_util.save_config_file_from_dict( + sweep_param_path, job.config + ) + os.environ[wandb.env.SWEEP_ID] = self._sweep_id + wandb.teardown() + + wandb.termlog(f"Agent Starting Run: {run_id} with config:") + for k, v in job.config.items(): + wandb.termlog("\t{}: {}".format(k, v["value"])) + + self._function() + wandb.finish() + except KeyboardInterrupt: + raise + except Exception as e: + wandb.finish(exit_code=1) + if self._run_status[run_id] == RunStatus.RUNNING: + self._run_status[run_id] = RunStatus.ERRORED + self._exceptions[run_id] = e + finally: + # clean up the environment changes made + os.environ.pop(wandb.env.RUN_ID, None) + os.environ.pop(wandb.env.SWEEP_ID, None) + os.environ.pop(wandb.env.SWEEP_PARAM_PATH, None) + + def run(self): + logger.info( + f"Starting sweep agent: entity={self._entity}, project={self._project}, count={self._count}" + ) + self._setup() + # self._main_thread = threading.Thread(target=self._run_jobs_from_queue) + self._heartbeat_thread = threading.Thread(target=self._heartbeat) + self._heartbeat_thread.daemon = True + # self._main_thread.start() + self._heartbeat_thread.start() + # self._main_thread.join() + self._run_jobs_from_queue() + + +def pyagent(sweep_id, function, entity=None, project=None, count=None): + """Generic agent entrypoint, used for CLI or jupyter. + + Args: + sweep_id (dict): Sweep ID generated by CLI or sweep API + function (func, optional): A function to call instead of the "program" + entity (str, optional): W&B Entity + project (str, optional): W&B Project + count (int, optional): the number of trials to run. + """ + if not callable(function): + raise TypeError("function parameter must be callable!") + agent = Agent( + sweep_id, + function=function, + entity=entity, + project=project, + count=count, + ) + agent.run() + + +_INSTANCES = 0 + + +def is_running(): + return bool(_INSTANCES) diff --git a/lib/python3.12/site-packages/wandb/analytics/__init__.py b/lib/python3.12/site-packages/wandb/analytics/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ab265f84acb6b30517cca1a08e24ccdfd90c50e8 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/analytics/__init__.py @@ -0,0 +1,3 @@ +__all__ = ("Sentry",) + +from .sentry import Sentry diff --git a/lib/python3.12/site-packages/wandb/analytics/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/analytics/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ec5c40a4c842c6a2f37fa33b062a80d85748dbf0 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/analytics/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/analytics/__pycache__/sentry.cpython-312.pyc b/lib/python3.12/site-packages/wandb/analytics/__pycache__/sentry.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b7216cce5ae5682da04be8fbac8f4d5481fd03f1 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/analytics/__pycache__/sentry.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/analytics/sentry.py b/lib/python3.12/site-packages/wandb/analytics/sentry.py new file mode 100644 index 0000000000000000000000000000000000000000..f433af9817abb2c68fcac38de0c4bb94a803df62 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/analytics/sentry.py @@ -0,0 +1,262 @@ +from __future__ import annotations + +__all__ = ("Sentry",) + + +import atexit +import functools +import os +import pathlib +import sys +from types import TracebackType +from typing import TYPE_CHECKING, Any, Callable, Literal +from urllib.parse import quote + +import sentry_sdk # type: ignore +import sentry_sdk.scope # type: ignore +import sentry_sdk.utils # type: ignore +from typing_extensions import Never + +import wandb +import wandb.env +import wandb.util + +if TYPE_CHECKING: + import wandb.sdk.internal.settings_static + +SENTRY_DEFAULT_DSN = ( + "https://2592b1968ea94cca9b5ef5e348e094a7@o151352.ingest.sentry.io/4504800232407040" +) + +SessionStatus = Literal["ok", "exited", "crashed", "abnormal"] + + +def _safe_noop(func: Callable) -> Callable: + """Decorator to ensure that Sentry methods do nothing if disabled and don't raise.""" + + @functools.wraps(func) + def wrapper(self: type[Sentry], *args: Any, **kwargs: Any) -> Any: + if self._disabled: + return None + try: + return func(self, *args, **kwargs) + except Exception as e: + # do not call self.exception here to avoid infinite recursion + if func.__name__ != "exception": + self.exception(f"Error in {func.__name__}: {e}") + return None + + return wrapper + + +class Sentry: + _disabled: bool + + def __init__(self) -> None: + self._disabled = not wandb.env.error_reporting_enabled() + self._sent_messages: set = set() + + self.dsn = os.environ.get(wandb.env.SENTRY_DSN, SENTRY_DEFAULT_DSN) + + self.scope: sentry_sdk.scope.Scope | None = None + + # ensure we always end the Sentry session + atexit.register(self.end_session) + + @property + def environment(self) -> str: + """Return the environment we're running in.""" + # check if we're in a git repo + is_git = pathlib.Path(__file__).parent.parent.parent.joinpath(".git").exists() + + # these match the environments for gorilla + return "development" if is_git else "production" + + @_safe_noop + def setup(self) -> None: + """Setup Sentry SDK. + + We use lower-level APIs (i.e., not sentry_sdk.init) here + to avoid the possibility of interfering with the user's + own Sentry SDK setup. + """ + client = sentry_sdk.Client( + dsn=self.dsn, + default_integrations=False, + environment=self.environment, + release=wandb.__version__, + ) + self.scope = sentry_sdk.get_global_scope().fork() + self.scope.clear() + self.scope.set_client(client) + + @_safe_noop + def message(self, message: str, repeat: bool = True) -> str | None: + """Send a message to Sentry.""" + if not repeat and message in self._sent_messages: + return None + self._sent_messages.add(message) + with sentry_sdk.scope.use_isolation_scope(self.scope): # type: ignore + return sentry_sdk.capture_message(message) # type: ignore + + @_safe_noop + def exception( + self, + exc: str + | BaseException + | tuple[ + type[BaseException] | None, + BaseException | None, + TracebackType | None, + ] + | None, + handled: bool = False, + status: SessionStatus | None = None, + ) -> str | None: + """Log an exception to Sentry.""" + if isinstance(exc, str): + exc_info = sentry_sdk.utils.exc_info_from_error(Exception(exc)) + elif isinstance(exc, BaseException): + exc_info = sentry_sdk.utils.exc_info_from_error(exc) + else: + exc_info = sys.exc_info() + + event, _ = sentry_sdk.utils.event_from_exception( + exc_info, + client_options=self.scope.get_client().options, # type: ignore + mechanism={"type": "generic", "handled": handled}, + ) + event_id = None + try: + with sentry_sdk.scope.use_isolation_scope(self.scope): # type: ignore + event_id = sentry_sdk.capture_event(event) # type: ignore + except Exception: + pass + + # if the status is not explicitly set, we'll set it to "crashed" if the exception + # was unhandled, or "errored" if it was handled + status = status or ("crashed" if not handled else "errored") # type: ignore + self.mark_session(status=status) + + client = self.scope.get_client() # type: ignore + if client is not None: + client.flush() + + return event_id + + def reraise(self, exc: Any) -> Never: + """Re-raise an exception after logging it to Sentry. + + Use this for top-level exceptions when you want the user to see the traceback. + + Must be called from within an exception handler. + """ + self.exception(exc) + # this will messily add this "reraise" function to the stack trace, + # but hopefully it's not too bad + raise exc.with_traceback(sys.exc_info()[2]) + + @_safe_noop + def start_session(self) -> None: + """Start a new session.""" + assert self.scope is not None + # get the current client and scope + session = self.scope._session + + # if there's no session, start one + if session is None: + self.scope.start_session() + + @_safe_noop + def end_session(self) -> None: + """End the current session.""" + assert self.scope is not None + # get the current client and scope + client = self.scope.get_client() + session = self.scope._session + + if session is not None and client is not None: + self.scope.end_session() + client.flush() + + @_safe_noop + def mark_session(self, status: SessionStatus | None = None) -> None: + """Mark the current session with a status.""" + assert self.scope is not None + session = self.scope._session + + if session is not None: + session.update(status=status) + + @_safe_noop + def configure_scope( + self, + tags: dict[str, Any] | None = None, + process_context: str | None = None, + ) -> None: + """Configure the Sentry scope for the current thread. + + This function should be called at the beginning of every thread that + will send events to Sentry. It sets the tags that will be applied to + all events sent from this thread. It also tries to start a session + if one doesn't already exist for this thread. + """ + assert self.scope is not None + settings_tags = ( + "entity", + "project", + "run_id", + "run_url", + "sweep_url", + "sweep_id", + "deployment", + "launch", + "_platform", + ) + + # set context + if process_context: + self.scope.set_tag("process_context", process_context) + + # apply settings tags + if tags is None: + return None + + for tag in settings_tags: + val = tags.get(tag, None) + if val not in (None, ""): + self.scope.set_tag(tag, val) + + if tags.get("_colab", None): + python_runtime = "colab" + elif tags.get("_jupyter", None): + python_runtime = "jupyter" + elif tags.get("_ipython", None): + python_runtime = "ipython" + else: + python_runtime = "python" + self.scope.set_tag("python_runtime", python_runtime) + + # Construct run_url and sweep_url given run_id and sweep_id + for obj in ("run", "sweep"): + obj_id, obj_url = f"{obj}_id", f"{obj}_url" + if tags.get(obj_url, None): + continue + + try: + app_url = wandb.util.app_url(tags["base_url"]) # type: ignore + entity, project = (quote(tags[k]) for k in ("entity", "project")) # type: ignore + self.scope.set_tag( + obj_url, + f"{app_url}/{entity}/{project}/{obj}s/{tags[obj_id]}", + ) + except Exception: + pass + + email = tags.get("email") + if email: + self.scope.user = {"email": email} + + # todo: add back the option to pass general tags see: c645f625d1c1a3db4a6b0e2aa8e924fee101904c (wandb/util.py) + + self.start_session() diff --git a/lib/python3.12/site-packages/wandb/apis/__init__.py b/lib/python3.12/site-packages/wandb/apis/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4daf1e8773ee37ed9af6499243b96ddb9b79aff8 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/__init__.py @@ -0,0 +1,48 @@ +"""api.""" + +from typing import Callable + +import requests +from urllib3.exceptions import InsecureRequestWarning + +import wandb +from wandb import env, util + + +def _disable_ssl() -> Callable[[], None]: + # Because third party libraries may also use requests, we monkey patch it globally + # and turn off urllib3 warnings instead printing a global warning to the user. + wandb.termwarn( + "Disabling SSL verification. Connections to this server are not verified and may be insecure!" + ) + + requests.packages.urllib3.disable_warnings(category=InsecureRequestWarning) + old_merge_environment_settings = requests.Session.merge_environment_settings + + def merge_environment_settings(self, url, proxies, stream, verify, cert): + settings = old_merge_environment_settings( + self, url, proxies, stream, verify, cert + ) + settings["verify"] = False + return settings + + requests.Session.merge_environment_settings = merge_environment_settings + + def reset(): + requests.Session.merge_environment_settings = old_merge_environment_settings + + return reset + + +if env.ssl_disabled(): + _disable_ssl() + + +reset_path = util.vendor_setup() + +from .internal import Api as InternalApi # noqa +from .public import Api as PublicApi # noqa + +reset_path() + +__all__ = ["InternalApi", "PublicApi"] diff --git a/lib/python3.12/site-packages/wandb/apis/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ea25148789867b263407a61535e8d29d052992b0 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/__pycache__/attrs.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/__pycache__/attrs.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..81ee5c898567d514d16471beb24d14f5e9cedab0 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/__pycache__/attrs.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/__pycache__/internal.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/__pycache__/internal.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0e077baa83921a85cf2abe3984ba0b0cd56cb90c Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/__pycache__/internal.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/__pycache__/normalize.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/__pycache__/normalize.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2fa6aa8705e67dda69c88246ac5065a9442bdf73 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/__pycache__/normalize.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/__pycache__/paginator.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/__pycache__/paginator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2a8319a4d2412ce7ff54f47b5b2665b4cd069ebd Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/__pycache__/paginator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/attrs.py b/lib/python3.12/site-packages/wandb/apis/attrs.py new file mode 100644 index 0000000000000000000000000000000000000000..9eb497af14f470afaefcccf33047f32ec22a9cdb --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/attrs.py @@ -0,0 +1,50 @@ +from typing import Any, MutableMapping + +import wandb + +from ..sdk.lib import ipython + + +class Attrs: + def __init__(self, attrs: MutableMapping[str, Any]): + self._attrs = attrs + + def snake_to_camel(self, string): + camel = "".join([i.title() for i in string.split("_")]) + return camel[0].lower() + camel[1:] + + def display(self, height=420, hidden=False) -> bool: + """Display this object in jupyter.""" + if wandb.run and wandb.run._settings.silent: + return False + + if not ipython.in_jupyter(): + return False + + html = self.to_html(height, hidden) + if html is None: + wandb.termwarn("This object does not support `.display()`") + return False + + try: + from IPython import display + except ImportError: + wandb.termwarn(".display() only works in jupyter environments") + return False + + display.display(display.HTML(html)) + return True + + def to_html(self, *args, **kwargs): + return None + + def __getattr__(self, name): + key = self.snake_to_camel(name) + if key == "user": + raise AttributeError + if key in self._attrs.keys(): + return self._attrs[key] + elif name in self._attrs.keys(): + return self._attrs[name] + else: + raise AttributeError(f"{repr(self)!r} object has no attribute {name!r}") diff --git a/lib/python3.12/site-packages/wandb/apis/importers/__init__.py b/lib/python3.12/site-packages/wandb/apis/importers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0266b9f204d4a5487c5d66a6d2487c3670055082 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/importers/__init__.py @@ -0,0 +1 @@ +from .internals.util import Namespace diff --git a/lib/python3.12/site-packages/wandb/apis/importers/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/importers/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..39050c7637db1b0190c14408cd8c02ca22b8c106 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/importers/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/importers/__pycache__/mlflow.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/importers/__pycache__/mlflow.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..daa6e3001bad8e56be4a35b31fff61e9175f95fe Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/importers/__pycache__/mlflow.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/importers/__pycache__/validation.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/importers/__pycache__/validation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..113f2b8b7157185c5411299979ad08076ea7a4ee Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/importers/__pycache__/validation.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/importers/__pycache__/wandb.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/importers/__pycache__/wandb.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..51060d8ecb00d2ab5a0cb97a127c08076fb08338 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/importers/__pycache__/wandb.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/importers/internals/__pycache__/internal.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/importers/internals/__pycache__/internal.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4cc7bccbbe71e42196ef224249713a5451194ef7 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/importers/internals/__pycache__/internal.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/importers/internals/__pycache__/protocols.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/importers/internals/__pycache__/protocols.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..777548b3b84909a9d3035e8c27f7132a62dd90d5 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/importers/internals/__pycache__/protocols.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/importers/internals/__pycache__/util.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/importers/internals/__pycache__/util.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a798c27ee7c7f500f33d35ee6ca07b1ca91f5a69 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/importers/internals/__pycache__/util.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/importers/internals/internal.py b/lib/python3.12/site-packages/wandb/apis/importers/internals/internal.py new file mode 100644 index 0000000000000000000000000000000000000000..3ae10924ace8a9305cb7aa8805a8f94af3c1a3aa --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/importers/internals/internal.py @@ -0,0 +1,382 @@ +import json +import logging +import math +import os +import queue +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Dict, Iterable, Optional + +import numpy as np +from google.protobuf.json_format import ParseDict +from tenacity import retry, stop_after_attempt, wait_random_exponential + +from wandb import Artifact +from wandb.proto import wandb_internal_pb2 as pb +from wandb.proto import wandb_settings_pb2 +from wandb.proto import wandb_telemetry_pb2 as telem_pb +from wandb.sdk.interface.interface import file_policy_to_enum +from wandb.sdk.interface.interface_queue import InterfaceQueue +from wandb.sdk.internal import context +from wandb.sdk.internal.sender import SendManager +from wandb.sdk.internal.settings_static import SettingsStatic +from wandb.util import coalesce, recursive_cast_dictlike_to_dict + +from .protocols import ImporterRun + +ROOT_DIR = "./wandb-importer" + + +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) + +if os.getenv("WANDB_IMPORTER_ENABLE_RICH_LOGGING"): + from rich.logging import RichHandler + + logger.addHandler(RichHandler(rich_tracebacks=True, tracebacks_show_locals=True)) +else: + console_handler = logging.StreamHandler() + console_handler.setLevel(logging.INFO) + + formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s") + console_handler.setFormatter(formatter) + + logger.addHandler(console_handler) + + +exp_retry = retry( + wait=wait_random_exponential(multiplier=1, max=10), stop=stop_after_attempt(3) +) + + +class AlternateSendManager(SendManager): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._send_artifact = exp_retry(self._send_artifact) + + +@dataclass(frozen=True) +class SendManagerConfig: + """Configure which parts of SendManager tooling to use.""" + + use_artifacts: bool = False + log_artifacts: bool = False + metadata: bool = False + files: bool = False + media: bool = False + code: bool = False + history: bool = False + summary: bool = False + terminal_output: bool = False + + +@dataclass +class RecordMaker: + run: ImporterRun + interface: InterfaceQueue = InterfaceQueue() + + @property + def run_dir(self) -> str: + p = Path(f"{ROOT_DIR}/{self.run.run_id()}/wandb") + p.mkdir(parents=True, exist_ok=True) + return f"{ROOT_DIR}/{self.run.run_id()}" + + def make_artifacts_only_records( + self, + artifacts: Optional[Iterable[Artifact]] = None, + used_artifacts: Optional[Iterable[Artifact]] = None, + ) -> Iterable[pb.Record]: + """Only make records required to upload artifacts. + + Escape hatch for adding extra artifacts to a run. + """ + yield self._make_run_record() + + if used_artifacts: + for art in used_artifacts: + yield self._make_artifact_record(art, use_artifact=True) + + if artifacts: + for art in artifacts: + yield self._make_artifact_record(art) + + def make_records( + self, + config: SendManagerConfig, + ) -> Iterable[pb.Record]: + """Make all the records that constitute a run.""" + yield self._make_run_record() + yield self._make_telem_record() + + include_artifacts = config.log_artifacts or config.use_artifacts + yield self._make_files_record( + include_artifacts, config.files, config.media, config.code + ) + + if config.use_artifacts: + if (used_artifacts := self.run.used_artifacts()) is not None: + for artifact in used_artifacts: + yield self._make_artifact_record(artifact, use_artifact=True) + + if config.log_artifacts: + if (artifacts := self.run.artifacts()) is not None: + for artifact in artifacts: + yield self._make_artifact_record(artifact) + + if config.history: + yield from self._make_history_records() + + if config.summary: + yield self._make_summary_record() + + if config.terminal_output: + if (lines := self.run.logs()) is not None: + for line in lines: + yield self._make_output_record(line) + + def _make_run_record(self) -> pb.Record: + run = pb.RunRecord() + run.run_id = self.run.run_id() + run.entity = self.run.entity() + run.project = self.run.project() + run.display_name = coalesce(self.run.display_name()) + run.notes = coalesce(self.run.notes(), "") + run.tags.extend(coalesce(self.run.tags(), [])) + run.start_time.FromMilliseconds(self.run.start_time()) + + host = self.run.host() + if host is not None: + run.host = host + + runtime = self.run.runtime() + if runtime is not None: + run.runtime = runtime + + run_group = self.run.run_group() + if run_group is not None: + run.run_group = run_group + + config = self.run.config() + if "_wandb" not in config: + config["_wandb"] = {} + + # how do I get this automatically? + config["_wandb"]["code_path"] = self.run.code_path() + config["_wandb"]["python_version"] = self.run.python_version() + config["_wandb"]["cli_version"] = self.run.cli_version() + + self.interface._make_config( + data=config, + obj=run.config, + ) # is there a better way? + return self.interface._make_record(run=run) + + def _make_output_record(self, line) -> pb.Record: + output_raw = pb.OutputRawRecord() + output_raw.output_type = pb.OutputRawRecord.OutputType.STDOUT + output_raw.line = line + return self.interface._make_record(output_raw=output_raw) + + def _make_summary_record(self) -> pb.Record: + d: dict = { + **self.run.summary(), + "_runtime": self.run.runtime(), # quirk of runtime -- it has to be here! + # '_timestamp': self.run.start_time()/1000, + } + d = recursive_cast_dictlike_to_dict(d) + summary = self.interface._make_summary_from_dict(d) + return self.interface._make_record(summary=summary) + + def _make_history_records(self) -> Iterable[pb.Record]: + for metrics in self.run.metrics(): + history = pb.HistoryRecord() + for k, v in metrics.items(): + item = history.item.add() + item.key = k + # There seems to be some conversion issue to breaks when we try to re-upload. + # np.NaN gets converted to float("nan"), which is not expected by our system. + # If this cast to string (!) is not done, the row will be dropped. + if (isinstance(v, float) and math.isnan(v)) or v == "NaN": + v = np.NaN + + if isinstance(v, bytes): + # it's a json string encoded as bytes + v = v.decode("utf-8") + else: + v = json.dumps(v) + + item.value_json = v + rec = self.interface._make_record(history=history) + yield rec + + def _make_files_record( + self, artifacts: bool, files: bool, media: bool, code: bool + ) -> pb.Record: + run_files = self.run.files() + metadata_fname = f"{self.run_dir}/files/wandb-metadata.json" + if not files or run_files is None: + # We'll always need a metadata file even if there are no other files to upload + metadata_fname = self._make_metadata_file() + run_files = [(metadata_fname, "end")] + files_record = pb.FilesRecord() + for path, policy in run_files: + if not artifacts and path.startswith("artifact/"): + continue + if not media and path.startswith("media/"): + continue + if not code and path.startswith("code/"): + continue + + # DirWatcher requires the path to start with media/ instead of the full path + if "media" in path: + p = Path(path) + path = str(p.relative_to(f"{self.run_dir}/files")) + f = files_record.files.add() + f.path = path + f.policy = file_policy_to_enum(policy) + + return self.interface._make_record(files=files_record) + + def _make_artifact_record( + self, artifact: Artifact, use_artifact=False + ) -> pb.Record: + proto = self.interface._make_artifact(artifact) + proto.run_id = str(self.run.run_id()) + proto.project = str(self.run.project()) + proto.entity = str(self.run.entity()) + proto.user_created = use_artifact + proto.use_after_commit = use_artifact + proto.finalize = True + + aliases = artifact._aliases + aliases += ["latest", "imported"] + + for alias in aliases: + proto.aliases.append(alias) + return self.interface._make_record(artifact=proto) + + def _make_telem_record(self) -> pb.Record: + telem = telem_pb.TelemetryRecord() + + feature = telem_pb.Feature() + feature.importer_mlflow = True + telem.feature.CopyFrom(feature) + + cli_version = self.run.cli_version() + if cli_version: + telem.cli_version = cli_version + + python_version = self.run.python_version() + if python_version: + telem.python_version = python_version + + return self.interface._make_record(telemetry=telem) + + def _make_metadata_file(self) -> str: + missing_text = "This data was not captured" + files_dir = f"{self.run_dir}/files" + os.makedirs(files_dir, exist_ok=True) + + d = {} + d["os"] = coalesce(self.run.os_version(), missing_text) + d["python"] = coalesce(self.run.python_version(), missing_text) + d["program"] = coalesce(self.run.program(), missing_text) + d["cuda"] = coalesce(self.run.cuda_version(), missing_text) + d["host"] = coalesce(self.run.host(), missing_text) + d["username"] = coalesce(self.run.username(), missing_text) + d["executable"] = coalesce(self.run.executable(), missing_text) + + gpus_used = self.run.gpus_used() + if gpus_used is not None: + d["gpu_devices"] = json.dumps(gpus_used) + d["gpu_count"] = json.dumps(len(gpus_used)) + + cpus_used = self.run.cpus_used() + if cpus_used is not None: + d["cpu_count"] = json.dumps(self.run.cpus_used()) + + mem_used = self.run.memory_used() + if mem_used is not None: + d["memory"] = json.dumps({"total": self.run.memory_used()}) + + fname = f"{files_dir}/wandb-metadata.json" + with open(fname, "w") as f: + f.write(json.dumps(d)) + return fname + + +def _make_settings( + root_dir: str, settings_override: Optional[Dict[str, Any]] = None +) -> SettingsStatic: + _settings_override = coalesce(settings_override, {}) + + default_settings: Dict[str, Any] = { + "x_files_dir": os.path.join(root_dir, "files"), + "root_dir": root_dir, + "sync_file": os.path.join(root_dir, "txlog.wandb"), + "resume": "false", + "program": None, + "ignore_globs": [], + "disable_job_creation": True, + "x_start_time": 0, + "_offline": None, + "x_sync": True, + "x_live_policy_rate_limit": 15, # matches dir_watcher + "x_live_policy_wait_time": 600, # matches dir_watcher + "x_file_stream_timeout_seconds": 60, + } + + combined_settings = {**default_settings, **_settings_override} + settings_message = wandb_settings_pb2.Settings() + ParseDict(combined_settings, settings_message) + + return SettingsStatic(settings_message) + + +def send_run( + run: ImporterRun, + *, + extra_arts: Optional[Iterable[Artifact]] = None, + extra_used_arts: Optional[Iterable[Artifact]] = None, + config: Optional[SendManagerConfig] = None, + overrides: Optional[Dict[str, Any]] = None, + settings_override: Optional[Dict[str, Any]] = None, +) -> None: + if config is None: + config = SendManagerConfig() + + # does this need to be here for pmap? + if overrides: + for k, v in overrides.items(): + # `lambda: v` won't work! + # https://stackoverflow.com/questions/10802002/why-deepcopy-doesnt-create-new-references-to-lambda-function + setattr(run, k, lambda v=v: v) + + rm = RecordMaker(run) + root_dir = rm.run_dir + + settings = _make_settings(root_dir, settings_override) + sm_record_q = queue.Queue() + # wm_record_q = queue.Queue() + result_q = queue.Queue() + interface = InterfaceQueue(record_q=sm_record_q) + context_keeper = context.ContextKeeper() + sm = AlternateSendManager( + settings, sm_record_q, result_q, interface, context_keeper + ) + + if extra_arts or extra_used_arts: + records = rm.make_artifacts_only_records(extra_arts, extra_used_arts) + else: + records = rm.make_records(config) + + for r in records: + logger.debug(f"Sending {r=}") + # In a future update, it might be good to write to a transaction log and have + # incremental uploads only send the missing records. + # wm.write(r) + + sm.send(r) + + sm.finish() + # wm.finish() diff --git a/lib/python3.12/site-packages/wandb/apis/importers/internals/protocols.py b/lib/python3.12/site-packages/wandb/apis/importers/internals/protocols.py new file mode 100644 index 0000000000000000000000000000000000000000..87d42e99175ce5d23db6e375a8cfcaa66568722a --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/importers/internals/protocols.py @@ -0,0 +1,103 @@ +import logging +from typing import ( + Any, + Dict, + Iterable, + List, + Literal, + Optional, + Protocol, + Tuple, + runtime_checkable, +) + +from wandb.sdk.artifacts.artifact import Artifact + +logger = logging.getLogger("import_logger") + +PathStr = str +Policy = Literal["now", "end", "live"] + + +@runtime_checkable +class ImporterRun(Protocol): + def run_id(self) -> str: ... # pragma: no cover + + def entity(self) -> str: ... # pragma: no cover + + def project(self) -> str: ... # pragma: no cover + + def config(self) -> Dict[str, Any]: ... # pragma: no cover + + def summary(self) -> Dict[str, float]: ... # pragma: no cover + + def metrics(self) -> Iterable[Dict[str, float]]: + """Metrics for the run. + + We expect metrics in this shape: + + [ + {'metric1': 1, 'metric2': 1, '_step': 0}, + {'metric1': 2, 'metric2': 4, '_step': 1}, + {'metric1': 3, 'metric2': 9, '_step': 2}, + ... + ] + + You can also submit metrics in this shape: + [ + {'metric1': 1, '_step': 0}, + {'metric2': 1, '_step': 0}, + {'metric1': 2, '_step': 1}, + {'metric2': 4, '_step': 1}, + ... + ] + """ + ... # pragma: no cover + + def run_group(self) -> Optional[str]: ... # pragma: no cover + + def job_type(self) -> Optional[str]: ... # pragma: no cover + + def display_name(self) -> str: ... # pragma: no cover + + def notes(self) -> Optional[str]: ... # pragma: no cover + + def tags(self) -> Optional[List[str]]: ... # pragma: no cover + + def artifacts(self) -> Optional[Iterable[Artifact]]: ... # pragma: no cover + + def used_artifacts(self) -> Optional[Iterable[Artifact]]: ... # pragma: no cover + + def os_version(self) -> Optional[str]: ... # pragma: no cover + + def python_version(self) -> Optional[str]: ... # pragma: no cover + + def cuda_version(self) -> Optional[str]: ... # pragma: no cover + + def program(self) -> Optional[str]: ... # pragma: no cover + + def host(self) -> Optional[str]: ... # pragma: no cover + + def username(self) -> Optional[str]: ... # pragma: no cover + + def executable(self) -> Optional[str]: ... # pragma: no cover + + def gpus_used(self) -> Optional[str]: ... # pragma: no cover + + def cpus_used(self) -> Optional[int]: ... # pragma: no cover + + def memory_used(self) -> Optional[int]: ... # pragma: no cover + + def runtime(self) -> Optional[int]: ... # pragma: no cover + + def start_time(self) -> Optional[int]: ... # pragma: no cover + + def code_path(self) -> Optional[str]: ... # pragma: no cover + + def cli_version(self) -> Optional[str]: ... # pragma: no cover + + def files( + self, + ) -> Optional[Iterable[Tuple[PathStr, Policy]]]: ... # pragma: no cover + + def logs(self) -> Optional[Iterable[str]]: ... # pragma: no cover diff --git a/lib/python3.12/site-packages/wandb/apis/importers/internals/util.py b/lib/python3.12/site-packages/wandb/apis/importers/internals/util.py new file mode 100644 index 0000000000000000000000000000000000000000..d77585f42e35c26ddc2607620e30888e8b007286 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/importers/internals/util.py @@ -0,0 +1,78 @@ +import logging +import sys +import traceback +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import dataclass +from typing import Iterable, Optional + + +@dataclass(frozen=True) +class Namespace: + """Configure an alternate entity/project at the dst server your data will end up in.""" + + entity: str + project: str + + @classmethod + def from_path(cls, path: str): + entity, project = path.split("/") + return cls(entity, project) + + @property + def path(self): + return f"{self.entity}/{self.project}" + + @property + def send_manager_overrides(self): + overrides = {} + if self.entity: + overrides["entity"] = self.entity + if self.project: + overrides["project"] = self.project + return overrides + + +logger = logging.getLogger("import_logger") + + +def parallelize( + func, + iterable: Iterable, + *args, + max_workers: Optional[int] = None, + raise_on_error: bool = False, + **kwargs, +): + def safe_func(*args, **kwargs): + try: + return func(*args, **kwargs) + except Exception as e: + _, _, exc_traceback = sys.exc_info() + traceback_details = traceback.extract_tb(exc_traceback) + filename = traceback_details[-1].filename + lineno = traceback_details[-1].lineno + logger.debug( + f"Exception: {func=} {args=} {kwargs=} {e=} {filename=} {lineno=}. {traceback_details=}" + ) + if raise_on_error: + raise + + results = [] + with ThreadPoolExecutor(max_workers) as exc: + futures = {exc.submit(safe_func, x, *args, **kwargs): x for x in iterable} + for future in as_completed(futures): + results.append(future.result()) + return results + + +def for_each( + func, iterable: Iterable, parallel: bool = True, max_workers: Optional[int] = None +): + if parallel: + return parallelize( + func, + iterable, + max_workers=max_workers, + ) + + return [func(x) for x in iterable] diff --git a/lib/python3.12/site-packages/wandb/apis/importers/mlflow.py b/lib/python3.12/site-packages/wandb/apis/importers/mlflow.py new file mode 100644 index 0000000000000000000000000000000000000000..3ae2d8c76ee1fa81d9b936f511fc5912378235f1 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/importers/mlflow.py @@ -0,0 +1,254 @@ +import itertools +import logging +import re +from collections import defaultdict +from typing import Any, Dict, Iterable, List, Optional, Tuple + +import mlflow +from packaging.version import Version # type: ignore + +import wandb +from wandb import Artifact + +from .internals import internal +from .internals.util import Namespace, for_each + +mlflow_version = Version(mlflow.__version__) + +logger = logging.getLogger("import_logger") + + +class MlflowRun: + def __init__(self, run, mlflow_client): + self.run = run + self.mlflow_client: mlflow.MlflowClient = mlflow_client + + def run_id(self) -> str: + return self.run.info.run_id + + def entity(self) -> str: + return self.run.info.user_id + + def project(self) -> str: + return "imported-from-mlflow" + + def config(self) -> Dict[str, Any]: + conf = self.run.data.params + + # Add tags here since mlflow supports very long tag names but we only support up to 64 chars + tags = { + k: v for k, v in self.run.data.tags.items() if not k.startswith("mlflow.") + } + return {**conf, "imported_mlflow_tags": tags} + + def summary(self) -> Dict[str, float]: + return self.run.data.metrics + + def metrics(self) -> Iterable[Dict[str, float]]: + d: Dict[int, Dict[str, float]] = defaultdict(dict) + for k in self.run.data.metrics.keys(): + metric = self.mlflow_client.get_metric_history(self.run.info.run_id, k) + for item in metric: + d[item.step][item.key] = item.value + + for k, v in d.items(): + yield {"_step": k, **v} + + def run_group(self) -> Optional[str]: + # this is nesting? Parent at `run.info.tags.get("mlflow.parentRunId")` + return f"Experiment {self.run.info.experiment_id}" + + def job_type(self) -> Optional[str]: + # Is this the right approach? + return f"User {self.run.info.user_id}" + + def display_name(self) -> str: + if mlflow_version < Version("1.30.0"): + return self.run.data.tags["mlflow.runName"] + return self.run.info.run_name + + def notes(self) -> Optional[str]: + return self.run.data.tags.get("mlflow.note.content") + + def tags(self) -> Optional[List[str]]: + ... + + # W&B tags are different than mlflow tags. + # The full mlflow tags are added to config under key `imported_mlflow_tags` instead + + def artifacts(self) -> Optional[Iterable[Artifact]]: # type: ignore + if mlflow_version < Version("2.0.0"): + dir_path = self.mlflow_client.download_artifacts( + run_id=self.run.info.run_id, + path="", + ) + else: + dir_path = mlflow.artifacts.download_artifacts(run_id=self.run.info.run_id) + + # Since mlflow doesn't have extra metadata about the artifacts, + # we just lump them all together into a single wandb.Artifact + artifact_name = self._handle_incompatible_strings(self.display_name()) + art = wandb.Artifact(artifact_name, "imported-artifacts") + art.add_dir(dir_path) + + return [art] + + def used_artifacts(self) -> Optional[Iterable[Artifact]]: # type: ignore + ... # pragma: no cover + + def os_version(self) -> Optional[str]: ... # pragma: no cover + + def python_version(self) -> Optional[str]: ... # pragma: no cover + + def cuda_version(self) -> Optional[str]: ... # pragma: no cover + + def program(self) -> Optional[str]: ... # pragma: no cover + + def host(self) -> Optional[str]: ... # pragma: no cover + + def username(self) -> Optional[str]: ... # pragma: no cover + + def executable(self) -> Optional[str]: ... # pragma: no cover + + def gpus_used(self) -> Optional[str]: ... # pragma: no cover + + def cpus_used(self) -> Optional[int]: # can we get the model? + ... # pragma: no cover + + def memory_used(self) -> Optional[int]: ... # pragma: no cover + + def runtime(self) -> Optional[int]: + end_time = ( + self.run.info.end_time // 1000 + if self.run.info.end_time is not None + else self.start_time() + ) + return end_time - self.start_time() + + def start_time(self) -> Optional[int]: + return self.run.info.start_time // 1000 + + def code_path(self) -> Optional[str]: ... # pragma: no cover + + def cli_version(self) -> Optional[str]: ... # pragma: no cover + + def files(self) -> Optional[Iterable[Tuple[str, str]]]: ... # pragma: no cover + + def logs(self) -> Optional[Iterable[str]]: ... # pragma: no cover + + @staticmethod + def _handle_incompatible_strings(s: str) -> str: + valid_chars = r"[^a-zA-Z0-9_\-\.]" + replacement = "__" + + return re.sub(valid_chars, replacement, s) + + +class MlflowImporter: + def __init__( + self, + dst_base_url: str, + dst_api_key: str, + mlflow_tracking_uri: str, + mlflow_registry_uri: Optional[str] = None, + *, + custom_api_kwargs: Optional[Dict[str, Any]] = None, + ) -> None: + self.dst_base_url = dst_base_url + self.dst_api_key = dst_api_key + + if custom_api_kwargs is None: + custom_api_kwargs = {"timeout": 600} + + self.dst_api = wandb.Api( + api_key=dst_api_key, + overrides={"base_url": dst_base_url}, + **custom_api_kwargs, + ) + self.mlflow_tracking_uri = mlflow_tracking_uri + mlflow.set_tracking_uri(self.mlflow_tracking_uri) + + if mlflow_registry_uri: + mlflow.set_registry_uri(mlflow_registry_uri) + + self.mlflow_client = mlflow.tracking.MlflowClient(mlflow_tracking_uri) + + def __repr__(self): + return f"" + + def collect_runs(self, *, limit: Optional[int] = None) -> Iterable[MlflowRun]: + if mlflow_version < Version("1.28.0"): + experiments = self.mlflow_client.list_experiments() + else: + experiments = self.mlflow_client.search_experiments() + + def _runs(): + for exp in experiments: + for run in self.mlflow_client.search_runs(exp.experiment_id): + yield MlflowRun(run, self.mlflow_client) + + runs = itertools.islice(_runs(), limit) + yield from runs + + def _import_run( + self, + run: MlflowRun, + *, + artifacts: bool = True, + namespace: Optional[Namespace] = None, + config: Optional[internal.SendManagerConfig] = None, + ) -> None: + if namespace is None: + namespace = Namespace(run.entity(), run.project()) + + if config is None: + config = internal.SendManagerConfig( + metadata=True, + files=True, + media=True, + code=True, + history=True, + summary=True, + terminal_output=True, + ) + + settings_override = { + "api_key": self.dst_api_key, + "base_url": self.dst_base_url, + "resume": "true", + "resumed": True, + } + + mlflow.set_tracking_uri(self.mlflow_tracking_uri) + internal.send_run( + run, + overrides=namespace.send_manager_overrides, + settings_override=settings_override, + config=config, + ) + + # in mlflow, the artifacts come with the runs, so import them together + if artifacts: + arts = list(run.artifacts()) + logger.debug(f"Importing history artifacts, {run=}") + internal.send_run( + run, + extra_arts=arts, + overrides=namespace.send_manager_overrides, + settings_override=settings_override, + config=internal.SendManagerConfig(log_artifacts=True), + ) + + def import_runs( + self, + runs: Iterable[MlflowRun], + *, + artifacts: bool = True, + namespace: Optional[Namespace] = None, + parallel: bool = True, + max_workers: Optional[int] = None, + ) -> None: + def _import_run_wrapped(run): + self._import_run(run, namespace=namespace, artifacts=artifacts) + + for_each(_import_run_wrapped, runs, parallel=parallel, max_workers=max_workers) diff --git a/lib/python3.12/site-packages/wandb/apis/importers/validation.py b/lib/python3.12/site-packages/wandb/apis/importers/validation.py new file mode 100644 index 0000000000000000000000000000000000000000..219c40285935b6681c819c2c6de006252a806aa7 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/importers/validation.py @@ -0,0 +1,108 @@ +import filecmp +import logging +import os + +import requests + +import wandb + +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) + + +def _compare_artifact_manifests( + src_art: wandb.Artifact, dst_art: wandb.Artifact +) -> list: + problems = [] + if isinstance(dst_art, wandb.CommError): + return ["commError"] + + if src_art.digest != dst_art.digest: + problems.append(f"digest mismatch {src_art.digest=}, {dst_art.digest=}") + + for name, src_entry in src_art.manifest.entries.items(): + dst_entry = dst_art.manifest.entries.get(name) + if dst_entry is None: + problems.append(f"missing manifest entry {name=}, {src_entry=}") + continue + + for attr in ["path", "digest", "size"]: + if getattr(src_entry, attr) != getattr(dst_entry, attr): + problems.append( + f"manifest entry mismatch {attr=}, {getattr(src_entry, attr)=}, {getattr(dst_entry, attr)=}" + ) + + return problems + + +def _compare_artifact_dirs(src_dir, dst_dir) -> list: + def compare(src_dir, dst_dir): + comparison = filecmp.dircmp(src_dir, dst_dir) + differences = { + "left_only": comparison.left_only, + "right_only": comparison.right_only, + "diff_files": comparison.diff_files, + "subdir_differences": {}, + } + + # Recursively find differences in subdirectories + for subdir in comparison.subdirs: + subdir_src = os.path.join(src_dir, subdir) + subdir_dst = os.path.join(dst_dir, subdir) + subdir_differences = compare(subdir_src, subdir_dst) + # If there are differences, add them to the result + if subdir_differences and any(subdir_differences.values()): + differences["subdir_differences"][subdir] = subdir_differences + + if all(not diff for diff in differences.values()): + return None + + return differences + + return compare(src_dir, dst_dir) + + +def _check_entries_are_downloadable(art): + entries = _collect_entries(art) + for entry in entries: + if not _check_entry_is_downloable(entry): + return False + return True + + +def _collect_entries(art): + has_next_page = True + cursor = None + entries = [] + while has_next_page: + attrs = art._fetch_file_urls(cursor) + has_next_page = attrs["pageInfo"]["hasNextPage"] + cursor = attrs["pageInfo"]["endCursor"] + for edge in attrs["edges"]: + name = edge["node"]["name"] + entry = art.get_entry(name) + entry._download_url = edge["node"]["directUrl"] + entries.append(entry) + return entries + + +def _check_entry_is_downloable(entry): + url = entry._download_url + expected_size = entry.size + + try: + resp = requests.head(url, allow_redirects=True) + except Exception: + logger.exception(f"Problem validating {entry=}") + return False + + if resp.status_code != 200: + return False + + actual_size = resp.headers.get("content-length", -1) + actual_size = int(actual_size) + + if expected_size == actual_size: + return True + + return False diff --git a/lib/python3.12/site-packages/wandb/apis/importers/wandb.py b/lib/python3.12/site-packages/wandb/apis/importers/wandb.py new file mode 100644 index 0000000000000000000000000000000000000000..b325668295a1d245c1af01ed5a0ab3a997796b65 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/importers/wandb.py @@ -0,0 +1,1608 @@ +"""Tooling for the W&B Importer.""" + +import itertools +import json +import logging +import numbers +import os +import re +import shutil +from dataclasses import dataclass, field +from datetime import datetime as dt +from pathlib import Path +from typing import Any, Dict, Iterable, Iterator, List, Optional, Tuple +from unittest.mock import patch + +import filelock +import polars as pl +import requests +import urllib3 +import wandb_workspaces.reports.v1 as wr +import yaml +from wandb_gql import gql +from wandb_workspaces.reports.v1 import Report + +import wandb +from wandb.apis.public import ArtifactCollection, Run +from wandb.apis.public.files import File +from wandb.util import coalesce, remove_keys_with_none_values + +from . import validation +from .internals import internal +from .internals.protocols import PathStr, Policy +from .internals.util import Namespace, for_each + +Artifact = wandb.Artifact +Api = wandb.Api +Project = wandb.apis.public.Project + +ARTIFACT_ERRORS_FNAME = "artifact_errors.jsonl" +ARTIFACT_SUCCESSES_FNAME = "artifact_successes.jsonl" +RUN_ERRORS_FNAME = "run_errors.jsonl" +RUN_SUCCESSES_FNAME = "run_successes.jsonl" + +ART_SEQUENCE_DUMMY_PLACEHOLDER = "__ART_SEQUENCE_DUMMY_PLACEHOLDER__" +RUN_DUMMY_PLACEHOLDER = "__RUN_DUMMY_PLACEHOLDER__" +ART_DUMMY_PLACEHOLDER_PATH = "__importer_temp__" +ART_DUMMY_PLACEHOLDER_TYPE = "__temp__" + +SRC_ART_PATH = "./artifacts/src" +DST_ART_PATH = "./artifacts/dst" + + +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) + +if os.getenv("WANDB_IMPORTER_ENABLE_RICH_LOGGING"): + from rich.logging import RichHandler + + logger.addHandler(RichHandler(rich_tracebacks=True, tracebacks_show_locals=True)) +else: + console_handler = logging.StreamHandler() + console_handler.setLevel(logging.INFO) + + formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s") + console_handler.setFormatter(formatter) + + logger.addHandler(console_handler) + + +@dataclass +class ArtifactSequence: + artifacts: Iterable[wandb.Artifact] + entity: str + project: str + type_: str + name: str + + def __iter__(self) -> Iterator: + return iter(self.artifacts) + + def __repr__(self) -> str: + return f"ArtifactSequence({self.identifier})" + + @property + def identifier(self) -> str: + return "/".join([self.entity, self.project, self.type_, self.name]) + + @classmethod + def from_collection(cls, collection: ArtifactCollection): + arts = collection.artifacts() + arts = sorted(arts, key=lambda a: int(a.version.lstrip("v"))) + return ArtifactSequence( + arts, + collection.entity, + collection.project, + collection.type, + collection.name, + ) + + +class WandbRun: + def __init__( + self, + run: Run, + *, + src_base_url: str, + src_api_key: str, + dst_base_url: str, + dst_api_key: str, + ) -> None: + self.run = run + self.api = wandb.Api( + api_key=src_api_key, + overrides={"base_url": src_base_url}, + ) + self.dst_api = wandb.Api( + api_key=dst_api_key, + overrides={"base_url": dst_base_url}, + ) + + # For caching + self._files: Optional[Iterable[Tuple[str, str]]] = None + self._artifacts: Optional[Iterable[Artifact]] = None + self._used_artifacts: Optional[Iterable[Artifact]] = None + self._parquet_history_paths: Optional[Iterable[str]] = None + + def __repr__(self) -> str: + s = os.path.join(self.entity(), self.project(), self.run_id()) + return f"WandbRun({s})" + + def run_id(self) -> str: + return self.run.id + + def entity(self) -> str: + return self.run.entity + + def project(self) -> str: + return self.run.project + + def config(self) -> Dict[str, Any]: + return self.run.config + + def summary(self) -> Dict[str, float]: + s = self.run.summary + return s + + def metrics(self) -> Iterable[Dict[str, float]]: + if self._parquet_history_paths is None: + self._parquet_history_paths = list(self._get_parquet_history_paths()) + + if self._parquet_history_paths: + rows = self._get_rows_from_parquet_history_paths() + else: + logger.warning( + "No parquet files detected; using scan history (this may not be reliable)" + ) + rows = self.run.scan_history() + + for row in rows: + row = remove_keys_with_none_values(row) + yield row + + def run_group(self) -> Optional[str]: + return self.run.group + + def job_type(self) -> Optional[str]: + return self.run.job_type + + def display_name(self) -> str: + return self.run.display_name + + def notes(self) -> Optional[str]: + # Notes includes the previous notes and serves as a catch-all for things we missed or can't add back + previous_link = f"Imported from: {self.run.url}" + previous_author = f"Author: {self.run.user.username}" + + header = [previous_link, previous_author] + previous_notes = self.run.notes or "" + + return "\n".join(header) + "\n---\n" + previous_notes + + def tags(self) -> Optional[List[str]]: + return self.run.tags + + def artifacts(self) -> Optional[Iterable[Artifact]]: + if self._artifacts is None: + _artifacts = [] + for art in self.run.logged_artifacts(): + a = _clone_art(art) + _artifacts.append(a) + self._artifacts = _artifacts + + yield from self._artifacts + + def used_artifacts(self) -> Optional[Iterable[Artifact]]: + if self._used_artifacts is None: + _used_artifacts = [] + for art in self.run.used_artifacts(): + a = _clone_art(art) + _used_artifacts.append(a) + self._used_artifacts = _used_artifacts + + yield from self._used_artifacts + + def os_version(self) -> Optional[str]: ... # pragma: no cover + + def python_version(self) -> Optional[str]: + return self._metadata_file().get("python") + + def cuda_version(self) -> Optional[str]: ... # pragma: no cover + + def program(self) -> Optional[str]: ... # pragma: no cover + + def host(self) -> Optional[str]: + return self._metadata_file().get("host") + + def username(self) -> Optional[str]: ... # pragma: no cover + + def executable(self) -> Optional[str]: ... # pragma: no cover + + def gpus_used(self) -> Optional[str]: ... # pragma: no cover + + def cpus_used(self) -> Optional[int]: # can we get the model? + ... # pragma: no cover + + def memory_used(self) -> Optional[int]: ... # pragma: no cover + + def runtime(self) -> Optional[int]: + wandb_runtime = self.run.summary.get("_wandb", {}).get("runtime") + base_runtime = self.run.summary.get("_runtime") + + if (t := coalesce(wandb_runtime, base_runtime)) is None: + return t + return int(t) + + def start_time(self) -> Optional[int]: + t = dt.fromisoformat(self.run.created_at).timestamp() * 1000 + return int(t) + + def code_path(self) -> Optional[str]: + path = self._metadata_file().get("codePath", "") + return f"code/{path}" + + def cli_version(self) -> Optional[str]: + return self._config_file().get("_wandb", {}).get("value", {}).get("cli_version") + + def files(self) -> Optional[Iterable[Tuple[PathStr, Policy]]]: + if self._files is None: + files_dir = f"{internal.ROOT_DIR}/{self.run_id()}/files" + _files = [] + for f in self.run.files(): + f: File + # These optimizations are intended to avoid rate limiting when importing many runs in parallel + # Don't carry over empty files + if f.size == 0: + continue + # Skip deadlist to avoid overloading S3 + if "wandb_manifest.json.deadlist" in f.name: + continue + + result = f.download(files_dir, exist_ok=True, api=self.api) + file_and_policy = (result.name, "end") + _files.append(file_and_policy) + self._files = _files + + yield from self._files + + def logs(self) -> Optional[Iterable[str]]: + log_files = self._find_all_in_files_regex(r"^.*output\.log$") + for path in log_files: + with open(path) as f: + yield from f.readlines() + + def _metadata_file(self) -> Dict[str, Any]: + if (fname := self._find_in_files("wandb-metadata.json")) is None: + return {} + + with open(fname) as f: + return json.loads(f.read()) + + def _config_file(self) -> Dict[str, Any]: + if (fname := self._find_in_files("config.yaml")) is None: + return {} + + with open(fname) as f: + return yaml.safe_load(f) or {} + + def _get_rows_from_parquet_history_paths(self) -> Iterable[Dict[str, Any]]: + # Unfortunately, it's not feasible to validate non-parquet history + if not (paths := self._get_parquet_history_paths()): + yield {} + return + + # Collect and merge parquet history + dfs = [ + pl.read_parquet(p) for path in paths for p in Path(path).glob("*.parquet") + ] + if "_step" in (df := _merge_dfs(dfs)): + df = df.with_columns(pl.col("_step").cast(pl.Int64)) + yield from df.iter_rows(named=True) + + def _get_parquet_history_paths(self) -> Iterable[str]: + if self._parquet_history_paths is None: + paths = [] + # self.artifacts() returns a copy of the artifacts; use this to get raw + for art in self.run.logged_artifacts(): + if art.type != "wandb-history": + continue + if ( + path := _download_art(art, root=f"{SRC_ART_PATH}/{art.name}") + ) is None: + continue + paths.append(path) + self._parquet_history_paths = paths + + yield from self._parquet_history_paths + + def _find_in_files(self, name: str) -> Optional[str]: + if files := self.files(): + for path, _ in files: + if name in path: + return path + return None + + def _find_all_in_files_regex(self, regex: str) -> Iterable[str]: + if files := self.files(): + for path, _ in files: + if re.match(regex, path): + yield path + + +class WandbImporter: + """Transfers runs, reports, and artifact sequences between W&B instances.""" + + def __init__( + self, + src_base_url: str, + src_api_key: str, + dst_base_url: str, + dst_api_key: str, + *, + custom_api_kwargs: Optional[Dict[str, Any]] = None, + ) -> None: + self.src_base_url = src_base_url + self.src_api_key = src_api_key + self.dst_base_url = dst_base_url + self.dst_api_key = dst_api_key + + if custom_api_kwargs is None: + custom_api_kwargs = {"timeout": 600} + + self.src_api = wandb.Api( + api_key=src_api_key, + overrides={"base_url": src_base_url}, + **custom_api_kwargs, + ) + self.dst_api = wandb.Api( + api_key=dst_api_key, + overrides={"base_url": dst_base_url}, + **custom_api_kwargs, + ) + + self.run_api_kwargs = { + "src_base_url": src_base_url, + "src_api_key": src_api_key, + "dst_base_url": dst_base_url, + "dst_api_key": dst_api_key, + } + + def __repr__(self): + return f"" # pragma: no cover + + def _import_run( + self, + run: WandbRun, + *, + namespace: Optional[Namespace] = None, + config: Optional[internal.SendManagerConfig] = None, + ) -> None: + """Import one WandbRun. + + Use `namespace` to specify alternate settings like where the run should be uploaded + """ + if namespace is None: + namespace = Namespace(run.entity(), run.project()) + + if config is None: + config = internal.SendManagerConfig( + metadata=True, + files=True, + media=True, + code=True, + history=True, + summary=True, + terminal_output=True, + ) + + settings_override = { + "api_key": self.dst_api_key, + "base_url": self.dst_base_url, + "resume": "true", + "resumed": True, + } + + # Send run with base config + logger.debug(f"Importing run, {run=}") + internal.send_run( + run, + overrides=namespace.send_manager_overrides, + settings_override=settings_override, + config=config, + ) + + if config.history: + # Send run again with history artifacts in case config history=True, artifacts=False + # The history artifact must come with the actual history data + + logger.debug(f"Collecting history artifacts, {run=}") + history_arts = [] + for art in run.run.logged_artifacts(): + if art.type != "wandb-history": + continue + logger.debug(f"Collecting history artifact {art.name=}") + new_art = _clone_art(art) + history_arts.append(new_art) + + logger.debug(f"Importing history artifacts, {run=}") + internal.send_run( + run, + extra_arts=history_arts, + overrides=namespace.send_manager_overrides, + settings_override=settings_override, + config=config, + ) + + def _delete_collection_in_dst( + self, + seq: ArtifactSequence, + namespace: Optional[Namespace] = None, + ): + """Deletes the equivalent artifact collection in destination. + + Intended to clear the destination when an uploaded artifact does not pass validation. + """ + entity = coalesce(namespace.entity, seq.entity) + project = coalesce(namespace.project, seq.project) + art_type = f"{entity}/{project}/{seq.type_}" + art_name = seq.name + + logger.info( + f"Deleting collection {entity=}, {project=}, {art_type=}, {art_name=}" + ) + try: + dst_collection = self.dst_api.artifact_collection(art_type, art_name) + except (wandb.CommError, ValueError): + logger.warning(f"Collection doesn't exist {art_type=}, {art_name=}") + return + + try: + dst_collection.delete() + except (wandb.CommError, ValueError) as e: + logger.warning( + f"Collection can't be deleted, {art_type=}, {art_name=}, {e=}" + ) + return + + def _import_artifact_sequence( + self, + seq: ArtifactSequence, + *, + namespace: Optional[Namespace] = None, + ) -> None: + """Import one artifact sequence. + + Use `namespace` to specify alternate settings like where the artifact sequence should be uploaded + """ + if not seq.artifacts: + # The artifact sequence has no versions. This usually means all artifacts versions were deleted intentionally, + # but it can also happen if the sequence represents run history and that run was deleted. + logger.warning(f"Artifact {seq=} has no artifacts, skipping.") + return + + if namespace is None: + namespace = Namespace(seq.entity, seq.project) + + settings_override = { + "api_key": self.dst_api_key, + "base_url": self.dst_base_url, + "resume": "true", + "resumed": True, + } + + send_manager_config = internal.SendManagerConfig(log_artifacts=True) + + # Delete any existing artifact sequence, otherwise versions will be out of order + # Unfortunately, you can't delete only part of the sequence because versions are "remembered" even after deletion + self._delete_collection_in_dst(seq, namespace) + + # Get a placeholder run for dummy artifacts we'll upload later + art = seq.artifacts[0] + run_or_dummy: Optional[Run] = _get_run_or_dummy_from_art(art, self.src_api) + + # Each `group_of_artifacts` is either: + # 1. A single "real" artifact in a list; or + # 2. A list of dummy artifacts that are uploaded together. + # This guarantees the real artifacts have the correct version numbers while allowing for parallel upload of dummies. + groups_of_artifacts = list(_make_groups_of_artifacts(seq)) + for i, group in enumerate(groups_of_artifacts, 1): + art = group[0] + if art.description == ART_SEQUENCE_DUMMY_PLACEHOLDER: + run = WandbRun(run_or_dummy, **self.run_api_kwargs) + else: + try: + wandb_run = art.logged_by() + except ValueError: + # The run used to exist but has since been deleted + # wandb_run = None + pass + + # Could be logged by None (rare) or ValueError + if wandb_run is None: + logger.warning( + f"Run for {art.name=} does not exist (deleted?), using {run_or_dummy=}" + ) + wandb_run = run_or_dummy + + new_art = _clone_art(art) + group = [new_art] + run = WandbRun(wandb_run, **self.run_api_kwargs) + + logger.info( + f"Uploading partial artifact {seq=}, {i}/{len(groups_of_artifacts)}" + ) + internal.send_run( + run, + extra_arts=group, + overrides=namespace.send_manager_overrides, + settings_override=settings_override, + config=send_manager_config, + ) + logger.info(f"Finished uploading {seq=}") + + # query it back and remove placeholders + self._remove_placeholders(seq) + + def _remove_placeholders(self, seq: ArtifactSequence) -> None: + try: + retry_arts_func = internal.exp_retry(self._dst_api.artifacts) + dst_arts = list(retry_arts_func(seq.type_, seq.name)) + except wandb.CommError: + logger.warning( + f"{seq=} does not exist in dst. Has it already been deleted?" + ) + return + except TypeError: + logger.exception("Problem getting dst versions (try again later).") + return + + for art in dst_arts: + if art.description != ART_SEQUENCE_DUMMY_PLACEHOLDER: + continue + if art.type in ("wandb-history", "job"): + continue + + try: + art.delete(delete_aliases=True) + except wandb.CommError as e: + if "cannot delete system managed artifact" in str(e): + logger.warning("Cannot delete system managed artifact") + else: + raise + + def _get_dst_art( + self, src_art: Run, entity: Optional[str] = None, project: Optional[str] = None + ) -> Artifact: + entity = coalesce(entity, src_art.entity) + project = coalesce(project, src_art.project) + name = src_art.name + + return self.dst_api._artifact(f"{entity}/{project}/{name}") + + def _get_run_problems( + self, src_run: Run, dst_run: Run, force_retry: bool = False + ) -> List[dict]: + problems = [] + + if force_retry: + problems.append("__force_retry__") + + if non_matching_metadata := self._compare_run_metadata(src_run, dst_run): + problems.append("metadata:" + str(non_matching_metadata)) + + if non_matching_summary := self._compare_run_summary(src_run, dst_run): + problems.append("summary:" + str(non_matching_summary)) + + # TODO: Compare files? + + return problems + + def _compare_run_metadata(self, src_run: Run, dst_run: Run) -> dict: + fname = "wandb-metadata.json" + # problems = {} + + src_f = src_run.file(fname) + if src_f.size == 0: + # the src was corrupted so no comparisons here will ever work + return {} + + dst_f = dst_run.file(fname) + try: + contents = wandb.util.download_file_into_memory( + dst_f.url, self.dst_api.api_key + ) + except urllib3.exceptions.ReadTimeoutError: + return {"Error checking": "Timeout"} + except requests.HTTPError as e: + if e.response.status_code == 404: + return {"Bad upload": f"File not found: {fname}"} + return {"http problem": f"{fname}: ({e})"} + + dst_meta = wandb.wandb_sdk.lib.json_util.loads(contents) + + non_matching = {} + if src_run.metadata: + for k, src_v in src_run.metadata.items(): + if k not in dst_meta: + non_matching[k] = {"src": src_v, "dst": "KEY NOT FOUND"} + continue + dst_v = dst_meta[k] + if src_v != dst_v: + non_matching[k] = {"src": src_v, "dst": dst_v} + + return non_matching + + def _compare_run_summary(self, src_run: Run, dst_run: Run) -> dict: + non_matching = {} + for k, src_v in src_run.summary.items(): + # These won't match between systems and that's ok + if isinstance(src_v, str) and src_v.startswith("wandb-client-artifact://"): + continue + if k in ("_wandb", "_runtime"): + continue + + src_v = _recursive_cast_to_dict(src_v) + + dst_v = dst_run.summary.get(k) + dst_v = _recursive_cast_to_dict(dst_v) + + if isinstance(src_v, dict) and isinstance(dst_v, dict): + for kk, sv in src_v.items(): + # These won't match between systems and that's ok + if isinstance(sv, str) and sv.startswith( + "wandb-client-artifact://" + ): + continue + dv = dst_v.get(kk) + if not _almost_equal(sv, dv): + non_matching[f"{k}-{kk}"] = {"src": sv, "dst": dv} + else: + if not _almost_equal(src_v, dst_v): + non_matching[k] = {"src": src_v, "dst": dst_v} + + return non_matching + + def _collect_failed_artifact_sequences(self) -> Iterable[ArtifactSequence]: + if (df := _read_ndjson(ARTIFACT_ERRORS_FNAME)) is None: + logger.debug(f"{ARTIFACT_ERRORS_FNAME=} is empty, returning nothing") + return + + unique_failed_sequences = df[ + ["src_entity", "src_project", "name", "type"] + ].unique() + + for row in unique_failed_sequences.iter_rows(named=True): + entity = row["src_entity"] + project = row["src_project"] + name = row["name"] + _type = row["type"] + + art_name = f"{entity}/{project}/{name}" + arts = self.src_api.artifacts(_type, art_name) + arts = sorted(arts, key=lambda a: int(a.version.lstrip("v"))) + arts = sorted(arts, key=lambda a: a.type) + + yield ArtifactSequence(arts, entity, project, _type, name) + + def _cleanup_dummy_runs( + self, + *, + namespaces: Optional[Iterable[Namespace]] = None, + api: Optional[Api] = None, + remapping: Optional[Dict[Namespace, Namespace]] = None, + ) -> None: + api = coalesce(api, self.dst_api) + namespaces = coalesce(namespaces, self._all_namespaces()) + + for ns in namespaces: + if remapping and ns in remapping: + ns = remapping[ns] + + logger.debug(f"Cleaning up, {ns=}") + try: + runs = list( + api.runs(ns.path, filters={"displayName": RUN_DUMMY_PLACEHOLDER}) + ) + except ValueError as e: + if "Could not find project" in str(e): + logger.exception("Could not find project, does it exist?") + continue + + for run in runs: + logger.debug(f"Deleting dummy {run=}") + run.delete(delete_artifacts=False) + + def _import_report( + self, report: Report, *, namespace: Optional[Namespace] = None + ) -> None: + """Import one wandb.Report. + + Use `namespace` to specify alternate settings like where the report should be uploaded + """ + if namespace is None: + namespace = Namespace(report.entity, report.project) + + entity = coalesce(namespace.entity, report.entity) + project = coalesce(namespace.project, report.project) + name = report.name + title = report.title + description = report.description + + api = self.dst_api + + # We shouldn't need to upsert the project for every report + logger.debug(f"Upserting {entity=}/{project=}") + try: + api.create_project(project, entity) + except requests.exceptions.HTTPError as e: + if e.response.status_code != 409: + logger.warning(f"Issue upserting {entity=}/{project=}, {e=}") + + logger.debug(f"Upserting report {entity=}, {project=}, {name=}, {title=}") + api.client.execute( + wr.report.UPSERT_VIEW, + variable_values={ + "id": None, # Is there any benefit for this to be the same as default report? + "name": name, + "entityName": entity, + "projectName": project, + "description": description, + "displayName": title, + "type": "runs", + "spec": json.dumps(report.spec), + }, + ) + + def _use_artifact_sequence( + self, + sequence: ArtifactSequence, + *, + namespace: Optional[Namespace] = None, + ): + if namespace is None: + namespace = Namespace(sequence.entity, sequence.project) + + settings_override = { + "api_key": self.dst_api_key, + "base_url": self.dst_base_url, + "resume": "true", + "resumed": True, + } + logger.debug(f"Using artifact sequence with {settings_override=}, {namespace=}") + + send_manager_config = internal.SendManagerConfig(use_artifacts=True) + + for art in sequence: + if (used_by := art.used_by()) is None: + continue + + for wandb_run in used_by: + run = WandbRun(wandb_run, **self.run_api_kwargs) + + internal.send_run( + run, + overrides=namespace.send_manager_overrides, + settings_override=settings_override, + config=send_manager_config, + ) + + def import_runs( + self, + *, + namespaces: Optional[Iterable[Namespace]] = None, + remapping: Optional[Dict[Namespace, Namespace]] = None, + parallel: bool = True, + incremental: bool = True, + max_workers: Optional[int] = None, + limit: Optional[int] = None, + metadata: bool = True, + files: bool = True, + media: bool = True, + code: bool = True, + history: bool = True, + summary: bool = True, + terminal_output: bool = True, + ): + logger.info("START: Import runs") + + logger.info("Setting up for import") + _create_files_if_not_exists() + _clear_fname(RUN_ERRORS_FNAME) + + logger.info("Collecting runs") + runs = list(self._collect_runs(namespaces=namespaces, limit=limit)) + + logger.info(f"Validating runs, {len(runs)=}") + self._validate_runs( + runs, + skip_previously_validated=incremental, + remapping=remapping, + ) + + logger.info("Collecting failed runs") + runs = list(self._collect_failed_runs()) + + logger.info(f"Importing runs, {len(runs)=}") + + def _import_run_wrapped(run): + namespace = Namespace(run.entity(), run.project()) + if remapping is not None and namespace in remapping: + namespace = remapping[namespace] + + config = internal.SendManagerConfig( + metadata=metadata, + files=files, + media=media, + code=code, + history=history, + summary=summary, + terminal_output=terminal_output, + ) + + logger.debug(f"Importing {run=}, {namespace=}, {config=}") + self._import_run(run, namespace=namespace, config=config) + logger.debug(f"Finished importing {run=}, {namespace=}, {config=}") + + for_each(_import_run_wrapped, runs, max_workers=max_workers, parallel=parallel) + logger.info("END: Importing runs") + + def import_reports( + self, + *, + namespaces: Optional[Iterable[Namespace]] = None, + limit: Optional[int] = None, + remapping: Optional[Dict[Namespace, Namespace]] = None, + ): + logger.info("START: Importing reports") + + logger.info("Collecting reports") + reports = self._collect_reports(namespaces=namespaces, limit=limit) + + logger.info("Importing reports") + + def _import_report_wrapped(report): + namespace = Namespace(report.entity, report.project) + if remapping is not None and namespace in remapping: + namespace = remapping[namespace] + + logger.debug(f"Importing {report=}, {namespace=}") + self._import_report(report, namespace=namespace) + logger.debug(f"Finished importing {report=}, {namespace=}") + + for_each(_import_report_wrapped, reports) + + logger.info("END: Importing reports") + + def import_artifact_sequences( + self, + *, + namespaces: Optional[Iterable[Namespace]] = None, + incremental: bool = True, + max_workers: Optional[int] = None, + remapping: Optional[Dict[Namespace, Namespace]] = None, + ): + """Import all artifact sequences from `namespaces`. + + Note: There is a known bug with the AWS backend where artifacts > 2048MB will fail to upload. This seems to be related to multipart uploads, but we don't have a fix yet. + """ + logger.info("START: Importing artifact sequences") + _clear_fname(ARTIFACT_ERRORS_FNAME) + + logger.info("Collecting artifact sequences") + seqs = list(self._collect_artifact_sequences(namespaces=namespaces)) + + logger.info("Validating artifact sequences") + self._validate_artifact_sequences( + seqs, + incremental=incremental, + remapping=remapping, + ) + + logger.info("Collecting failed artifact sequences") + seqs = list(self._collect_failed_artifact_sequences()) + + logger.info(f"Importing artifact sequences, {len(seqs)=}") + + def _import_artifact_sequence_wrapped(seq): + namespace = Namespace(seq.entity, seq.project) + if remapping is not None and namespace in remapping: + namespace = remapping[namespace] + + logger.debug(f"Importing artifact sequence {seq=}, {namespace=}") + self._import_artifact_sequence(seq, namespace=namespace) + logger.debug(f"Finished importing artifact sequence {seq=}, {namespace=}") + + for_each(_import_artifact_sequence_wrapped, seqs, max_workers=max_workers) + + # it's safer to just use artifact on all seqs to make sure we don't miss anything + # For seqs that have already been used, this is a no-op. + logger.debug(f"Using artifact sequences, {len(seqs)=}") + + def _use_artifact_sequence_wrapped(seq): + namespace = Namespace(seq.entity, seq.project) + if remapping is not None and namespace in remapping: + namespace = remapping[namespace] + + logger.debug(f"Using artifact sequence {seq=}, {namespace=}") + self._use_artifact_sequence(seq, namespace=namespace) + logger.debug(f"Finished using artifact sequence {seq=}, {namespace=}") + + for_each(_use_artifact_sequence_wrapped, seqs, max_workers=max_workers) + + # Artifacts whose parent runs have been deleted should have that run deleted in the + # destination as well + + logger.info("Cleaning up dummy runs") + self._cleanup_dummy_runs( + namespaces=namespaces, + remapping=remapping, + ) + + logger.info("END: Importing artifact sequences") + + def import_all( + self, + *, + runs: bool = True, + artifacts: bool = True, + reports: bool = True, + namespaces: Optional[Iterable[Namespace]] = None, + incremental: bool = True, + remapping: Optional[Dict[Namespace, Namespace]] = None, + ): + logger.info(f"START: Importing all, {runs=}, {artifacts=}, {reports=}") + if runs: + self.import_runs( + namespaces=namespaces, + incremental=incremental, + remapping=remapping, + ) + + if reports: + self.import_reports( + namespaces=namespaces, + remapping=remapping, + ) + + if artifacts: + self.import_artifact_sequences( + namespaces=namespaces, + incremental=incremental, + remapping=remapping, + ) + + logger.info("END: Importing all") + + def _validate_run( + self, + src_run: Run, + *, + remapping: Optional[Dict[Namespace, Namespace]] = None, + ) -> None: + namespace = Namespace(src_run.entity, src_run.project) + if remapping is not None and namespace in remapping: + namespace = remapping[namespace] + + dst_entity = namespace.entity + dst_project = namespace.project + run_id = src_run.id + + try: + dst_run = self.dst_api.run(f"{dst_entity}/{dst_project}/{run_id}") + except wandb.CommError: + problems = [f"run does not exist in dst at {dst_entity=}/{dst_project=}"] + else: + problems = self._get_run_problems(src_run, dst_run) + + d = { + "src_entity": src_run.entity, + "src_project": src_run.project, + "dst_entity": dst_entity, + "dst_project": dst_project, + "run_id": run_id, + } + if problems: + d["problems"] = problems + fname = RUN_ERRORS_FNAME + else: + fname = RUN_SUCCESSES_FNAME + + with filelock.FileLock("runs.lock"): + with open(fname, "a") as f: + f.write(json.dumps(d) + "\n") + + def _filter_previously_checked_runs( + self, + runs: Iterable[Run], + *, + remapping: Optional[Dict[Namespace, Namespace]] = None, + ) -> Iterable[Run]: + if (df := _read_ndjson(RUN_SUCCESSES_FNAME)) is None: + logger.debug(f"{RUN_SUCCESSES_FNAME=} is empty, yielding all runs") + yield from runs + return + + data = [] + for r in runs: + namespace = Namespace(r.entity, r.project) + if remapping is not None and namespace in remapping: + namespace = remapping[namespace] + + data.append( + { + "src_entity": r.entity, + "src_project": r.project, + "dst_entity": namespace.entity, + "dst_project": namespace.project, + "run_id": r.id, + "data": r, + } + ) + df2 = pl.DataFrame(data) + logger.debug(f"Starting with {len(runs)=} in namespaces") + + results = df2.join( + df, + how="anti", + on=["src_entity", "src_project", "dst_entity", "dst_project", "run_id"], + ) + logger.debug(f"After filtering out already successful runs, {len(results)=}") + + if not results.is_empty(): + results = results.filter(~results["run_id"].is_null()) + results = results.unique( + ["src_entity", "src_project", "dst_entity", "dst_project", "run_id"] + ) + + for r in results.iter_rows(named=True): + yield r["data"] + + def _validate_artifact( + self, + src_art: Artifact, + dst_entity: str, + dst_project: str, + download_files_and_compare: bool = False, + check_entries_are_downloadable: bool = True, + ): + problems = [] + + # These patterns of artifacts are special and should not be validated + ignore_patterns = [ + r"^job-(.*?)\.py(:v\d+)?$", + # r"^run-.*-history(?:\:v\d+)?$$", + ] + for pattern in ignore_patterns: + if re.search(pattern, src_art.name): + return (src_art, dst_entity, dst_project, problems) + + try: + dst_art = self._get_dst_art(src_art, dst_entity, dst_project) + except Exception: + problems.append("destination artifact not found") + return (src_art, dst_entity, dst_project, problems) + + try: + logger.debug("Comparing artifact manifests") + except Exception as e: + problems.append( + f"Problem getting problems! problem with {src_art.entity=}, {src_art.project=}, {src_art.name=} {e=}" + ) + else: + problems += validation._compare_artifact_manifests(src_art, dst_art) + + if check_entries_are_downloadable: + # validation._check_entries_are_downloadable(src_art) + validation._check_entries_are_downloadable(dst_art) + + if download_files_and_compare: + logger.debug(f"Downloading {src_art=}") + try: + src_dir = _download_art(src_art, root=f"{SRC_ART_PATH}/{src_art.name}") + except requests.HTTPError as e: + problems.append( + f"Invalid download link for src {src_art.entity=}, {src_art.project=}, {src_art.name=}, {e}" + ) + + logger.debug(f"Downloading {dst_art=}") + try: + dst_dir = _download_art(dst_art, root=f"{DST_ART_PATH}/{dst_art.name}") + except requests.HTTPError as e: + problems.append( + f"Invalid download link for dst {dst_art.entity=}, {dst_art.project=}, {dst_art.name=}, {e}" + ) + else: + logger.debug(f"Comparing artifact dirs {src_dir=}, {dst_dir=}") + if problem := validation._compare_artifact_dirs(src_dir, dst_dir): + problems.append(problem) + + return (src_art, dst_entity, dst_project, problems) + + def _validate_runs( + self, + runs: Iterable[WandbRun], + *, + skip_previously_validated: bool = True, + remapping: Optional[Dict[Namespace, Namespace]] = None, + ): + base_runs = [r.run for r in runs] + if skip_previously_validated: + base_runs = list( + self._filter_previously_checked_runs( + base_runs, + remapping=remapping, + ) + ) + + def _validate_run(run): + logger.debug(f"Validating {run=}") + self._validate_run(run, remapping=remapping) + logger.debug(f"Finished validating {run=}") + + for_each(_validate_run, base_runs) + + def _collect_failed_runs(self): + if (df := _read_ndjson(RUN_ERRORS_FNAME)) is None: + logger.debug(f"{RUN_ERRORS_FNAME=} is empty, returning nothing") + return + + unique_failed_runs = df[ + ["src_entity", "src_project", "dst_entity", "dst_project", "run_id"] + ].unique() + + for row in unique_failed_runs.iter_rows(named=True): + src_entity = row["src_entity"] + src_project = row["src_project"] + # dst_entity = row["dst_entity"] + # dst_project = row["dst_project"] + run_id = row["run_id"] + + run = self.src_api.run(f"{src_entity}/{src_project}/{run_id}") + yield WandbRun(run, **self.run_api_kwargs) + + def _filter_previously_checked_artifacts(self, seqs: Iterable[ArtifactSequence]): + if (df := _read_ndjson(ARTIFACT_SUCCESSES_FNAME)) is None: + logger.info( + f"{ARTIFACT_SUCCESSES_FNAME=} is empty, yielding all artifact sequences" + ) + for seq in seqs: + yield from seq.artifacts + return + + for seq in seqs: + for art in seq: + try: + logged_by = _get_run_or_dummy_from_art(art, self.src_api) + except requests.HTTPError: + logger.exception(f"Failed to get run, skipping: {art=}") + continue + + if art.type == "wandb-history" and isinstance(logged_by, _DummyRun): + logger.debug(f"Skipping history artifact {art=}") + # We can never upload valid history for a deleted run, so skip it + continue + + entity = art.entity + project = art.project + _type = art.type + name, ver = _get_art_name_ver(art) + + filtered_df = df.filter( + (df["src_entity"] == entity) + & (df["src_project"] == project) + & (df["name"] == name) + & (df["version"] == ver) + & (df["type"] == _type) + ) + + # not in file, so not verified yet, don't filter out + if len(filtered_df) == 0: + yield art + + def _validate_artifact_sequences( + self, + seqs: Iterable[ArtifactSequence], + *, + incremental: bool = True, + download_files_and_compare: bool = False, + check_entries_are_downloadable: bool = True, + remapping: Optional[Dict[Namespace, Namespace]] = None, + ): + if incremental: + logger.info("Validating in incremental mode") + + def filtered_sequences(): + for seq in seqs: + if not seq.artifacts: + continue + + art = seq.artifacts[0] + try: + logged_by = _get_run_or_dummy_from_art(art, self.src_api) + except requests.HTTPError: + logger.exception( + f"Validate Artifact http error: {art.entity=}," + f" {art.project=}, {art.name=}" + ) + continue + + if art.type == "wandb-history" and isinstance(logged_by, _DummyRun): + # We can never upload valid history for a deleted run, so skip it + continue + + yield seq + + artifacts = self._filter_previously_checked_artifacts(filtered_sequences()) + else: + logger.info("Validating in non-incremental mode") + artifacts = [art for seq in seqs for art in seq.artifacts] + + def _validate_artifact_wrapped(args): + art, entity, project = args + if ( + remapping is not None + and (namespace := Namespace(entity, project)) in remapping + ): + remapped_ns = remapping[namespace] + entity = remapped_ns.entity + project = remapped_ns.project + + logger.debug(f"Validating {art=}, {entity=}, {project=}") + result = self._validate_artifact( + art, + entity, + project, + download_files_and_compare=download_files_and_compare, + check_entries_are_downloadable=check_entries_are_downloadable, + ) + logger.debug(f"Finished validating {art=}, {entity=}, {project=}") + return result + + args = ((art, art.entity, art.project) for art in artifacts) + art_problems = for_each(_validate_artifact_wrapped, args) + for art, dst_entity, dst_project, problems in art_problems: + name, ver = _get_art_name_ver(art) + d = { + "src_entity": art.entity, + "src_project": art.project, + "dst_entity": dst_entity, + "dst_project": dst_project, + "name": name, + "version": ver, + "type": art.type, + } + + if problems: + d["problems"] = problems + fname = ARTIFACT_ERRORS_FNAME + else: + fname = ARTIFACT_SUCCESSES_FNAME + + with open(fname, "a") as f: + f.write(json.dumps(d) + "\n") + + def _collect_runs( + self, + *, + namespaces: Optional[Iterable[Namespace]] = None, + limit: Optional[int] = None, + skip_ids: Optional[List[str]] = None, + start_date: Optional[str] = None, + api: Optional[Api] = None, + ) -> Iterable[WandbRun]: + api = coalesce(api, self.src_api) + namespaces = coalesce(namespaces, self._all_namespaces()) + + filters: Dict[str, Any] = {} + if skip_ids is not None: + filters["name"] = {"$nin": skip_ids} + if start_date is not None: + filters["createdAt"] = {"$gte": start_date} + + def _runs(): + for ns in namespaces: + logger.debug(f"Collecting runs from {ns=}") + for run in api.runs(ns.path, filters=filters): + yield WandbRun(run, **self.run_api_kwargs) + + runs = itertools.islice(_runs(), limit) + yield from runs + + def _all_namespaces( + self, *, entity: Optional[str] = None, api: Optional[Api] = None + ): + api = coalesce(api, self.src_api) + entity = coalesce(entity, api.default_entity) + projects = api.projects(entity) + for p in projects: + yield Namespace(p.entity, p.name) + + def _collect_reports( + self, + *, + namespaces: Optional[Iterable[Namespace]] = None, + limit: Optional[int] = None, + api: Optional[Api] = None, + ): + api = coalesce(api, self.src_api) + namespaces = coalesce(namespaces, self._all_namespaces()) + + wandb.login(key=self.src_api_key, host=self.src_base_url) + + def reports(): + for ns in namespaces: + for r in api.reports(ns.path): + yield wr.Report.from_url(r.url, api=api) + + yield from itertools.islice(reports(), limit) + + def _collect_artifact_sequences( + self, + *, + namespaces: Optional[Iterable[Namespace]] = None, + limit: Optional[int] = None, + api: Optional[Api] = None, + ): + api = coalesce(api, self.src_api) + namespaces = coalesce(namespaces, self._all_namespaces()) + + def artifact_sequences(): + for ns in namespaces: + logger.debug(f"Collecting artifact sequences from {ns=}") + types = [] + try: + types = [t for t in api.artifact_types(ns.path)] + except Exception: + logger.exception("Failed to get artifact types.") + + for t in types: + collections = [] + + # Skip history because it's really for run history + if t.name == "wandb-history": + continue + + try: + collections = t.collections() + except Exception: + logger.exception("Failed to get artifact collections.") + + for c in collections: + if c.is_sequence(): + yield ArtifactSequence.from_collection(c) + + seqs = itertools.islice(artifact_sequences(), limit) + unique_sequences = {seq.identifier: seq for seq in seqs} + yield from unique_sequences.values() + + +def _get_art_name_ver(art: Artifact) -> Tuple[str, int]: + name, ver = art.name.split(":v") + return name, int(ver) + + +def _make_dummy_art(name: str, _type: str, ver: int): + art = Artifact(name, ART_DUMMY_PLACEHOLDER_TYPE) + art._type = _type + art._description = ART_SEQUENCE_DUMMY_PLACEHOLDER + + p = Path(ART_DUMMY_PLACEHOLDER_PATH) + p.mkdir(parents=True, exist_ok=True) + + # dummy file with different name to prevent dedupe + fname = p / str(ver) + with open(fname, "w"): + pass + art.add_file(fname) + + return art + + +def _make_groups_of_artifacts(seq: ArtifactSequence, start: int = 0): + prev_ver = start - 1 + for art in seq: + name, ver = _get_art_name_ver(art) + + # If there's a gap between versions, fill with dummy artifacts + if ver - prev_ver > 1: + yield [_make_dummy_art(name, art.type, v) for v in range(prev_ver + 1, ver)] + + # Then yield the actual artifact + # Must always be a list of one artifact to guarantee ordering + yield [art] + prev_ver = ver + + +def _recursive_cast_to_dict(obj): + if isinstance(obj, list): + return [_recursive_cast_to_dict(item) for item in obj] + elif isinstance(obj, dict) or hasattr(obj, "items"): + new_dict = {} + for key, value in obj.items(): + new_dict[key] = _recursive_cast_to_dict(value) + return new_dict + else: + return obj + + +def _almost_equal(x, y, eps=1e-6): + if isinstance(x, dict) and isinstance(y, dict): + if x.keys() != y.keys(): + return False + return all(_almost_equal(x[k], y[k], eps) for k in x) + + if isinstance(x, numbers.Number) and isinstance(y, numbers.Number): + return abs(x - y) < eps + + if type(x) is not type(y): + return False + + return x == y + + +@dataclass +class _DummyUser: + username: str = "" + + +@dataclass +class _DummyRun: + entity: str = "" + project: str = "" + run_id: str = RUN_DUMMY_PLACEHOLDER + id: str = RUN_DUMMY_PLACEHOLDER + display_name: str = RUN_DUMMY_PLACEHOLDER + notes: str = "" + url: str = "" + group: str = "" + created_at: str = "2000-01-01" + user: _DummyUser = field(default_factory=_DummyUser) + tags: list = field(default_factory=list) + summary: dict = field(default_factory=dict) + config: dict = field(default_factory=dict) + + def files(self): + return [] + + +def _read_ndjson(fname: str) -> Optional[pl.DataFrame]: + try: + df = pl.read_ndjson(fname) + except FileNotFoundError: + return None + except RuntimeError as e: + # No runs previously checked + if "empty string is not a valid JSON value" in str(e): + return None + if "error parsing ndjson" in str(e): + return None + raise + + return df + + +def _get_run_or_dummy_from_art(art: Artifact, api=None): + run = None + + try: + run = art.logged_by() + except ValueError as e: + logger.warning( + f"Can't log artifact because run doesn't exist, {art=}, {run=}, {e=}" + ) + + if run is not None: + return run + + query = gql( + """ + query ArtifactCreatedBy( + $id: ID! + ) { + artifact(id: $id) { + createdBy { + ... on Run { + name + project { + name + entityName + } + } + } + } + } + """ + ) + response = api.client.execute(query, variable_values={"id": art.id}) + creator = response.get("artifact", {}).get("createdBy", {}) + run = _DummyRun( + entity=art.entity, + project=art.project, + run_id=creator.get("name", RUN_DUMMY_PLACEHOLDER), + id=creator.get("name", RUN_DUMMY_PLACEHOLDER), + ) + return run + + +def _clear_fname(fname: str) -> None: + old_fname = f"{internal.ROOT_DIR}/{fname}" + new_fname = f"{internal.ROOT_DIR}/prev_{fname}" + + logger.debug(f"Moving {old_fname=} to {new_fname=}") + try: + shutil.copy2(old_fname, new_fname) + except FileNotFoundError: + # this is just to make a copy of the last iteration, so its ok if the src doesn't exist + pass + + with open(fname, "w"): + pass + + +def _download_art(art: Artifact, root: str) -> Optional[str]: + try: + with patch("click.echo"): + return art.download(root=root, skip_cache=True) + except Exception: + logger.exception(f"Error downloading artifact {art=}") + + +def _clone_art(art: Artifact, root: Optional[str] = None): + if root is None: + # Currently, we would only ever clone a src artifact to move it to dst. + root = f"{SRC_ART_PATH}/{art.name}" + + if (path := _download_art(art, root=root)) is None: + raise ValueError(f"Problem downloading {art=}") + + name, _ = art.name.split(":v") + + # Hack: skip naming validation check for wandb-* types + new_art = Artifact(name, ART_DUMMY_PLACEHOLDER_TYPE) + new_art._type = art.type + new_art._created_at = art.created_at + + new_art._aliases = art.aliases + new_art._description = art.description + + with patch("click.echo"): + new_art.add_dir(path) + + return new_art + + +def _create_files_if_not_exists() -> None: + fnames = [ + ARTIFACT_ERRORS_FNAME, + ARTIFACT_SUCCESSES_FNAME, + RUN_ERRORS_FNAME, + RUN_SUCCESSES_FNAME, + ] + + for fname in fnames: + logger.debug(f"Creating {fname=} if not exists") + with open(fname, "a"): + pass + + +def _merge_dfs(dfs: List[pl.DataFrame]) -> pl.DataFrame: + # Ensure there are DataFrames in the list + if len(dfs) == 0: + return pl.DataFrame() + + if len(dfs) == 1: + return dfs[0] + + merged_df = dfs[0] + for df in dfs[1:]: + merged_df = merged_df.join(df, how="outer", on=["_step"]) + col_pairs = [ + (c, f"{c}_right") + for c in merged_df.columns + if f"{c}_right" in merged_df.columns + ] + for col, right in col_pairs: + new_col = merged_df[col].fill_null(merged_df[right]) + merged_df = merged_df.with_columns(new_col).drop(right) + + return merged_df diff --git a/lib/python3.12/site-packages/wandb/apis/internal.py b/lib/python3.12/site-packages/wandb/apis/internal.py new file mode 100644 index 0000000000000000000000000000000000000000..4452213e2228324b561552c555e6a9359a45a102 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/internal.py @@ -0,0 +1,239 @@ +from typing import Any + +from wandb.sdk.internal.internal_api import Api as InternalApi + + +class Api: + """Internal proxy to the official internal API.""" + + # TODO: Move these methods to PublicApi. + + def __init__(self, *args: Any, **kwargs: Any) -> None: + self._api_args = args + self._api_kwargs = kwargs + self._api = None + + def __getstate__(self): + """Use for serializing. + + self._api is not serializable, so it's dropped + """ + state = self.__dict__.copy() + del state["_api"] + return state + + def __setstate__(self, state): + """Used for deserializing. + + Don't need to set self._api because it's constructed when needed. + """ + self.__dict__.update(state) + self._api = None + + @property + def api(self) -> InternalApi: + # This is a property in order to delay construction of Internal API + # for as long as possible. If constructed in constructor, then the + # whole InternalAPI is started when simply importing wandb. + if self._api is None: + self._api = InternalApi(*self._api_args, **self._api_kwargs) + return self._api + + @property + def api_key(self): + return self.api.api_key + + @property + def is_authenticated(self): + return self.api.access_token is not None or self.api.api_key is not None + + @property + def api_url(self): + return self.api.api_url + + @property + def app_url(self): + return self.api.app_url + + @property + def default_entity(self): + return self.api.default_entity + + @property + def git(self): + return self.api.git + + def validate_api_key(self) -> bool: + """Returns whether the API key stored on initialization is valid.""" + return self.api.validate_api_key() + + def file_current(self, *args): + return self.api.file_current(*args) + + def download_file(self, *args, **kwargs): + return self.api.download_file(*args, **kwargs) + + def download_write_file(self, *args, **kwargs): + return self.api.download_write_file(*args, **kwargs) + + def set_current_run_id(self, run_id): + return self.api.set_current_run_id(run_id) + + def viewer(self): + return self.api.viewer() + + def max_cli_version(self): + return self.api.max_cli_version() + + def viewer_server_info(self): + return self.api.viewer_server_info() + + def list_projects(self, entity=None): + return self.api.list_projects(entity=entity) + + def format_project(self, project): + return self.api.format_project(project) + + def upsert_project(self, project, id=None, description=None, entity=None): + return self.api.upsert_project( + project, id=id, description=description, entity=entity + ) + + def upsert_run(self, *args, **kwargs): + return self.api.upsert_run(*args, **kwargs) + + def settings(self, *args, **kwargs): + return self.api.settings(*args, **kwargs) + + def clear_setting( + self, key: str, globally: bool = False, persist: bool = False + ) -> None: + return self.api.clear_setting(key, globally, persist) + + def set_setting( + self, key: str, value: Any, globally: bool = False, persist: bool = False + ) -> None: + return self.api.set_setting(key, value, globally, persist) + + def parse_slug(self, *args, **kwargs): + return self.api.parse_slug(*args, **kwargs) + + def download_url(self, *args, **kwargs): + return self.api.download_url(*args, **kwargs) + + def download_urls(self, *args, **kwargs): + return self.api.download_urls(*args, **kwargs) + + def create_anonymous_api_key(self) -> str: + return self.api.create_anonymous_api_key() + + def push(self, *args, **kwargs): + return self.api.push(*args, **kwargs) + + def sweep(self, *args, **kwargs): + return self.api.sweep(*args, **kwargs) + + def upsert_sweep(self, *args, **kwargs): + return self.api.upsert_sweep(*args, **kwargs) + + def set_sweep_state(self, *args, **kwargs): + return self.api.set_sweep_state(*args, **kwargs) + + def get_sweep_state(self, *args, **kwargs): + return self.api.get_sweep_state(*args, **kwargs) + + def stop_sweep(self, *args, **kwargs): + return self.api.stop_sweep(*args, **kwargs) + + def cancel_sweep(self, *args, **kwargs): + return self.api.cancel_sweep(*args, **kwargs) + + def pause_sweep(self, *args, **kwargs): + return self.api.pause_sweep(*args, **kwargs) + + def resume_sweep(self, *args, **kwargs): + return self.api.resume_sweep(*args, **kwargs) + + def register_agent(self, *args, **kwargs): + return self.api.register_agent(*args, **kwargs) + + def agent_heartbeat(self, *args, **kwargs): + return self.api.agent_heartbeat(*args, **kwargs) + + def use_artifact(self, *args, **kwargs): + return self.api.use_artifact(*args, **kwargs) + + def create_artifact(self, *args, **kwargs): + return self.api.create_artifact(*args, **kwargs) + + def complete_multipart_upload_artifact(self, *args, **kwargs): + return self.api.complete_multipart_upload_artifact(*args, **kwargs) + + def run_config(self, *args, **kwargs): + return self.api.run_config(*args, **kwargs) + + def upload_file_retry(self, *args, **kwargs): + return self.api.upload_file_retry(*args, **kwargs) + + def upload_multipart_file_chunk_retry(self, *args, **kwargs): + return self.api.upload_multipart_file_chunk_retry(*args, **kwargs) + + def get_run_info(self, *args, **kwargs): + return self.api.get_run_info(*args, **kwargs) + + def get_run_state(self, *args, **kwargs): + return self.api.get_run_state(*args, **kwargs) + + def entity_is_team(self, *args, **kwargs): + return self.api.entity_is_team(*args, **kwargs) + + def get_project_run_queues(self, *args, **kwargs): + return self.api.get_project_run_queues(*args, **kwargs) + + def push_to_run_queue(self, *args, **kwargs): + return self.api.push_to_run_queue(*args, **kwargs) + + def pop_from_run_queue(self, *args, **kwargs): + return self.api.pop_from_run_queue(*args, **kwargs) + + def ack_run_queue_item(self, *args, **kwargs): + return self.api.ack_run_queue_item(*args, **kwargs) + + def create_launch_agent(self, *args, **kwargs): + return self.api.create_launch_agent(*args, **kwargs) + + def create_default_resource_config(self, *args, **kwargs): + return self.api.create_default_resource_config(*args, **kwargs) + + def create_run_queue(self, *args, **kwargs): + return self.api.create_run_queue(*args, **kwargs) + + def upsert_run_queue(self, *args, **kwargs): + return self.api.upsert_run_queue(*args, **kwargs) + + def create_custom_chart(self, *args, **kwargs): + return self.api.create_custom_chart(*args, **kwargs) + + def update_launch_agent_status(self, *args, **kwargs): + return self.api.update_launch_agent_status(*args, **kwargs) + + def launch_agent_introspection(self, *args, **kwargs): + return self.api.launch_agent_introspection(*args, **kwargs) + + def fail_run_queue_item_introspection(self, *args, **kwargs): + return self.api.fail_run_queue_item_introspection(*args, **kwargs) + + def fail_run_queue_item(self, *args, **kwargs): + return self.api.fail_run_queue_item(*args, **kwargs) + + def update_run_queue_item_warning(self, *args, **kwargs): + return self.api.update_run_queue_item_warning(*args, **kwargs) + + def get_launch_agent(self, *args, **kwargs): + return self.api.get_launch_agent(*args, **kwargs) + + def stop_run(self, *args, **kwargs): + return self.api.stop_run(*args, **kwargs) + + +__all__ = ["Api"] diff --git a/lib/python3.12/site-packages/wandb/apis/normalize.py b/lib/python3.12/site-packages/wandb/apis/normalize.py new file mode 100644 index 0000000000000000000000000000000000000000..23327aecbf2d62492e267932a57f10abe056c892 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/normalize.py @@ -0,0 +1,81 @@ +"""normalize.""" + +import ast +import sys +from functools import wraps +from typing import Callable, TypeVar + +import requests +from wandb_gql.client import RetryError + +from wandb import env +from wandb.errors import CommError, Error +from wandb.util import parse_backend_error_messages + +_F = TypeVar("_F", bound=Callable) + + +def normalize_exceptions(func: _F) -> _F: + """Function decorator for catching common errors and re-raising as wandb.Error.""" + + @wraps(func) + def wrapper(*args, **kwargs): + message = "Whoa, you found a bug." + try: + return func(*args, **kwargs) + + except requests.HTTPError as error: + errors = parse_backend_error_messages(error.response) + status = error.response.status_code + + if errors: + message = f"HTTP {status}: {'; '.join(errors)}" + elif error.response.text: + message = f"HTTP {status}: {error.response.text}" + elif error.response.reason: + # Visually different to distinguish backend errors from + # standard HTTP status descriptions. + message = f"HTTP {status} ({error.response.reason})" + else: + message = f"HTTP {status}" + + raise CommError(message, error) + + except RetryError as err: + if ( + "response" in dir(err.last_exception) + and err.last_exception.response is not None + ): + try: + message = err.last_exception.response.json().get( + "errors", [{"message": message}] + )[0]["message"] + except ValueError: + message = err.last_exception.response.text + else: + message = err.last_exception + + if env.is_debug(): + raise err.last_exception.with_traceback(sys.exc_info()[2]) + else: + raise CommError(message, err.last_exception).with_traceback( + sys.exc_info()[2] + ) + except Error: + raise + except Exception as err: + # gql raises server errors with dict's as strings... + if len(err.args) > 0: + payload = err.args[0] + else: + payload = err + if str(payload).startswith("{"): + message = ast.literal_eval(str(payload))["message"] + else: + message = str(err) + if env.is_debug(): + raise + else: + raise CommError(message, err).with_traceback(sys.exc_info()[2]) + + return wrapper diff --git a/lib/python3.12/site-packages/wandb/apis/paginator.py b/lib/python3.12/site-packages/wandb/apis/paginator.py new file mode 100644 index 0000000000000000000000000000000000000000..d858b82aa0ecd63f22ff8273a046deec7d6d66ee --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/paginator.py @@ -0,0 +1,138 @@ +from __future__ import annotations + +from abc import abstractmethod +from typing import ( + TYPE_CHECKING, + Any, + ClassVar, + Iterator, + Mapping, + Protocol, + Sized, + TypeVar, + overload, +) + +import wandb + +if TYPE_CHECKING: + from wandb_graphql.language.ast import Document + +T = TypeVar("T") + + +# Structural type hint for the client instance +class _Client(Protocol): + def execute(self, *args: Any, **kwargs: Any) -> dict[str, Any]: ... + + +class Paginator(Iterator[T]): + """An iterator for paginated objects from GraphQL requests.""" + + QUERY: ClassVar[Document | None] = None + + def __init__( + self, + client: _Client, + variables: Mapping[str, Any], + per_page: int = 50, # We don't allow unbounded paging + ): + self.client: _Client = client + + # shallow copy partly guards against mutating the original input + self.variables: dict[str, Any] = dict(variables) + + self.per_page: int = per_page + self.objects: list[T] = [] + self.index: int = -1 + self.last_response: object | None = None + + def __iter__(self) -> Iterator[T]: + self.index = -1 + return self + + @property + @abstractmethod + def more(self) -> bool: + """Whether there are more pages to be fetched.""" + raise NotImplementedError + + @property + @abstractmethod + def cursor(self) -> str | None: + """The start cursor to use for the next fetched page.""" + raise NotImplementedError + + @abstractmethod + def convert_objects(self) -> list[T]: + """Convert the last fetched response data into the iterated objects.""" + raise NotImplementedError + + def update_variables(self) -> None: + """Update the query variables for the next page fetch.""" + self.variables.update({"perPage": self.per_page, "cursor": self.cursor}) + + def _update_response(self) -> None: + """Fetch and store the response data for the next page.""" + self.last_response = self.client.execute( + self.QUERY, variable_values=self.variables + ) + + def _load_page(self) -> bool: + """Fetch the next page, if any, returning True and storing the response if there was one.""" + if not self.more: + return False + self.update_variables() + self._update_response() + self.objects.extend(self.convert_objects()) + return True + + @overload + def __getitem__(self, index: int) -> T: ... + @overload + def __getitem__(self, index: slice) -> list[T]: ... + + def __getitem__(self, index: int | slice) -> T | list[T]: + loaded = True + stop = index.stop if isinstance(index, slice) else index + while loaded and stop > len(self.objects) - 1: + loaded = self._load_page() + return self.objects[index] + + def __next__(self) -> T: + self.index += 1 + if len(self.objects) <= self.index: + if not self._load_page(): + raise StopIteration + if len(self.objects) <= self.index: + raise StopIteration + return self.objects[self.index] + + next = __next__ + + +class SizedPaginator(Paginator[T], Sized): + """A Paginator for objects with a known total count.""" + + @property + def length(self) -> int | None: + wandb.termwarn( + ( + "`.length` is deprecated and will be removed in a future version. " + "Use `len(...)` instead." + ), + repeat=False, + ) + return len(self) + + def __len__(self) -> int: + if self._length is None: + self._load_page() + if self._length is None: + raise ValueError("Object doesn't provide length") + return self._length + + @property + @abstractmethod + def _length(self) -> int | None: + raise NotImplementedError diff --git a/lib/python3.12/site-packages/wandb/apis/public/__init__.py b/lib/python3.12/site-packages/wandb/apis/public/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b753463661d701fd2f7496de9e1c4767a4e6508a --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/public/__init__.py @@ -0,0 +1,35 @@ +from wandb.apis.public.api import Api, RetryingClient, requests +from wandb.apis.public.artifacts import ( + ArtifactCollection, + ArtifactCollections, + ArtifactFiles, + Artifacts, + ArtifactType, + ArtifactTypes, + RunArtifacts, +) +from wandb.apis.public.automations import Automations +from wandb.apis.public.files import FILE_FRAGMENT, File, Files +from wandb.apis.public.history import HistoryScan, SampledHistoryScan +from wandb.apis.public.integrations import SlackIntegrations, WebhookIntegrations +from wandb.apis.public.jobs import ( + Job, + QueuedRun, + RunQueue, + RunQueueAccessType, + RunQueuePrioritizationMode, + RunQueueResourceType, +) +from wandb.apis.public.projects import PROJECT_FRAGMENT, Project, Projects +from wandb.apis.public.query_generator import QueryGenerator +from wandb.apis.public.registries.registry import Registry +from wandb.apis.public.reports import ( + BetaReport, + PanelMetricsHelper, + PythonMongoishQueryGenerator, + Reports, +) +from wandb.apis.public.runs import RUN_FRAGMENT, Run, Runs +from wandb.apis.public.sweeps import SWEEP_FRAGMENT, Sweep +from wandb.apis.public.teams import Member, Team +from wandb.apis.public.users import User diff --git a/lib/python3.12/site-packages/wandb/apis/public/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1fcfb72925ae78fa2508f90724c099723c7bb1f0 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/public/__pycache__/api.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/api.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..268db76dbe375fb9606d4ec6fe60316d7e0266a5 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/api.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/public/__pycache__/artifacts.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/artifacts.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e927a6df20eaefe3e882e01a71847ea58caf42e7 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/artifacts.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/public/__pycache__/automations.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/automations.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a6008c85864cb6a552834fe7da86fe3ffe334189 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/automations.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/public/__pycache__/const.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/const.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bf9ade5dab88c5833059565abf9e889995b60402 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/const.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/public/__pycache__/files.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/files.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ff6b80108614e82414eded0993486cd4b660e7c7 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/files.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/public/__pycache__/history.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/history.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cd6486c7ba30f864b63b884825b6b4624ad38121 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/history.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/public/__pycache__/integrations.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/integrations.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..167447bb4610bf27bfbbaeda23fc0728cd204165 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/integrations.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/public/__pycache__/jobs.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/jobs.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b31f6d69a6f88e44638eaf9d2faf7a90a6bc5528 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/jobs.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/public/__pycache__/projects.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/projects.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..469a1102eb7bebd93549fd2e8281984fd607410d Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/projects.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/public/__pycache__/query_generator.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/query_generator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e0a2bc8dba95eb0c224c682fc035de6eb8b7c6f5 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/query_generator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/public/__pycache__/reports.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/reports.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4908c8f328504461547355e46b13a63fdf1b375d Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/reports.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/public/__pycache__/runs.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/runs.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b581eb3747d5f12666a9e0d479c4ac5012abf81e Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/runs.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/public/__pycache__/sweeps.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/sweeps.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..000f1618e1c3f9b71a54198447c109dcb3f58e44 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/sweeps.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/public/__pycache__/teams.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/teams.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..431aece4fbeaf032f3cb3c3b580216923bf3d3bb Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/teams.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/public/__pycache__/users.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/users.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1edb38a7032491f857383c28bf5372a102fe56e3 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/users.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/public/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..210f86263a77a22d5f8ed082d10a4fea12edd53a Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/public/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/public/api.py b/lib/python3.12/site-packages/wandb/apis/public/api.py new file mode 100644 index 0000000000000000000000000000000000000000..f380d7b39e1bc346b3ccf0f4050f201598964a5d --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/public/api.py @@ -0,0 +1,2284 @@ +"""Use the Public API to export or update data that you have saved to W&B. + +Before using this API, you'll want to log data from your script — check the +[Quickstart](https://docs.wandb.ai/quickstart) for more details. + +You might use the Public API to + - update metadata or metrics for an experiment after it has been completed, + - pull down your results as a dataframe for post-hoc analysis in a Jupyter notebook, or + - check your saved model artifacts for those tagged as `ready-to-deploy`. + +For more on using the Public API, check out [our guide](https://docs.wandb.com/guides/track/public-api-guide). +""" + +import json +import logging +import os +import urllib +from http import HTTPStatus +from typing import ( + TYPE_CHECKING, + Any, + Dict, + Iterator, + List, + Literal, + Optional, + Set, + Union, +) + +import requests +from pydantic import ValidationError +from typing_extensions import Unpack +from wandb_gql import Client, gql +from wandb_gql.client import RetryError + +import wandb +from wandb import env, util +from wandb._iterutils import one +from wandb.apis import public +from wandb.apis.normalize import normalize_exceptions +from wandb.apis.public.const import RETRY_TIMEDELTA +from wandb.apis.public.registries._utils import fetch_org_entity_from_organization +from wandb.apis.public.registries.registries_search import Registries +from wandb.apis.public.registries.registry import Registry +from wandb.apis.public.utils import ( + PathType, + fetch_org_from_settings_or_entity, + gql_compat, + parse_org_from_registry_path, +) +from wandb.proto.wandb_deprecated import Deprecated +from wandb.proto.wandb_internal_pb2 import ServerFeature +from wandb.sdk.artifacts._validators import is_artifact_registry_project +from wandb.sdk.internal.internal_api import Api as InternalApi +from wandb.sdk.internal.thread_local_settings import _thread_local_api_settings +from wandb.sdk.launch.utils import LAUNCH_DEFAULT_PROJECT +from wandb.sdk.lib import retry, runid +from wandb.sdk.lib.deprecate import deprecate +from wandb.sdk.lib.gql_request import GraphQLSession + +if TYPE_CHECKING: + from wandb.automations import ( + ActionType, + Automation, + EventType, + Integration, + NewAutomation, + SlackIntegration, + WebhookIntegration, + ) + from wandb.automations._utils import WriteAutomationsKwargs + +logger = logging.getLogger(__name__) + + +class RetryingClient: + INFO_QUERY = gql( + """ + query ServerInfo{ + serverInfo { + cliVersionInfo + latestLocalVersionInfo { + outOfDate + latestVersionString + versionOnThisInstanceString + } + } + } + """ + ) + + def __init__(self, client: Client): + self._server_info = None + self._client = client + + @property + def app_url(self): + return util.app_url(self._client.transport.url.replace("/graphql", "")) + "/" + + @retry.retriable( + retry_timedelta=RETRY_TIMEDELTA, + check_retry_fn=util.no_retry_auth, + retryable_exceptions=(RetryError, requests.RequestException), + ) + def execute( + self, *args, **kwargs + ): # User not encouraged to use this class directly + try: + return self._client.execute(*args, **kwargs) + except requests.exceptions.ReadTimeout: + if "timeout" not in kwargs: + timeout = self._client.transport.default_timeout + wandb.termwarn( + f"A graphql request initiated by the public wandb API timed out (timeout={timeout} sec). " + f"Create a new API with an integer timeout larger than {timeout}, e.g., `api = wandb.Api(timeout={timeout + 10})` " + f"to increase the graphql timeout." + ) + raise + + @property + def server_info(self): + if self._server_info is None: + self._server_info = self.execute(self.INFO_QUERY).get("serverInfo") + return self._server_info + + def version_supported( + self, min_version: str + ) -> bool: # User not encouraged to use this class directly + from packaging.version import parse + + return parse(min_version) <= parse( + self.server_info["cliVersionInfo"]["max_cli_version"] + ) + + +class Api: + """Used for querying the wandb server. + + Examples: + Most common way to initialize + >>> wandb.Api() + + Args: + overrides: (dict) You can set `base_url` if you are using a wandb server + other than https://api.wandb.ai. + You can also set defaults for `entity`, `project`, and `run`. + """ + + _HTTP_TIMEOUT = env.get_http_timeout(19) + DEFAULT_ENTITY_QUERY = gql( + """ + query Viewer{ + viewer { + id + entity + } + } + """ + ) + + VIEWER_QUERY = gql( + """ + query Viewer{ + viewer { + id + flags + entity + username + email + admin + apiKeys { + edges { + node { + id + name + description + } + } + } + teams { + edges { + node { + name + } + } + } + } + } + """ + ) + USERS_QUERY = gql( + """ + query SearchUsers($query: String) { + users(query: $query) { + edges { + node { + id + flags + entity + admin + email + deletedAt + username + apiKeys { + edges { + node { + id + name + description + } + } + } + teams { + edges { + node { + name + } + } + } + } + } + } + } + """ + ) + + CREATE_PROJECT = gql( + """ + mutation upsertModel( + $description: String + $entityName: String + $id: String + $name: String + $framework: String + $access: String + $views: JSONString + ) { + upsertModel( + input: { + description: $description + entityName: $entityName + id: $id + name: $name + framework: $framework + access: $access + views: $views + } + ) { + project { + id + name + entityName + description + access + views + } + model { + id + name + entityName + description + access + views + } + inserted + } + } + """ + ) + + def __init__( + self, + overrides: Optional[Dict[str, Any]] = None, + timeout: Optional[int] = None, + api_key: Optional[str] = None, + ) -> None: + self.settings = InternalApi().settings() + + _overrides = overrides or {} + self.settings.update(_overrides) + self.settings["base_url"] = self.settings["base_url"].rstrip("/") + if "organization" in _overrides: + self.settings["organization"] = _overrides["organization"] + if "username" in _overrides and "entity" not in _overrides: + wandb.termwarn( + 'Passing "username" to Api is deprecated. please use "entity" instead.' + ) + self.settings["entity"] = _overrides["username"] + + self._api_key = api_key + if self.api_key is None and _thread_local_api_settings.cookies is None: + wandb.login(host=_overrides.get("base_url")) + + self._viewer = None + self._projects = {} + self._runs = {} + self._sweeps = {} + self._reports = {} + self._default_entity = None + self._timeout = timeout if timeout is not None else self._HTTP_TIMEOUT + auth = None + if not _thread_local_api_settings.cookies: + auth = ("api", self.api_key) + proxies = self.settings.get("_proxies") or json.loads( + os.environ.get("WANDB__PROXIES", "{}") + ) + self._base_client = Client( + transport=GraphQLSession( + headers={ + "User-Agent": self.user_agent, + "Use-Admin-Privileges": "true", + **(_thread_local_api_settings.headers or {}), + }, + use_json=True, + # this timeout won't apply when the DNS lookup fails. in that case, it will be 60s + # https://bugs.python.org/issue22889 + timeout=self._timeout, + auth=auth, + url="{}/graphql".format(self.settings["base_url"]), + cookies=_thread_local_api_settings.cookies, + proxies=proxies, + ) + ) + self._client = RetryingClient(self._base_client) + + def create_project(self, name: str, entity: str) -> None: + """Create a new project. + + Args: + name: (str) The name of the new project. + entity: (str) The entity of the new project. + """ + self.client.execute(self.CREATE_PROJECT, {"entityName": entity, "name": name}) + + def create_run( + self, + *, + run_id: Optional[str] = None, + project: Optional[str] = None, + entity: Optional[str] = None, + ) -> "public.Run": + """Create a new run. + + Args: + run_id: (str, optional) The ID to assign to the run, if given. The run ID is automatically generated by + default, so in general, you do not need to specify this and should only do so at your own risk. + project: (str, optional) If given, the project of the new run. + entity: (str, optional) If given, the entity of the new run. + + Returns: + The newly created `Run`. + """ + if entity is None: + entity = self.default_entity + return public.Run.create(self, run_id=run_id, project=project, entity=entity) + + def create_run_queue( + self, + name: str, + type: "public.RunQueueResourceType", + entity: Optional[str] = None, + prioritization_mode: Optional["public.RunQueuePrioritizationMode"] = None, + config: Optional[dict] = None, + template_variables: Optional[dict] = None, + ) -> "public.RunQueue": + """Create a new run queue (launch). + + Args: + name: (str) Name of the queue to create + type: (str) Type of resource to be used for the queue. One of "local-container", "local-process", "kubernetes", "sagemaker", or "gcp-vertex". + entity: (str) Optional name of the entity to create the queue. If None, will use the configured or default entity. + prioritization_mode: (str) Optional version of prioritization to use. Either "V0" or None + config: (dict) Optional default resource configuration to be used for the queue. Use handlebars (eg. `{{var}}`) to specify template variables. + template_variables: (dict) A dictionary of template variable schemas to be used with the config. Expected format of: + `{ + "var-name": { + "schema": { + "type": ("string", "number", or "integer"), + "default": (optional value), + "minimum": (optional minimum), + "maximum": (optional maximum), + "enum": [..."(options)"] + } + } + }` + + Returns: + The newly created `RunQueue` + + Raises: + ValueError if any of the parameters are invalid + wandb.Error on wandb API errors + """ + # TODO(np): Need to check server capabilities for this feature + # 0. assert params are valid/normalized + if entity is None: + entity = self.settings["entity"] or self.default_entity + if entity is None: + raise ValueError( + "entity must be passed as a parameter, or set in settings" + ) + + if len(name) == 0: + raise ValueError("name must be non-empty") + if len(name) > 64: + raise ValueError("name must be less than 64 characters") + + if type not in [ + "local-container", + "local-process", + "kubernetes", + "sagemaker", + "gcp-vertex", + ]: + raise ValueError( + "resource_type must be one of 'local-container', 'local-process', 'kubernetes', 'sagemaker', or 'gcp-vertex'" + ) + + if prioritization_mode: + prioritization_mode = prioritization_mode.upper() + if prioritization_mode not in ["V0"]: + raise ValueError("prioritization_mode must be 'V0' if present") + + if config is None: + config = {} + + # 1. create required default launch project in the entity + self.create_project(LAUNCH_DEFAULT_PROJECT, entity) + + api = InternalApi( + default_settings={ + "entity": entity, + "project": self.project(LAUNCH_DEFAULT_PROJECT), + }, + retry_timedelta=RETRY_TIMEDELTA, + ) + + # 2. create default resource config, receive config id + config_json = json.dumps({"resource_args": {type: config}}) + create_config_result = api.create_default_resource_config( + entity, type, config_json, template_variables + ) + if not create_config_result["success"]: + raise wandb.Error("failed to create default resource config") + config_id = create_config_result["defaultResourceConfigID"] + + # 3. create run queue + create_queue_result = api.create_run_queue( + entity, + LAUNCH_DEFAULT_PROJECT, + name, + "PROJECT", + prioritization_mode, + config_id, + ) + if not create_queue_result["success"]: + raise wandb.Error("failed to create run queue") + + return public.RunQueue( + client=self.client, + name=name, + entity=entity, + prioritization_mode=prioritization_mode, + _access="PROJECT", + _default_resource_config_id=config_id, + _default_resource_config=config, + ) + + def create_custom_chart( + self, + entity: str, + name: str, + display_name: str, + spec_type: Literal["vega2"], + access: Literal["private", "public"], + spec: Union[str, dict], + ) -> str: + """Create a custom chart preset and return its id. + + Args: + entity: The entity (user or team) that owns the chart + name: Unique identifier for the chart preset + display_name: Human-readable name shown in the UI + spec_type: Type of specification. Must be "vega2" for Vega-Lite v2 specifications. + access: Access level for the chart: + - "private": Chart is only accessible to the entity that created it + - "public": Chart is publicly accessible + spec: The Vega/Vega-Lite specification as a dictionary or JSON string + + Returns: + The ID of the created chart preset in the format "entity/name" + + Raises: + wandb.Error: If chart creation fails + UnsupportedError: If the server doesn't support custom charts + + Example: + ```python + import wandb + + api = wandb.Api() + + # Define a simple bar chart specification + vega_spec = { + "$schema": "https://vega.github.io/schema/vega-lite/v6.json", + "mark": "bar", + "data": {"name": "wandb"}, + "encoding": { + "x": {"field": "${field:x}", "type": "ordinal"}, + "y": {"field": "${field:y}", "type": "quantitative"}, + }, + } + + # Create the custom chart + chart_id = api.create_custom_chart( + entity="my-team", + name="my-bar-chart", + display_name="My Custom Bar Chart", + spec_type="vega2", + access="private", + spec=vega_spec, + ) + + # Use with wandb.plot_table() + chart = wandb.plot_table( + vega_spec_name=chart_id, + data_table=my_table, + fields={"x": "category", "y": "value"}, + ) + ``` + """ + # Convert user-facing lowercase access to backend uppercase + backend_access = access.upper() + + api = InternalApi(retry_timedelta=RETRY_TIMEDELTA) + result = api.create_custom_chart( + entity=entity, + name=name, + display_name=display_name, + spec_type=spec_type, + access=backend_access, + spec=spec, + ) + if result is None or result.get("chart") is None: + raise wandb.Error("failed to create custom chart") + return result["chart"]["id"] + + def upsert_run_queue( + self, + name: str, + resource_config: dict, + resource_type: "public.RunQueueResourceType", + entity: Optional[str] = None, + template_variables: Optional[dict] = None, + external_links: Optional[dict] = None, + prioritization_mode: Optional["public.RunQueuePrioritizationMode"] = None, + ): + """Upsert a run queue (launch). + + Args: + name: (str) Name of the queue to create + entity: (str) Optional name of the entity to create the queue. If None, will use the configured or default entity. + resource_config: (dict) Optional default resource configuration to be used for the queue. Use handlebars (eg. `{{var}}`) to specify template variables. + resource_type: (str) Type of resource to be used for the queue. One of "local-container", "local-process", "kubernetes", "sagemaker", or "gcp-vertex". + template_variables: (dict) A dictionary of template variable schemas to be used with the config. Expected format of: + `{ + "var-name": { + "schema": { + "type": ("string", "number", or "integer"), + "default": (optional value), + "minimum": (optional minimum), + "maximum": (optional maximum), + "enum": [..."(options)"] + } + } + }` + external_links: (dict) Optional dictionary of external links to be used with the queue. Expected format of: + `{ + "name": "url" + }` + prioritization_mode: (str) Optional version of prioritization to use. Either "V0" or None + + Returns: + The upserted `RunQueue`. + + Raises: + ValueError if any of the parameters are invalid + wandb.Error on wandb API errors + """ + if entity is None: + entity = self.settings["entity"] or self.default_entity + if entity is None: + raise ValueError( + "entity must be passed as a parameter, or set in settings" + ) + + if len(name) == 0: + raise ValueError("name must be non-empty") + if len(name) > 64: + raise ValueError("name must be less than 64 characters") + + prioritization_mode = prioritization_mode or "DISABLED" + prioritization_mode = prioritization_mode.upper() + if prioritization_mode not in ["V0", "DISABLED"]: + raise ValueError( + "prioritization_mode must be 'V0' or 'DISABLED' if present" + ) + + if resource_type not in [ + "local-container", + "local-process", + "kubernetes", + "sagemaker", + "gcp-vertex", + ]: + raise ValueError( + "resource_type must be one of 'local-container', 'local-process', 'kubernetes', 'sagemaker', or 'gcp-vertex'" + ) + + self.create_project(LAUNCH_DEFAULT_PROJECT, entity) + api = InternalApi( + default_settings={ + "entity": entity, + "project": self.project(LAUNCH_DEFAULT_PROJECT), + }, + retry_timedelta=RETRY_TIMEDELTA, + ) + # User provides external_links as a dict with name: url format + # but backend stores it as a list of dicts with url and label keys. + external_links = external_links or {} + external_links = { + "links": [ + { + "label": key, + "url": value, + } + for key, value in external_links.items() + ] + } + upsert_run_queue_result = api.upsert_run_queue( + name, + entity, + resource_type, + {"resource_args": {resource_type: resource_config}}, + template_variables=template_variables, + external_links=external_links, + prioritization_mode=prioritization_mode, + ) + if not upsert_run_queue_result["success"]: + raise wandb.Error("failed to create run queue") + schema_errors = ( + upsert_run_queue_result.get("configSchemaValidationErrors") or [] + ) + for error in schema_errors: + wandb.termwarn(f"resource config validation: {error}") + + return public.RunQueue( + client=self.client, + name=name, + entity=entity, + ) + + def create_user(self, email, admin=False): + """Create a new user. + + Args: + email: (str) The email address of the user + admin: (bool) Whether this user should be a global instance admin + + Returns: + A `User` object + """ + return public.User.create(self, email, admin) + + def sync_tensorboard(self, root_dir, run_id=None, project=None, entity=None): + """Sync a local directory containing tfevent files to wandb.""" + from wandb.sync import SyncManager # TODO: circular import madness + + run_id = run_id or runid.generate_id() + project = project or self.settings.get("project") or "uncategorized" + entity = entity or self.default_entity + # TODO: pipe through log_path to inform the user how to debug + sm = SyncManager( + project=project, + entity=entity, + run_id=run_id, + mark_synced=False, + app_url=self.client.app_url, + view=False, + verbose=False, + sync_tensorboard=True, + ) + sm.add(root_dir) + sm.start() + while not sm.is_done(): + _ = sm.poll() + return self.run("/".join([entity, project, run_id])) + + @property + def client(self) -> RetryingClient: + return self._client + + @property + def user_agent(self) -> str: + return "W&B Public Client {}".format(wandb.__version__) + + @property + def api_key(self) -> Optional[str]: + # just use thread local api key if it's set + if _thread_local_api_settings.api_key: + return _thread_local_api_settings.api_key + if self._api_key is not None: + return self._api_key + auth = requests.utils.get_netrc_auth(self.settings["base_url"]) + key = None + if auth: + key = auth[-1] + # Environment should take precedence + if os.getenv("WANDB_API_KEY"): + key = os.environ["WANDB_API_KEY"] + self._api_key = key # memoize key + return key + + @property + def default_entity(self) -> Optional[str]: + if self._default_entity is None: + res = self._client.execute(self.DEFAULT_ENTITY_QUERY) + self._default_entity = (res.get("viewer") or {}).get("entity") + return self._default_entity + + @property + def viewer(self) -> "public.User": + if self._viewer is None: + self._viewer = public.User( + self._client, self._client.execute(self.VIEWER_QUERY).get("viewer") + ) + self._default_entity = self._viewer.entity + return self._viewer + + def flush(self): + """Flush the local cache. + + The api object keeps a local cache of runs, so if the state of the run may + change while executing your script you must clear the local cache with + `api.flush()` to get the latest values associated with the run. + """ + self._runs = {} + + def from_path(self, path): + """Return a run, sweep, project or report from a path. + + Examples: + ``` + project = api.from_path("my_project") + team_project = api.from_path("my_team/my_project") + run = api.from_path("my_team/my_project/runs/id") + sweep = api.from_path("my_team/my_project/sweeps/id") + report = api.from_path("my_team/my_project/reports/My-Report-Vm11dsdf") + ``` + + Args: + path: (str) The path to the project, run, sweep or report + + Returns: + A `Project`, `Run`, `Sweep`, or `BetaReport` instance. + + Raises: + wandb.Error if path is invalid or the object doesn't exist + """ + parts = path.strip("/ ").split("/") + if len(parts) == 1: + return self.project(path) + elif len(parts) == 2: + return self.project(parts[1], parts[0]) + elif len(parts) == 3: + return self.run(path) + elif len(parts) == 4: + if parts[2].startswith("run"): + return self.run(path) + elif parts[2].startswith("sweep"): + return self.sweep(path) + elif parts[2].startswith("report"): + if "--" not in parts[-1]: + if "-" in parts[-1]: + raise wandb.Error( + "Invalid report path, should be team/project/reports/Name--XXXX" + ) + else: + parts[-1] = "--" + parts[-1] + name, id = parts[-1].split("--") + return public.BetaReport( + self.client, + { + "displayName": urllib.parse.unquote(name.replace("-", " ")), + "id": id, + "spec": "{}", + }, + parts[0], + parts[1], + ) + raise wandb.Error( + "Invalid path, should be TEAM/PROJECT/TYPE/ID where TYPE is runs, sweeps, or reports" + ) + + def _parse_project_path(self, path): + """Return project and entity for project specified by path.""" + project = self.settings["project"] or "uncategorized" + entity = self.settings["entity"] or self.default_entity + if path is None: + return entity, project + parts = path.split("/", 1) + if len(parts) == 1: + return entity, path + return parts + + def _parse_path(self, path): + """Parse url, filepath, or docker paths. + + Allows paths in the following formats: + - url: entity/project/runs/id + - path: entity/project/id + - docker: entity/project:id + + Entity is optional and will fall back to the current logged-in user. + """ + project = self.settings["project"] or "uncategorized" + entity = self.settings["entity"] or self.default_entity + parts = ( + path.replace("/runs/", "/").replace("/sweeps/", "/").strip("/ ").split("/") + ) + if ":" in parts[-1]: + id = parts[-1].split(":")[-1] + parts[-1] = parts[-1].split(":")[0] + elif parts[-1]: + id = parts[-1] + if len(parts) == 1 and project != "uncategorized": + pass + elif len(parts) > 1: + project = parts[1] + if entity and id == project: + project = parts[0] + else: + entity = parts[0] + if len(parts) == 3: + entity = parts[0] + else: + project = parts[0] + return entity, project, id + + def _parse_artifact_path(self, path): + """Return project, entity and artifact name for project specified by path.""" + project = self.settings["project"] or "uncategorized" + entity = self.settings["entity"] or self.default_entity + if path is None: + return entity, project + + path, colon, alias = path.partition(":") + full_alias = colon + alias + + parts = path.split("/") + if len(parts) > 3: + raise ValueError("Invalid artifact path: {}".format(path)) + elif len(parts) == 1: + return entity, project, path + full_alias + elif len(parts) == 2: + return entity, parts[0], parts[1] + full_alias + parts[-1] += full_alias + return parts + + def projects( + self, entity: Optional[str] = None, per_page: int = 200 + ) -> "public.Projects": + """Get projects for a given entity. + + Args: + entity: (str) Name of the entity requested. If None, will fall back to the + default entity passed to `Api`. If no default entity, will raise a `ValueError`. + per_page: (int) Sets the page size for query pagination. Usually there is no reason to change this. + + Returns: + A `Projects` object which is an iterable collection of `Project` objects. + """ + if entity is None: + entity = self.settings["entity"] or self.default_entity + if entity is None: + raise ValueError( + "entity must be passed as a parameter, or set in settings" + ) + if entity not in self._projects: + self._projects[entity] = public.Projects( + self.client, entity, per_page=per_page + ) + return self._projects[entity] + + def project(self, name: str, entity: Optional[str] = None) -> "public.Project": + """Return the `Project` with the given name (and entity, if given). + + Args: + name: (str) The project name. + entity: (str) Name of the entity requested. If None, will fall back to the + default entity passed to `Api`. If no default entity, will raise a `ValueError`. + + Returns: + A `Project` object. + """ + # For registry artifacts, capture potential org user inputted before resolving entity + org = entity if is_artifact_registry_project(name) else "" + + if entity is None: + entity = self.settings["entity"] or self.default_entity + + # For registry artifacts, resolve org-based entity + if is_artifact_registry_project(name): + settings_entity = self.settings["entity"] or self.default_entity + entity = InternalApi()._resolve_org_entity_name( + entity=settings_entity, organization=org + ) + return public.Project(self.client, entity, name, {}) + + def reports( + self, path: str = "", name: Optional[str] = None, per_page: int = 50 + ) -> "public.Reports": + """Get reports for a given project path. + + WARNING: This api is in beta and will likely change in a future release + + Args: + path: (str) path to project the report resides in, should be in the form: "entity/project" + name: (str, optional) optional name of the report requested. + per_page: (int) Sets the page size for query pagination. Usually there is no reason to change this. + + Returns: + A `Reports` object which is an iterable collection of `BetaReport` objects. + """ + entity, project, _ = self._parse_path(path + "/fake_run") + + if name: + name = urllib.parse.unquote(name) + key = "/".join([entity, project, str(name)]) + else: + key = "/".join([entity, project]) + + if key not in self._reports: + self._reports[key] = public.Reports( + self.client, + public.Project(self.client, entity, project, {}), + name=name, + per_page=per_page, + ) + return self._reports[key] + + def create_team(self, team, admin_username=None): + """Create a new team. + + Args: + team: (str) The name of the team + admin_username: (str) optional username of the admin user of the team, defaults to the current user. + + Returns: + A `Team` object + """ + return public.Team.create(self, team, admin_username) + + def team(self, team: str) -> "public.Team": + """Return the matching `Team` with the given name. + + Args: + team: (str) The name of the team. + + Returns: + A `Team` object. + """ + return public.Team(self.client, team) + + def user(self, username_or_email: str) -> Optional["public.User"]: + """Return a user from a username or email address. + + Note: This function only works for Local Admins, if you are trying to get your own user object, please use `api.viewer`. + + Args: + username_or_email: (str) The username or email address of the user + + Returns: + A `User` object or None if a user couldn't be found + """ + res = self._client.execute(self.USERS_QUERY, {"query": username_or_email}) + if len(res["users"]["edges"]) == 0: + return None + elif len(res["users"]["edges"]) > 1: + wandb.termwarn( + "Found multiple users, returning the first user matching {}".format( + username_or_email + ) + ) + return public.User(self._client, res["users"]["edges"][0]["node"]) + + def users(self, username_or_email: str) -> List["public.User"]: + """Return all users from a partial username or email address query. + + Note: This function only works for Local Admins, if you are trying to get your own user object, please use `api.viewer`. + + Args: + username_or_email: (str) The prefix or suffix of the user you want to find + + Returns: + An array of `User` objects + """ + res = self._client.execute(self.USERS_QUERY, {"query": username_or_email}) + return [ + public.User(self._client, edge["node"]) for edge in res["users"]["edges"] + ] + + def runs( + self, + path: Optional[str] = None, + filters: Optional[Dict[str, Any]] = None, + order: str = "+created_at", + per_page: int = 50, + include_sweeps: bool = True, + ): + """Return a set of runs from a project that match the filters provided. + + Fields you can filter by include: + - `createdAt`: The timestamp when the run was created. (in ISO 8601 format, e.g. "2023-01-01T12:00:00Z") + - `displayName`: The human-readable display name of the run. (e.g. "eager-fox-1") + - `duration`: The total runtime of the run in seconds. + - `group`: The group name used to organize related runs together. + - `host`: The hostname where the run was executed. + - `jobType`: The type of job or purpose of the run. + - `name`: The unique identifier of the run. (e.g. "a1b2cdef") + - `state`: The current state of the run. + - `tags`: The tags associated with the run. + - `username`: The username of the user who initiated the run + + Additionally, you can filter by items in the run config or summary metrics. + Such as `config.experiment_name`, `summary_metrics.loss`, etc. + + For more complex filtering, you can use MongoDB query operators. + For details, see: https://docs.mongodb.com/manual/reference/operator/query + The following operations are supported: + - `$and` + - `$or` + - `$nor` + - `$eq` + - `$ne` + - `$gt` + - `$gte` + - `$lt` + - `$lte` + - `$in` + - `$nin` + - `$exists` + - `$regex` + + + Examples: + Find runs in my_project where config.experiment_name has been set to "foo" + ``` + api.runs( + path="my_entity/my_project", + filters={"config.experiment_name": "foo"}, + ) + ``` + + Find runs in my_project where config.experiment_name has been set to "foo" or "bar" + ``` + api.runs( + path="my_entity/my_project", + filters={ + "$or": [ + {"config.experiment_name": "foo"}, + {"config.experiment_name": "bar"}, + ] + }, + ) + ``` + + Find runs in my_project where config.experiment_name matches a regex (anchors are not supported) + ``` + api.runs( + path="my_entity/my_project", + filters={"config.experiment_name": {"$regex": "b.*"}}, + ) + ``` + + Find runs in my_project where the run name matches a regex (anchors are not supported) + ``` + api.runs( + path="my_entity/my_project", + filters={"display_name": {"$regex": "^foo.*"}}, + ) + ``` + + Find runs in my_project where config.experiment contains a nested field "category" with value "testing" + ``` + api.runs( + path="my_entity/my_project", + filters={"config.experiment.category": "testing"}, + ) + ``` + + Find runs in my_project with a loss value of 0.5 nested in a dictionary under model1 in the summary metrics + ``` + api.runs( + path="my_entity/my_project", + filters={"summary_metrics.model1.loss": 0.5}, + ) + ``` + + Find runs in my_project sorted by ascending loss + ``` + api.runs(path="my_entity/my_project", order="+summary_metrics.loss") + ``` + + Args: + path: (str) path to project, should be in the form: "entity/project" + filters: (dict) queries for specific runs using the MongoDB query language. + You can filter by run properties such as config.key, summary_metrics.key, state, entity, createdAt, etc. + For example: `{"config.experiment_name": "foo"}` would find runs with a config entry + of experiment name set to "foo" + order: (str) Order can be `created_at`, `heartbeat_at`, `config.*.value`, or `summary_metrics.*`. + If you prepend order with a + order is ascending. + If you prepend order with a - order is descending (default). + The default order is run.created_at from oldest to newest. + per_page: (int) Sets the page size for query pagination. + include_sweeps: (bool) Whether to include the sweep runs in the results. + + Returns: + A `Runs` object, which is an iterable collection of `Run` objects. + """ + entity, project = self._parse_project_path(path) + filters = filters or {} + key = (path or "") + str(filters) + str(order) + if not self._runs.get(key): + self._runs[key] = public.Runs( + self.client, + entity, + project, + filters=filters, + order=order, + per_page=per_page, + include_sweeps=include_sweeps, + ) + return self._runs[key] + + @normalize_exceptions + def run(self, path=""): + """Return a single run by parsing path in the form entity/project/run_id. + + Args: + path: (str) path to run in the form `entity/project/run_id`. + If `api.entity` is set, this can be in the form `project/run_id` + and if `api.project` is set this can just be the run_id. + + Returns: + A `Run` object. + """ + entity, project, run_id = self._parse_path(path) + if not self._runs.get(path): + self._runs[path] = public.Run(self.client, entity, project, run_id) + return self._runs[path] + + def queued_run( + self, + entity, + project, + queue_name, + run_queue_item_id, + project_queue=None, + priority=None, + ): + """Return a single queued run based on the path. + + Parses paths of the form entity/project/queue_id/run_queue_item_id. + """ + return public.QueuedRun( + self.client, + entity, + project, + queue_name, + run_queue_item_id, + project_queue=project_queue, + priority=priority, + ) + + def run_queue( + self, + entity, + name, + ): + """Return the named `RunQueue` for entity. + + To create a new `RunQueue`, use `wandb.Api().create_run_queue(...)`. + """ + return public.RunQueue( + self.client, + name, + entity, + ) + + @normalize_exceptions + def sweep(self, path=""): + """Return a sweep by parsing path in the form `entity/project/sweep_id`. + + Args: + path: (str, optional) path to sweep in the form entity/project/sweep_id. If `api.entity` + is set, this can be in the form project/sweep_id and if `api.project` is set + this can just be the sweep_id. + + Returns: + A `Sweep` object. + """ + entity, project, sweep_id = self._parse_path(path) + if not self._sweeps.get(path): + self._sweeps[path] = public.Sweep(self.client, entity, project, sweep_id) + return self._sweeps[path] + + @normalize_exceptions + def artifact_types(self, project: Optional[str] = None) -> "public.ArtifactTypes": + """Return a collection of matching artifact types. + + Args: + project: (str, optional) If given, a project name or path to filter on. + + Returns: + An iterable `ArtifactTypes` object. + """ + project_path = project + entity, project = self._parse_project_path(project_path) + # If its a Registry project, the entity is considered to be an org instead + if is_artifact_registry_project(project): + settings_entity = self.settings["entity"] or self.default_entity + org = parse_org_from_registry_path(project_path, PathType.PROJECT) + entity = InternalApi()._resolve_org_entity_name( + entity=settings_entity, organization=org + ) + return public.ArtifactTypes(self.client, entity, project) + + @normalize_exceptions + def artifact_type( + self, type_name: str, project: Optional[str] = None + ) -> "public.ArtifactType": + """Return the matching `ArtifactType`. + + Args: + type_name: (str) The name of the artifact type to retrieve. + project: (str, optional) If given, a project name or path to filter on. + + Returns: + An `ArtifactType` object. + """ + project_path = project + entity, project = self._parse_project_path(project_path) + # If its an Registry artifact, the entity is an org instead + if is_artifact_registry_project(project): + org = parse_org_from_registry_path(project_path, PathType.PROJECT) + settings_entity = self.settings["entity"] or self.default_entity + entity = InternalApi()._resolve_org_entity_name( + entity=settings_entity, organization=org + ) + return public.ArtifactType(self.client, entity, project, type_name) + + @normalize_exceptions + def artifact_collections( + self, project_name: str, type_name: str, per_page: int = 50 + ) -> "public.ArtifactCollections": + """Return a collection of matching artifact collections. + + Args: + project_name: (str) The name of the project to filter on. + type_name: (str) The name of the artifact type to filter on. + per_page: (int) Sets the page size for query pagination. Usually there is no reason to change this. + + Returns: + An iterable `ArtifactCollections` object. + """ + entity, project = self._parse_project_path(project_name) + # If iterating through Registry project, the entity is considered to be an org instead + if is_artifact_registry_project(project): + org = parse_org_from_registry_path(project_name, PathType.PROJECT) + settings_entity = self.settings["entity"] or self.default_entity + entity = InternalApi()._resolve_org_entity_name( + entity=settings_entity, organization=org + ) + return public.ArtifactCollections( + self.client, entity, project, type_name, per_page=per_page + ) + + @normalize_exceptions + def artifact_collection( + self, type_name: str, name: str + ) -> "public.ArtifactCollection": + """Return a single artifact collection by type and parsing path in the form `entity/project/name`. + + Args: + type_name: (str) The type of artifact collection to fetch. + name: (str) An artifact collection name. May be prefixed with entity/project. + + Returns: + An `ArtifactCollection` object. + """ + entity, project, collection_name = self._parse_artifact_path(name) + # If its an Registry artifact, the entity is considered to be an org instead + if is_artifact_registry_project(project): + org = parse_org_from_registry_path(name, PathType.ARTIFACT) + settings_entity = self.settings["entity"] or self.default_entity + entity = InternalApi()._resolve_org_entity_name( + entity=settings_entity, organization=org + ) + + if entity is None: + raise ValueError( + "Could not determine entity. Please include the entity as part of the collection name path." + ) + + return public.ArtifactCollection( + self.client, entity, project, collection_name, type_name + ) + + @normalize_exceptions + def artifact_versions(self, type_name, name, per_page=50): + """Deprecated, use `artifacts(type_name, name)` instead.""" + deprecate( + field_name=Deprecated.api__artifact_versions, + warning_message=( + "Api.artifact_versions(type_name, name) is deprecated, " + "use Api.artifacts(type_name, name) instead." + ), + ) + return self.artifacts(type_name, name, per_page=per_page) + + @normalize_exceptions + def artifacts( + self, + type_name: str, + name: str, + per_page: int = 50, + tags: Optional[List[str]] = None, + ) -> "public.Artifacts": + """Return an `Artifacts` collection from the given parameters. + + Args: + type_name: (str) The type of artifacts to fetch. + name: (str) An artifact collection name. May be prefixed with entity/project. + per_page: (int) Sets the page size for query pagination. Usually there is no reason to change this. + tags: (list[str], optional) Only return artifacts with all of these tags. + + Returns: + An iterable `Artifacts` object. + """ + entity, project, collection_name = self._parse_artifact_path(name) + # If its an Registry project, the entity is considered to be an org instead + if is_artifact_registry_project(project): + org = parse_org_from_registry_path(name, PathType.ARTIFACT) + settings_entity = self.settings["entity"] or self.default_entity + entity = InternalApi()._resolve_org_entity_name( + entity=settings_entity, organization=org + ) + return public.Artifacts( + self.client, + entity, + project, + collection_name, + type_name, + per_page=per_page, + tags=tags, + ) + + @normalize_exceptions + def _artifact( + self, name: str, type: Optional[str] = None, enable_tracking: bool = False + ): + if name is None: + raise ValueError("You must specify name= to fetch an artifact.") + entity, project, artifact_name = self._parse_artifact_path(name) + + # If its an Registry artifact, the entity is an org instead + if is_artifact_registry_project(project): + organization = ( + name.split("/")[0] + if name.count("/") == 2 + else self.settings["organization"] + ) + # set entity to match the settings since in above code it was potentially set to an org + settings_entity = self.settings["entity"] or self.default_entity + # Registry artifacts are under the org entity. Because we offer a shorthand and alias for this path, + # we need to fetch the org entity to for the user behind the scenes. + entity = InternalApi()._resolve_org_entity_name( + entity=settings_entity, organization=organization + ) + + if entity is None: + raise ValueError( + "Could not determine entity. Please include the entity as part of the artifact name path." + ) + + artifact = wandb.Artifact._from_name( + entity=entity, + project=project, + name=artifact_name, + client=self.client, + enable_tracking=enable_tracking, + ) + if type is not None and artifact.type != type: + raise ValueError( + f"type {type} specified but this artifact is of type {artifact.type}" + ) + return artifact + + @normalize_exceptions + def artifact(self, name: str, type: Optional[str] = None): + """Return a single artifact by parsing path in the form `project/name` or `entity/project/name`. + + Args: + name: (str) An artifact name. May be prefixed with project/ or entity/project/. + If no entity is specified in the name, the Run or API setting's entity is used. + Valid names can be in the following forms: + name:version + name:alias + type: (str, optional) The type of artifact to fetch. + + Returns: + An `Artifact` object. + + Raises: + ValueError: If the artifact name is not specified. + ValueError: If the artifact type is specified but does not match the type of the fetched artifact. + + Note: + This method is intended for external use only. Do not call `api.artifact()` within the wandb repository code. + """ + return self._artifact(name=name, type=type, enable_tracking=True) + + @normalize_exceptions + def job(self, name: Optional[str], path: Optional[str] = None) -> "public.Job": + """Return a `Job` from the given parameters. + + Args: + name: (str) The job name. + path: (str, optional) If given, the root path in which to download the job artifact. + + Returns: + A `Job` object. + """ + if name is None: + raise ValueError("You must specify name= to fetch a job.") + elif name.count("/") != 2 or ":" not in name: + raise ValueError( + "Invalid job specification. A job must be of the form: //:" + ) + return public.Job(self, name, path) + + @normalize_exceptions + def list_jobs(self, entity: str, project: str) -> List[Dict[str, Any]]: + """Return a list of jobs, if any, for the given entity and project. + + Args: + entity: (str) The entity for the listed job(s). + project: (str) The project for the listed job(s). + + Returns: + A list of matching jobs. + """ + if entity is None: + raise ValueError("Specify an entity when listing jobs") + if project is None: + raise ValueError("Specify a project when listing jobs") + + query = gql( + """ + query ArtifactOfType( + $entityName: String!, + $projectName: String!, + $artifactTypeName: String!, + ) { + project(name: $projectName, entityName: $entityName) { + artifactType(name: $artifactTypeName) { + artifactCollections { + edges { + node { + artifacts { + edges { + node { + id + state + aliases { + alias + } + artifactSequence { + name + } + } + } + } + } + } + } + } + } + } + """ + ) + + try: + artifact_query = self._client.execute( + query, + { + "projectName": project, + "entityName": entity, + "artifactTypeName": "job", + }, + ) + + if not artifact_query or not artifact_query["project"]: + wandb.termerror( + f"Project: '{project}' not found in entity: '{entity}' or access denied." + ) + return [] + + if artifact_query["project"]["artifactType"] is None: + return [] + + artifacts = artifact_query["project"]["artifactType"][ + "artifactCollections" + ]["edges"] + + return [x["node"]["artifacts"] for x in artifacts] + except requests.exceptions.HTTPError: + return False + + @normalize_exceptions + def artifact_exists(self, name: str, type: Optional[str] = None) -> bool: + """Return whether an artifact version exists within a specified project and entity. + + Args: + name: (str) An artifact name. May be prefixed with entity/project. + If entity or project is not specified, it will be inferred from the override params if populated. + Otherwise, entity will be pulled from the user settings and project will default to "uncategorized". + Valid names can be in the following forms: + name:version + name:alias + type: (str, optional) The type of artifact + + Returns: + True if the artifact version exists, False otherwise. + """ + try: + self._artifact(name, type) + except wandb.errors.CommError: + return False + + return True + + @normalize_exceptions + def artifact_collection_exists(self, name: str, type: str) -> bool: + """Return whether an artifact collection exists within a specified project and entity. + + Args: + name: (str) An artifact collection name. May be prefixed with entity/project. + If entity or project is not specified, it will be inferred from the override params if populated. + Otherwise, entity will be pulled from the user settings and project will default to "uncategorized". + type: (str) The type of artifact collection + + Returns: + True if the artifact collection exists, False otherwise. + """ + try: + self.artifact_collection(type, name) + except wandb.errors.CommError: + return False + + return True + + def registries( + self, + organization: Optional[str] = None, + filter: Optional[Dict[str, Any]] = None, + ) -> Registries: + """Returns a Registry iterator. + + Use the iterator to search and filter registries, collections, + or artifact versions across your organization's registry. + + Examples: + Find all registries with the names that contain "model" + ```python + import wandb + + api = wandb.Api() # specify an org if your entity belongs to multiple orgs + api.registries(filter={"name": {"$regex": "model"}}) + ``` + + Find all collections in the registries with the name "my_collection" and the tag "my_tag" + ```python + api.registries().collections( + filter={"name": "my_collection", "tag": "my_tag"} + ) + ``` + + Find all artifact versions in the registries with a collection name that contains "my_collection" and a version that has the alias "best" + ```python + api.registries().collections( + filter={"name": {"$regex": "my_collection"}} + ).versions(filter={"alias": "best"}) + ``` + + Find all artifact versions in the registries that contain "model" and have the tag "prod" or alias "best" + ```python + api.registries(filter={"name": {"$regex": "model"}}).versions( + filter={"$or": [{"tag": "prod"}, {"alias": "best"}]} + ) + ``` + + Args: + organization: (str, optional) The organization of the registry to fetch. + If not specified, use the organization specified in the user's settings. + filter: (dict, optional) MongoDB-style filter to apply to each object in the registry iterator. + Fields available to filter for collections are + `name`, `description`, `created_at`, `updated_at`. + Fields available to filter for collections are + `name`, `tag`, `description`, `created_at`, `updated_at` + Fields available to filter for versions are + `tag`, `alias`, `created_at`, `updated_at`, `metadata` + + Returns: + A registry iterator. + """ + if not InternalApi()._server_supports(ServerFeature.ARTIFACT_REGISTRY_SEARCH): + raise RuntimeError( + "Registry search API is not enabled on this wandb server version. " + "Please upgrade your server version or contact support at support@wandb.com." + ) + + organization = organization or fetch_org_from_settings_or_entity( + self.settings, self.default_entity + ) + return Registries(self.client, organization, filter) + + def registry(self, name: str, organization: Optional[str] = None) -> Registry: + """Return a registry given a registry name. + + Args: + name: The name of the registry. This is without the `wandb-registry-` + prefix. + organization: The organization of the registry. + If no organization is set in the settings, the organization will be + fetched from the entity if the entity only belongs to one + organization. + + Returns: + A registry object. + + Examples: + Fetch and update a registry + ```python + import wandb + + api = wandb.Api() + registry = api.registry(name="my-registry", organization="my-org") + registry.description = "This is an updated description" + registry.save() + ``` + """ + if not InternalApi()._server_supports(ServerFeature.ARTIFACT_REGISTRY_SEARCH): + raise RuntimeError( + "api.registry() is not enabled on this wandb server version. " + "Please upgrade your server version or contact support at support@wandb.com." + ) + organization = organization or fetch_org_from_settings_or_entity( + self.settings, self.default_entity + ) + org_entity = fetch_org_entity_from_organization(self.client, organization) + registry = Registry(self.client, organization, org_entity, name) + registry.load() + return registry + + def create_registry( + self, + name: str, + visibility: Literal["organization", "restricted"], + organization: Optional[str] = None, + description: Optional[str] = None, + artifact_types: Optional[List[str]] = None, + ) -> Registry: + """Create a new registry. + + Args: + name: The name of the registry. Name must be unique within the organization. + visibility: The visibility of the registry. + organization: Anyone in the organization can view this registry. You can + edit their roles later from the settings in the UI. + restricted: Only invited members via the UI can access this registry. + Public sharing is disabled. + organization: The organization of the registry. + If no organization is set in the settings, the organization will be + fetched from the entity if the entity only belongs to one organization. + description: The description of the registry. + artifact_types: The accepted artifact types of the registry. A type is no + more than 128 characters and do not include characters `/` or `:`. If + not specified, all types are accepted. + Allowed types added to the registry cannot be removed later. + + Returns: + A registry object. + + Examples: + ```python + import wandb + + api = wandb.Api() + registry = api.create_registry( + name="my-registry", + visibility="restricted", + organization="my-org", + description="This is a test registry", + artifact_types=["model"], + ) + ``` + """ + if not InternalApi()._server_supports( + ServerFeature.INCLUDE_ARTIFACT_TYPES_IN_REGISTRY_CREATION + ): + raise RuntimeError( + "create_registry api is not enabled on this wandb server version. " + "Please upgrade your server version or contact support at support@wandb.com." + ) + + organization = organization or fetch_org_from_settings_or_entity( + self.settings, self.default_entity + ) + + try: + existing_registry = self.registry(name=name, organization=organization) + except ValueError: + existing_registry = None + if existing_registry: + raise ValueError( + f"Registry {name!r} already exists in organization {organization!r}," + " please use a different name." + ) + + return Registry.create( + self.client, + organization, + name, + visibility, + description, + artifact_types, + ) + + def integrations( + self, + entity: Optional[str] = None, + *, + per_page: int = 50, + ) -> Iterator["Integration"]: + """Return an iterator of all integrations for an entity. + + Args: + entity: The entity (e.g. team name) for which to + fetch integrations. If not provided, the user's default entity + will be used. + per_page: Number of integrations to fetch per page. + Defaults to 50. Usually there is no reason to change this. + + Yields: + Iterator[SlackIntegration | WebhookIntegration]: An iterator of any supported integrations. + """ + from wandb.apis.public.integrations import Integrations + + params = {"entityName": entity or self.default_entity} + return Integrations(client=self.client, variables=params, per_page=per_page) + + def webhook_integrations( + self, entity: Optional[str] = None, *, per_page: int = 50 + ) -> Iterator["WebhookIntegration"]: + """Returns an iterator of webhook integrations for an entity. + + Args: + entity: The entity (e.g. team name) for which to + fetch integrations. If not provided, the user's default entity + will be used. + per_page: Number of integrations to fetch per page. + Defaults to 50. Usually there is no reason to change this. + + Yields: + Iterator[WebhookIntegration]: An iterator of webhook integrations. + + Examples: + Get all registered webhook integrations for the team "my-team": + ```python + import wandb + + api = wandb.Api() + webhook_integrations = api.webhook_integrations(entity="my-team") + ``` + + Find only webhook integrations that post requests to "https://my-fake-url.com": + ```python + webhook_integrations = api.webhook_integrations(entity="my-team") + my_webhooks = [ + ig + for ig in webhook_integrations + if ig.url_endpoint.startswith("https://my-fake-url.com") + ] + ``` + """ + from wandb.apis.public.integrations import WebhookIntegrations + + params = {"entityName": entity or self.default_entity} + return WebhookIntegrations( + client=self.client, variables=params, per_page=per_page + ) + + def slack_integrations( + self, *, entity: Optional[str] = None, per_page: int = 50 + ) -> Iterator["SlackIntegration"]: + """Returns an iterator of Slack integrations for an entity. + + Args: + entity: The entity (e.g. team name) for which to + fetch integrations. If not provided, the user's default entity + will be used. + per_page: Number of integrations to fetch per page. + Defaults to 50. Usually there is no reason to change this. + + Yields: + Iterator[SlackIntegration]: An iterator of Slack integrations. + + Examples: + Get all registered Slack integrations for the team "my-team": + ```python + import wandb + + api = wandb.Api() + slack_integrations = api.slack_integrations(entity="my-team") + ``` + + Find only Slack integrations that post to channel names starting with "team-alerts-": + ```python + slack_integrations = api.slack_integrations(entity="my-team") + team_alert_integrations = [ + ig + for ig in slack_integrations + if ig.channel_name.startswith("team-alerts-") + ] + ``` + """ + from wandb.apis.public.integrations import SlackIntegrations + + params = {"entityName": entity or self.default_entity} + return SlackIntegrations( + client=self.client, variables=params, per_page=per_page + ) + + def _supports_automation( + self, + *, + event: Optional["EventType"] = None, + action: Optional["ActionType"] = None, + ) -> bool: + """Returns whether the server recognizes the automation event and/or action.""" + from wandb.automations._utils import ( + ALWAYS_SUPPORTED_ACTIONS, + ALWAYS_SUPPORTED_EVENTS, + ) + + api = InternalApi() + supports_event = ( + (event is None) + or (event in ALWAYS_SUPPORTED_EVENTS) + or api._server_supports(f"AUTOMATION_EVENT_{event.value}") + ) + supports_action = ( + (action is None) + or (action in ALWAYS_SUPPORTED_ACTIONS) + or api._server_supports(f"AUTOMATION_ACTION_{action.value}") + ) + return supports_event and supports_action + + def _omitted_automation_fragments(self) -> Set[str]: + """Returns the names of unsupported automation-related fragments. + + Older servers won't recognize newer GraphQL types, so a valid request may + unnecessarily error out because it won't recognize fragments defined on those types. + + So e.g. if a server does not support `NO_OP` action types, then the following need to be + removed from the body of the GraphQL request: + + - Fragment definition: + ``` + fragment NoOpActionFields on NoOpTriggeredAction { + noOp + } + ``` + + - Fragment spread in selection set: + ``` + { + ...NoOpActionFields + # ... other fields ... + } + ``` + """ + from wandb.automations import ActionType + from wandb.automations._generated import ( + GenericWebhookActionFields, + NoOpActionFields, + NotificationActionFields, + QueueJobActionFields, + ) + + # Note: we can't currently define this as a constant outside the method + # and still keep it nearby in this module, because it relies on pydantic v2-only imports + fragment_names: dict[ActionType, str] = { + ActionType.NO_OP: NoOpActionFields.__name__, + ActionType.QUEUE_JOB: QueueJobActionFields.__name__, + ActionType.NOTIFICATION: NotificationActionFields.__name__, + ActionType.GENERIC_WEBHOOK: GenericWebhookActionFields.__name__, + } + + return set( + name + for action in ActionType + if (not self._supports_automation(action=action)) + and (name := fragment_names.get(action)) + ) + + def automation( + self, + name: str, + *, + entity: Optional[str] = None, + ) -> "Automation": + """Returns the only Automation matching the parameters. + + Args: + name: The name of the automation to fetch. + entity: The entity to fetch the automation for. + + Raises: + ValueError: If zero or multiple Automations match the search criteria. + + Examples: + Get an existing automation named "my-automation": + + ```python + import wandb + + api = wandb.Api() + automation = api.automation(name="my-automation") + ``` + + Get an existing automation named "other-automation", from the entity "my-team": + + ```python + automation = api.automation(name="other-automation", entity="my-team") + ``` + """ + return one( + self.automations(entity=entity, name=name), + too_short=ValueError("No automations found"), + too_long=ValueError("Multiple automations found"), + ) + + def automations( + self, + entity: Optional[str] = None, + *, + name: Optional[str] = None, + per_page: int = 50, + ) -> Iterator["Automation"]: + """Returns an iterator over all Automations that match the given parameters. + + If no parameters are provided, the returned iterator will contain all + Automations that the user has access to. + + Args: + entity: The entity to fetch the automations for. + name: The name of the automation to fetch. + per_page: The number of automations to fetch per page. + Defaults to 50. Usually there is no reason to change this. + + Returns: + A list of automations. + + Examples: + Fetch all existing automations for the entity "my-team": + + ```python + import wandb + + api = wandb.Api() + automations = api.automations(entity="my-team") + ``` + """ + from wandb.apis.public.automations import Automations + from wandb.automations._generated import ( + GET_AUTOMATIONS_BY_ENTITY_GQL, + GET_AUTOMATIONS_GQL, + ) + + # For now, we need to use different queries depending on whether entity is given + variables = {"entityName": entity} + if entity is None: + gql_str = GET_AUTOMATIONS_GQL # Automations for viewer + else: + gql_str = GET_AUTOMATIONS_BY_ENTITY_GQL # Automations for entity + + # If needed, rewrite the GraphQL field selection set to omit unsupported fields/fragments/types + omit_fragments = self._omitted_automation_fragments() + query = gql_compat(gql_str, omit_fragments=omit_fragments) + iterator = Automations( + client=self.client, variables=variables, per_page=per_page, _query=query + ) + + # FIXME: this is crude, move this client-side filtering logic into backend + if name is not None: + iterator = filter(lambda x: x.name == name, iterator) + yield from iterator + + @normalize_exceptions + def create_automation( + self, + obj: "NewAutomation", + *, + fetch_existing: bool = False, + **kwargs: Unpack["WriteAutomationsKwargs"], + ) -> "Automation": + """Create a new Automation. + + Args: + obj: + The automation to create. + fetch_existing: + If True, and a conflicting automation already exists, attempt + to fetch the existing automation instead of raising an error. + **kwargs: + Any additional values to assign to the automation before + creating it. If given, these will override any values that may + already be set on the automation: + - `name`: The name of the automation. + - `description`: The description of the automation. + - `enabled`: Whether the automation is enabled. + - `scope`: The scope of the automation. + - `event`: The event that triggers the automation. + - `action`: The action that is triggered by the automation. + + Returns: + The saved Automation. + + Examples: + Create a new automation named "my-automation" that sends a Slack notification + when a run within a specific project logs a metric exceeding a custom threshold: + + ```python + import wandb + from wandb.automations import OnRunMetric, RunEvent, SendNotification + + api = wandb.Api() + + project = api.project("my-project", entity="my-team") + + # Use the first Slack integration for the team + slack_hook = next(api.slack_integrations(entity="my-team")) + + event = OnRunMetric( + scope=project, + filter=RunEvent.metric("custom-metric") > 10, + ) + action = SendNotification.from_integration(slack_hook) + + automation = api.create_automation( + event >> action, + name="my-automation", + description="Send a Slack message whenever 'custom-metric' exceeds 10.", + ) + ``` + """ + from wandb.automations import Automation + from wandb.automations._generated import CREATE_AUTOMATION_GQL, CreateAutomation + from wandb.automations._utils import prepare_to_create + + gql_input = prepare_to_create(obj, **kwargs) + + if not self._supports_automation( + event=(event := gql_input.triggering_event_type), + action=(action := gql_input.triggered_action_type), + ): + raise ValueError( + f"Automation event or action ({event!r} -> {action!r}) " + "is not supported on this wandb server version. " + "Please upgrade your server version, or contact support at " + "support@wandb.com." + ) + + # If needed, rewrite the GraphQL field selection set to omit unsupported fields/fragments/types + omit_fragments = self._omitted_automation_fragments() + mutation = gql_compat(CREATE_AUTOMATION_GQL, omit_fragments=omit_fragments) + variables = {"params": gql_input.model_dump(exclude_none=True)} + + name = gql_input.name + try: + data = self.client.execute(mutation, variable_values=variables) + except requests.HTTPError as e: + status = HTTPStatus(e.response.status_code) + if status is HTTPStatus.CONFLICT: # 409 + if fetch_existing: + wandb.termlog(f"Automation {name!r} exists. Fetching it instead.") + return self.automation(name=name) + + raise ValueError( + f"Automation {name!r} exists. Unable to create another with the same name." + ) from None + raise + + try: + result = CreateAutomation.model_validate(data).result + except ValidationError as e: + msg = f"Invalid response while creating automation {name!r}" + raise RuntimeError(msg) from e + + if (result is None) or (result.trigger is None): + msg = f"Empty response while creating automation {name!r}" + raise RuntimeError(msg) + + return Automation.model_validate(result.trigger) + + @normalize_exceptions + def update_automation( + self, + obj: "Automation", + *, + create_missing: bool = False, + **kwargs: Unpack["WriteAutomationsKwargs"], + ) -> "Automation": + """Update an existing automation. + + Args: + obj: The automation to update. Must be an existing automation. + create_missing (bool): + If True, and the automation does not exist, create it. + **kwargs: + Any additional values to assign to the automation before + updating it. If given, these will override any values that may + already be set on the automation: + - `name`: The name of the automation. + - `description`: The description of the automation. + - `enabled`: Whether the automation is enabled. + - `scope`: The scope of the automation. + - `event`: The event that triggers the automation. + - `action`: The action that is triggered by the automation. + + Returns: + The updated automation. + + Examples: + Disable and edit the description of an existing automation ("my-automation"): + + ```python + import wandb + + api = wandb.Api() + + automation = api.automation(name="my-automation") + automation.enabled = False + automation.description = "Kept for reference, but no longer used." + + updated_automation = api.update_automation(automation) + ``` + + OR: + + ```python + import wandb + + api = wandb.Api() + + automation = api.automation(name="my-automation") + + updated_automation = api.update_automation( + automation, + enabled=False, + description="Kept for reference, but no longer used.", + ) + ``` + """ + from wandb.automations import ActionType, Automation + from wandb.automations._generated import UPDATE_AUTOMATION_GQL, UpdateAutomation + from wandb.automations._utils import prepare_to_update + + # Check if the server even supports updating automations. + # + # NOTE: Unfortunately, there is no current server feature flag for this. As a workaround, + # we check whether the server supports the NO_OP action, which is a reasonably safe proxy + # for whether it supports updating automations. + if not self._supports_automation(action=ActionType.NO_OP): + raise RuntimeError( + "Updating existing automations is not enabled on this wandb server version. " + "Please upgrade your server version, or contact support at support@wandb.com." + ) + + gql_input = prepare_to_update(obj, **kwargs) + + if not self._supports_automation( + event=(event := gql_input.triggering_event_type), + action=(action := gql_input.triggered_action_type), + ): + raise ValueError( + f"Automation event or action ({event.value} -> {action.value}) " + "is not supported on this wandb server version. " + "Please upgrade your server version, or contact support at " + "support@wandb.com." + ) + + # If needed, rewrite the GraphQL field selection set to omit unsupported fields/fragments/types + omit_fragments = self._omitted_automation_fragments() + mutation = gql_compat(UPDATE_AUTOMATION_GQL, omit_fragments=omit_fragments) + variables = {"params": gql_input.model_dump(exclude_none=True)} + + name = gql_input.name + try: + data = self.client.execute(mutation, variable_values=variables) + except requests.HTTPError as e: + status = HTTPStatus(e.response.status_code) + if status is HTTPStatus.NOT_FOUND: # 404 + if create_missing: + wandb.termlog(f"Automation {name!r} not found. Creating it.") + return self.create_automation(obj) + + raise ValueError( + f"Automation {name!r} not found. Unable to edit it." + ) from e + + # Not a (known) recoverable HTTP error + wandb.termerror(f"Got response status {status!r}: {e.response.text!r}") + raise + + try: + result = UpdateAutomation.model_validate(data).result + except ValidationError as e: + msg = f"Invalid response while updating automation {name!r}" + raise RuntimeError(msg) from e + + if (result is None) or (result.trigger is None): + msg = f"Empty response while updating automation {name!r}" + raise RuntimeError(msg) + + return Automation.model_validate(result.trigger) + + @normalize_exceptions + def delete_automation(self, obj: Union["Automation", str]) -> Literal[True]: + """Delete an automation. + + Args: + obj: The automation to delete, or its ID. + + Returns: + True if the automation was deleted successfully. + """ + from wandb.automations._generated import DELETE_AUTOMATION_GQL, DeleteAutomation + from wandb.automations._utils import extract_id + + id_ = extract_id(obj) + mutation = gql(DELETE_AUTOMATION_GQL) + variables = {"id": id_} + + data = self.client.execute(mutation, variable_values=variables) + + try: + result = DeleteAutomation.model_validate(data).result + except ValidationError as e: + msg = f"Invalid response while deleting automation {id_!r}" + raise RuntimeError(msg) from e + + if result is None: + msg = f"Empty response while deleting automation {id_!r}" + raise RuntimeError(msg) + + if not result.success: + raise RuntimeError(f"Failed to delete automation: {id_!r}") + + return result.success diff --git a/lib/python3.12/site-packages/wandb/apis/public/artifacts.py b/lib/python3.12/site-packages/wandb/apis/public/artifacts.py new file mode 100644 index 0000000000000000000000000000000000000000..9fb3a10bd9928192aea3d59d1cacc6cc050d1027 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/public/artifacts.py @@ -0,0 +1,852 @@ +"""Public API: artifacts.""" + +from __future__ import annotations + +import json +import re +from copy import copy +from typing import TYPE_CHECKING, Any, Iterable, Literal, Mapping, Sequence + +from typing_extensions import override +from wandb_gql import Client, gql + +import wandb +from wandb.apis import public +from wandb.apis.normalize import normalize_exceptions +from wandb.apis.paginator import Paginator, SizedPaginator +from wandb.errors.term import termlog +from wandb.proto.wandb_deprecated import Deprecated +from wandb.proto.wandb_internal_pb2 import ServerFeature +from wandb.sdk.artifacts._generated import ( + ARTIFACT_COLLECTION_MEMBERSHIP_FILES_GQL, + ARTIFACT_VERSION_FILES_GQL, + CREATE_ARTIFACT_COLLECTION_TAG_ASSIGNMENTS_GQL, + DELETE_ARTIFACT_COLLECTION_TAG_ASSIGNMENTS_GQL, + DELETE_ARTIFACT_PORTFOLIO_GQL, + DELETE_ARTIFACT_SEQUENCE_GQL, + MOVE_ARTIFACT_COLLECTION_GQL, + PROJECT_ARTIFACT_COLLECTION_GQL, + PROJECT_ARTIFACT_COLLECTIONS_GQL, + PROJECT_ARTIFACT_TYPE_GQL, + PROJECT_ARTIFACT_TYPES_GQL, + PROJECT_ARTIFACTS_GQL, + RUN_INPUT_ARTIFACTS_GQL, + RUN_OUTPUT_ARTIFACTS_GQL, + UPDATE_ARTIFACT_PORTFOLIO_GQL, + UPDATE_ARTIFACT_SEQUENCE_GQL, + ArtifactCollectionMembershipFiles, + ArtifactCollectionsFragment, + ArtifactsFragment, + ArtifactTypeFragment, + ArtifactTypesFragment, + ArtifactVersionFiles, + FilesFragment, + ProjectArtifactCollection, + ProjectArtifactCollections, + ProjectArtifacts, + ProjectArtifactType, + ProjectArtifactTypes, + RunInputArtifactsProjectRunInputArtifacts, + RunOutputArtifactsProjectRunOutputArtifacts, +) +from wandb.sdk.artifacts._graphql_fragments import omit_artifact_fields +from wandb.sdk.artifacts._validators import ( + SOURCE_ARTIFACT_COLLECTION_TYPE, + validate_artifact_name, + validate_artifact_type, +) +from wandb.sdk.internal.internal_api import Api as InternalApi +from wandb.sdk.lib import deprecate + +from .utils import gql_compat + +if TYPE_CHECKING: + from wandb.sdk.artifacts.artifact import Artifact + + from . import RetryingClient, Run + + +class ArtifactTypes(Paginator["ArtifactType"]): + QUERY = gql(PROJECT_ARTIFACT_TYPES_GQL) + + last_response: ArtifactTypesFragment | None + + def __init__( + self, + client: Client, + entity: str, + project: str, + per_page: int = 50, + ): + self.entity = entity + self.project = project + + variable_values = { + "entityName": entity, + "projectName": project, + } + super().__init__(client, variable_values, per_page) + + @override + def _update_response(self) -> None: + """Fetch and validate the response data for the current page.""" + data = self.client.execute(self.QUERY, variable_values=self.variables) + result = ProjectArtifactTypes.model_validate(data) + + # Extract the inner `*Connection` result for faster/easier access. + if not ((proj := result.project) and (conn := proj.artifact_types)): + raise ValueError(f"Unable to parse {type(self).__name__!r} response data") + + self.last_response = ArtifactTypesFragment.model_validate(conn) + + @property + def _length(self) -> None: + # TODO + return None + + @property + def more(self) -> bool: + if self.last_response is None: + return True + return self.last_response.page_info.has_next_page + + @property + def cursor(self) -> str | None: + if self.last_response is None: + return None + return self.last_response.edges[-1].cursor + + def update_variables(self) -> None: + self.variables.update({"cursor": self.cursor}) + + def convert_objects(self) -> list[ArtifactType]: + if self.last_response is None: + return [] + + return [ + ArtifactType( + client=self.client, + entity=self.entity, + project=self.project, + type_name=node.name, + attrs=node.model_dump(exclude_unset=True), + ) + for edge in self.last_response.edges + if edge.node and (node := ArtifactTypeFragment.model_validate(edge.node)) + ] + + +class ArtifactType: + def __init__( + self, + client: Client, + entity: str, + project: str, + type_name: str, + attrs: Mapping[str, Any] | None = None, + ): + self.client = client + self.entity = entity + self.project = project + self.type = type_name + self._attrs = attrs + if self._attrs is None: + self.load() + + def load(self) -> Mapping[str, Any]: + data: Mapping[str, Any] | None = self.client.execute( + gql(PROJECT_ARTIFACT_TYPE_GQL), + variable_values={ + "entityName": self.entity, + "projectName": self.project, + "artifactTypeName": self.type, + }, + ) + result = ProjectArtifactType.model_validate(data) + if not ((proj := result.project) and (artifact_type := proj.artifact_type)): + raise ValueError(f"Could not find artifact type {self.type}") + + self._attrs = artifact_type.model_dump(exclude_unset=True) + return self._attrs + + @property + def id(self) -> str: + return self._attrs["id"] + + @property + def name(self) -> str: + return self._attrs["name"] + + @normalize_exceptions + def collections(self, per_page: int = 50) -> ArtifactCollections: + """Artifact collections.""" + return ArtifactCollections(self.client, self.entity, self.project, self.type) + + def collection(self, name: str) -> ArtifactCollection: + return ArtifactCollection( + self.client, self.entity, self.project, name, self.type + ) + + def __repr__(self) -> str: + return f"" + + +class ArtifactCollections(SizedPaginator["ArtifactCollection"]): + last_response: ArtifactCollectionsFragment | None + + def __init__( + self, + client: Client, + entity: str, + project: str, + type_name: str, + per_page: int = 50, + ): + self.entity = entity + self.project = project + self.type_name = type_name + + variable_values = { + "entityName": entity, + "projectName": project, + "artifactTypeName": type_name, + } + + if server_supports_artifact_collections_gql_edges(client): + rename_fields = None + else: + rename_fields = {"artifactCollections": "artifactSequences"} + + self.QUERY = gql_compat( + PROJECT_ARTIFACT_COLLECTIONS_GQL, rename_fields=rename_fields + ) + + super().__init__(client, variable_values, per_page) + + @override + def _update_response(self) -> None: + """Fetch and validate the response data for the current page.""" + data = self.client.execute(self.QUERY, variable_values=self.variables) + result = ProjectArtifactCollections.model_validate(data) + + # Extract the inner `*Connection` result for faster/easier access. + if not ( + (proj := result.project) + and (type_ := proj.artifact_type) + and (conn := type_.artifact_collections) + ): + raise ValueError(f"Unable to parse {type(self).__name__!r} response data") + + self.last_response = ArtifactCollectionsFragment.model_validate(conn) + + @property + def _length(self) -> int: + if self.last_response is None: + self._load_page() + return self.last_response.total_count + + @property + def more(self): + if self.last_response is None: + return True + return self.last_response.page_info.has_next_page + + @property + def cursor(self): + if self.last_response is None: + return None + return self.last_response.edges[-1].cursor + + def update_variables(self) -> None: + self.variables.update({"cursor": self.cursor}) + + def convert_objects(self) -> list[ArtifactCollection]: + if self.last_response is None: + return [] + + return [ + ArtifactCollection( + client=self.client, + entity=self.entity, + project=self.project, + name=node.name, + type=self.type_name, + ) + for edge in self.last_response.edges + if (node := edge.node) + ] + + +class ArtifactCollection: + def __init__( + self, + client: Client, + entity: str, + project: str, + name: str, + type: str, + organization: str | None = None, + attrs: Mapping[str, Any] | None = None, + is_sequence: bool | None = None, + ): + self.client = client + self.entity = entity + self.project = project + self._name = validate_artifact_name(name) + self._saved_name = name + self._type = type + self._saved_type = type + self._attrs = attrs + if is_sequence is not None: + self._is_sequence = is_sequence + if (attrs is None) or (is_sequence is None): + self.load() + self._aliases = [a["node"]["alias"] for a in self._attrs["aliases"]["edges"]] + self._description = self._attrs["description"] + self._created_at = self._attrs["createdAt"] + self._tags = [a["node"]["name"] for a in self._attrs["tags"]["edges"]] + self._saved_tags = copy(self._tags) + self.organization = organization + + @property + def id(self) -> str: + return self._attrs["id"] + + @normalize_exceptions + def artifacts(self, per_page: int = 50) -> Artifacts: + """Artifacts.""" + return Artifacts( + client=self.client, + entity=self.entity, + project=self.project, + collection_name=self._saved_name, + type=self._saved_type, + per_page=per_page, + ) + + @property + def aliases(self) -> list[str]: + """Artifact Collection Aliases.""" + return self._aliases + + @property + def created_at(self) -> str: + return self._created_at + + def load(self): + if server_supports_artifact_collections_gql_edges(self.client): + rename_fields = None + else: + rename_fields = {"artifactCollection": "artifactSequence"} + + response = self.client.execute( + gql_compat(PROJECT_ARTIFACT_COLLECTION_GQL, rename_fields=rename_fields), + variable_values={ + "entityName": self.entity, + "projectName": self.project, + "artifactTypeName": self._saved_type, + "artifactCollectionName": self._saved_name, + }, + ) + + result = ProjectArtifactCollection.model_validate(response) + + if not ( + result.project + and (proj := result.project) + and (type_ := proj.artifact_type) + and (collection := type_.artifact_collection) + ): + raise ValueError(f"Could not find artifact type {self._saved_type}") + + sequence = type_.artifact_sequence + self._is_sequence = ( + sequence is not None + ) and sequence.typename__ == SOURCE_ARTIFACT_COLLECTION_TYPE + + if self._attrs is None: + self._attrs = collection.model_dump(exclude_unset=True) + return self._attrs + + @normalize_exceptions + def change_type(self, new_type: str) -> None: + """Deprecated, change type directly with `save` instead.""" + deprecate.deprecate( + field_name=Deprecated.artifact_collection__change_type, + warning_message="ArtifactCollection.change_type(type) is deprecated, use ArtifactCollection.save() instead.", + ) + + if self._saved_type != new_type: + try: + validate_artifact_type(self._saved_type, self.name) + except ValueError as e: + raise ValueError( + f"The current type '{self._saved_type!r}' is an internal type and cannot be changed." + ) from e + + # Check that the new type is not going to conflict with internal types + validate_artifact_type(new_type, self.name) + + if not self.is_sequence(): + raise ValueError("Artifact collection needs to be a sequence") + termlog( + f"Changing artifact collection type of {self._saved_type} to {new_type}" + ) + self.client.execute( + gql(MOVE_ARTIFACT_COLLECTION_GQL), + variable_values={ + "artifactSequenceID": self.id, + "destinationArtifactTypeName": new_type, + }, + ) + self._saved_type = new_type + self._type = new_type + + def is_sequence(self) -> bool: + """Return whether the artifact collection is a sequence.""" + return self._is_sequence + + @normalize_exceptions + def delete(self) -> None: + """Delete the entire artifact collection.""" + self.client.execute( + gql( + DELETE_ARTIFACT_SEQUENCE_GQL + if self.is_sequence() + else DELETE_ARTIFACT_PORTFOLIO_GQL + ), + variable_values={"id": self.id}, + ) + + @property + def description(self) -> str: + """A description of the artifact collection.""" + return self._description + + @description.setter + def description(self, description: str | None) -> None: + self._description = description + + @property + def tags(self) -> list[str]: + """The tags associated with the artifact collection.""" + return self._tags + + @tags.setter + def tags(self, tags: list[str]) -> None: + if any(not re.match(r"^[-\w]+([ ]+[-\w]+)*$", tag) for tag in tags): + raise ValueError( + "Tags must only contain alphanumeric characters or underscores separated by spaces or hyphens" + ) + self._tags = tags + + @property + def name(self) -> str: + """The name of the artifact collection.""" + return self._name + + @name.setter + def name(self, name: str) -> None: + self._name = validate_artifact_name(name) + + @property + def type(self): + """The type of the artifact collection.""" + return self._type + + @type.setter + def type(self, type: list[str]) -> None: + if not self.is_sequence(): + raise ValueError( + "Type can only be changed if the artifact collection is a sequence." + ) + self._type = type + + def _update_collection(self) -> None: + self.client.execute( + gql( + UPDATE_ARTIFACT_SEQUENCE_GQL + if self.is_sequence() + else UPDATE_ARTIFACT_PORTFOLIO_GQL + ), + variable_values={ + "id": self.id, + "name": self.name, + "description": self.description, + }, + ) + self._saved_name = self._name + + def _update_collection_type(self) -> None: + self.client.execute( + gql(MOVE_ARTIFACT_COLLECTION_GQL), + variable_values={ + "artifactSequenceID": self.id, + "destinationArtifactTypeName": self.type, + }, + ) + self._saved_type = self._type + + def _add_tags(self, tags_to_add: Iterable[str]) -> None: + self.client.execute( + gql(CREATE_ARTIFACT_COLLECTION_TAG_ASSIGNMENTS_GQL), + variable_values={ + "entityName": self.entity, + "projectName": self.project, + "artifactCollectionName": self._saved_name, + "tags": [{"tagName": tag} for tag in tags_to_add], + }, + ) + + def _delete_tags(self, tags_to_delete: Iterable[str]) -> None: + self.client.execute( + gql(DELETE_ARTIFACT_COLLECTION_TAG_ASSIGNMENTS_GQL), + variable_values={ + "entityName": self.entity, + "projectName": self.project, + "artifactCollectionName": self._saved_name, + "tags": [{"tagName": tag} for tag in tags_to_delete], + }, + ) + + @normalize_exceptions + def save(self) -> None: + """Persist any changes made to the artifact collection.""" + if self._saved_type != self.type: + try: + validate_artifact_type(self.type, self._name) + except ValueError as e: + raise ValueError(f"Failed to save artifact collection: {e}") from e + try: + validate_artifact_type(self._saved_type, self._name) + except ValueError as e: + raise ValueError( + f"Failed to save artifact collection '{self._name}': " + f"The current type '{self._saved_type!r}' is an internal type and cannot be changed." + ) from e + + self._update_collection() + + if self.is_sequence() and (self._saved_type != self._type): + self._update_collection_type() + + current_tags = set(self._tags) + saved_tags = set(self._saved_tags) + if tags_to_add := (current_tags - saved_tags): + self._add_tags(tags_to_add) + if tags_to_delete := (saved_tags - current_tags): + self._delete_tags(tags_to_delete) + self._saved_tags = copy(self._tags) + + def __repr__(self) -> str: + return f"" + + +class Artifacts(SizedPaginator["Artifact"]): + """An iterable collection of artifact versions associated with a project and optional filter. + + This is generally used indirectly via the `Api`.artifact_versions method. + """ + + last_response: ArtifactsFragment | None + + def __init__( + self, + client: Client, + entity: str, + project: str, + collection_name: str, + type: str, + filters: Mapping[str, Any] | None = None, + order: str | None = None, + per_page: int = 50, + tags: str | list[str] | None = None, + ): + self.entity = entity + self.collection_name = collection_name + self.type = type + self.project = project + self.filters = {"state": "COMMITTED"} if filters is None else filters + self.tags = [tags] if isinstance(tags, str) else tags + self.order = order + variables = { + "project": self.project, + "entity": self.entity, + "order": self.order, + "type": self.type, + "collection": self.collection_name, + "filters": json.dumps(self.filters), + } + + if server_supports_artifact_collections_gql_edges(client): + rename_fields = None + else: + rename_fields = {"artifactCollection": "artifactSequence"} + + self.QUERY = gql_compat( + PROJECT_ARTIFACTS_GQL, + omit_fields=omit_artifact_fields(api=InternalApi()), + rename_fields=rename_fields, + ) + + super().__init__(client, variables, per_page) + + @override + def _update_response(self) -> None: + data = self.client.execute(self.QUERY, variable_values=self.variables) + result = ProjectArtifacts.model_validate(data) + + # Extract the inner `*Connection` result for faster/easier access. + if not ( + (proj := result.project) + and (type_ := proj.artifact_type) + and (collection := type_.artifact_collection) + and (conn := collection.artifacts) + ): + raise ValueError(f"Unable to parse {type(self).__name__!r} response data") + + self.last_response = ArtifactsFragment.model_validate(conn) + + @property + def _length(self) -> int: + if self.last_response is None: + self._load_page() + return self.last_response.total_count + + @property + def more(self) -> bool: + if self.last_response is None: + return True + return self.last_response.page_info.has_next_page + + @property + def cursor(self) -> str | None: + if self.last_response is None: + return None + return self.last_response.edges[-1].cursor + + def convert_objects(self) -> list[Artifact]: + if self.last_response is None: + return [] + + artifact_edges = (edge for edge in self.last_response.edges if edge.node) + artifacts = ( + wandb.Artifact._from_attrs( + entity=self.entity, + project=self.project, + name=f"{self.collection_name}:{edge.version}", + attrs=edge.node.model_dump(exclude_unset=True), + client=self.client, + ) + for edge in artifact_edges + ) + required_tags = set(self.tags or []) + return [art for art in artifacts if required_tags.issubset(art.tags)] + + +class RunArtifacts(SizedPaginator["Artifact"]): + last_response: ( + RunOutputArtifactsProjectRunOutputArtifacts + | RunInputArtifactsProjectRunInputArtifacts + ) + + #: The pydantic model used to parse the (inner part of the) raw response. + _response_cls: type[ + RunOutputArtifactsProjectRunOutputArtifacts + | RunInputArtifactsProjectRunInputArtifacts + ] + + def __init__( + self, + client: Client, + run: Run, + mode: Literal["logged", "used"] = "logged", + per_page: int = 50, + ): + self.run = run + + if mode == "logged": + self.run_key = "outputArtifacts" + self.QUERY = gql_compat( + RUN_OUTPUT_ARTIFACTS_GQL, + omit_fields=omit_artifact_fields(api=InternalApi()), + ) + self._response_cls = RunOutputArtifactsProjectRunOutputArtifacts + elif mode == "used": + self.run_key = "inputArtifacts" + self.QUERY = gql_compat( + RUN_INPUT_ARTIFACTS_GQL, + omit_fields=omit_artifact_fields(api=InternalApi()), + ) + self._response_cls = RunInputArtifactsProjectRunInputArtifacts + else: + raise ValueError("mode must be logged or used") + + variable_values = { + "entity": run.entity, + "project": run.project, + "runName": run.id, + } + super().__init__(client, variable_values, per_page) + + @override + def _update_response(self) -> None: + data = self.client.execute(self.QUERY, variable_values=self.variables) + + # Extract the inner `*Connection` result for faster/easier access. + inner_data = data["project"]["run"][self.run_key] + self.last_response = self._response_cls.model_validate(inner_data) + + @property + def _length(self) -> int: + if self.last_response is None: + self._load_page() + return self.last_response.total_count + + @property + def more(self) -> bool: + if self.last_response is None: + return True + return self.last_response.page_info.has_next_page + + @property + def cursor(self) -> str | None: + if self.last_response is None: + return None + return self.last_response.edges[-1].cursor + + def convert_objects(self) -> list[Artifact]: + if self.last_response is None: + return [] + + return [ + wandb.Artifact._from_attrs( + entity=proj.entity_name, + project=proj.name, + name=f"{artifact_seq.name}:v{node.version_index}", + attrs=node.model_dump(exclude_unset=True), + client=self.client, + ) + for edge in self.last_response.edges + if (node := edge.node) + and (artifact_seq := node.artifact_sequence) + and (proj := artifact_seq.project) + ] + + +class ArtifactFiles(SizedPaginator["public.File"]): + last_response: FilesFragment | None + + def __init__( + self, + client: Client, + artifact: Artifact, + names: Sequence[str] | None = None, + per_page: int = 50, + ): + self.query_via_membership = InternalApi()._server_supports( + ServerFeature.ARTIFACT_COLLECTION_MEMBERSHIP_FILES + ) + self.artifact = artifact + + if self.query_via_membership: + query_str = ARTIFACT_COLLECTION_MEMBERSHIP_FILES_GQL + variables = { + "entityName": artifact.entity, + "projectName": artifact.project, + "artifactName": artifact.name.split(":")[0], + "artifactVersionIndex": artifact.version, + "fileNames": names, + } + else: + query_str = ARTIFACT_VERSION_FILES_GQL + variables = { + "entityName": artifact.source_entity, + "projectName": artifact.source_project, + "artifactName": artifact.source_name, + "artifactTypeName": artifact.type, + "fileNames": names, + } + + # The server must advertise at least SDK 0.12.21 + # to get storagePath + if not client.version_supported("0.12.21"): + self.QUERY = gql_compat(query_str, omit_fields={"storagePath"}) + else: + self.QUERY = gql(query_str) + + super().__init__(client, variables, per_page) + + @override + def _update_response(self) -> None: + data = self.client.execute(self.QUERY, variable_values=self.variables) + + # Extract the inner `*Connection` result for faster/easier access. + if self.query_via_membership: + result = ArtifactCollectionMembershipFiles.model_validate(data) + conn = result.project.artifact_collection.artifact_membership.files + else: + result = ArtifactVersionFiles.model_validate(data) + conn = result.project.artifact_type.artifact.files + + if conn is None: + raise ValueError(f"Unable to parse {type(self).__name__!r} response data") + + self.last_response = FilesFragment.model_validate(conn) + + @property + def path(self) -> list[str]: + return [self.artifact.entity, self.artifact.project, self.artifact.name] + + @property + def _length(self) -> int: + if self.last_response is None: + self._load_page() + return self.artifact.file_count + + @property + def more(self) -> bool: + if self.last_response is None: + return True + return self.last_response.page_info.has_next_page + + @property + def cursor(self) -> str | None: + if self.last_response is None: + return None + return self.last_response.edges[-1].cursor + + def update_variables(self) -> None: + self.variables.update({"fileLimit": self.per_page, "fileCursor": self.cursor}) + + def convert_objects(self) -> list[public.File]: + if self.last_response is None: + return [] + + return [ + public.File( + client=self.client, + attrs=node.model_dump(exclude_unset=True), + ) + for edge in self.last_response.edges + if (node := edge.node) + ] + + def __repr__(self) -> str: + path_str = "/".join(self.path) + return f"" + + +def server_supports_artifact_collections_gql_edges( + client: RetryingClient, warn: bool = False +) -> bool: + # TODO: Validate this version + # Edges were merged into core on Mar 2, 2022: https://github.com/wandb/core/commit/81c90b29eaacfe0a96dc1ebd83c53560ca763e8b + # CLI version was bumped to "0.12.11" on Mar 3, 2022: https://github.com/wandb/core/commit/328396fa7c89a2178d510a1be9c0d4451f350d7b + supported = client.version_supported("0.12.11") # edges were merged on + if not supported and warn: + # First local release to include the above is 0.9.50: https://github.com/wandb/local/releases/tag/0.9.50 + wandb.termwarn( + "W&B Local Server version does not support ArtifactCollection gql edges; falling back to using legacy ArtifactSequence. Please update server to at least version 0.9.50." + ) + return supported diff --git a/lib/python3.12/site-packages/wandb/apis/public/automations.py b/lib/python3.12/site-packages/wandb/apis/public/automations.py new file mode 100644 index 0000000000000000000000000000000000000000..eaa19b5471061626e540d01e967b0f5d15285d4e --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/public/automations.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +from itertools import chain +from typing import TYPE_CHECKING, Any, Iterable, Mapping + +from pydantic import ValidationError +from typing_extensions import override +from wandb_graphql.language.ast import Document + +from wandb.apis.paginator import Paginator, _Client + +if TYPE_CHECKING: + from wandb.automations import Automation + from wandb.automations._generated import ProjectConnectionFields + + +class Automations(Paginator["Automation"]): + last_response: ProjectConnectionFields | None + _query: Document + + def __init__( + self, + client: _Client, + variables: Mapping[str, Any], + per_page: int = 50, + _query: Document | None = None, + ): + super().__init__(client, variables, per_page=per_page) + if _query is None: + raise RuntimeError(f"Query required for {type(self).__qualname__}") + self._query = _query + + @property + def more(self) -> bool: + """Whether there are more items to fetch.""" + if self.last_response is None: + return True + return self.last_response.page_info.has_next_page + + @property + def cursor(self) -> str | None: + """The start cursor to use for the next page.""" + if self.last_response is None: + return None + return self.last_response.page_info.end_cursor + + @override + def _update_response(self) -> None: + """Fetch the raw response data for the current page.""" + from wandb.automations._generated import ProjectConnectionFields + + data: dict[str, Any] = self.client.execute( + self._query, variable_values=self.variables + ) + try: + page_data = data["searchScope"]["projects"] + self.last_response = ProjectConnectionFields.model_validate(page_data) + except (LookupError, AttributeError, ValidationError) as e: + raise ValueError("Unexpected response data") from e + + def convert_objects(self) -> Iterable[Automation]: + """Parse the page data into a list of objects.""" + from wandb.automations import Automation + + page = self.last_response + return [ + Automation.model_validate(obj) + for obj in chain.from_iterable(edge.node.triggers for edge in page.edges) + ] diff --git a/lib/python3.12/site-packages/wandb/apis/public/const.py b/lib/python3.12/site-packages/wandb/apis/public/const.py new file mode 100644 index 0000000000000000000000000000000000000000..77af2da729d553a25f97a864546639f520b8817a --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/public/const.py @@ -0,0 +1,4 @@ +import datetime + +# Only retry requests for 20 seconds in the public api +RETRY_TIMEDELTA = datetime.timedelta(seconds=20) diff --git a/lib/python3.12/site-packages/wandb/apis/public/files.py b/lib/python3.12/site-packages/wandb/apis/public/files.py new file mode 100644 index 0000000000000000000000000000000000000000..771cdc51bbe9cf756d4a8f632f5725ca8964c9f0 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/public/files.py @@ -0,0 +1,263 @@ +"""Public API: files.""" + +import io +import os +from typing import Optional + +import requests +from wandb_gql import gql +from wandb_gql.client import RetryError + +import wandb +from wandb import util +from wandb.apis.attrs import Attrs +from wandb.apis.normalize import normalize_exceptions +from wandb.apis.paginator import SizedPaginator +from wandb.apis.public import utils +from wandb.apis.public.api import Api +from wandb.apis.public.const import RETRY_TIMEDELTA +from wandb.sdk.lib import retry + +FILE_FRAGMENT = """fragment RunFilesFragment on Run { + files(names: $fileNames, after: $fileCursor, first: $fileLimit) { + edges { + node { + id + name + url(upload: $upload) + directUrl + sizeBytes + mimetype + updatedAt + md5 + } + cursor + } + pageInfo { + endCursor + hasNextPage + } + } +}""" + + +class Files(SizedPaginator["File"]): + """An iterable collection of `File` objects.""" + + QUERY = gql( + """ + query RunFiles($project: String!, $entity: String!, $name: String!, $fileCursor: String, + $fileLimit: Int = 50, $fileNames: [String] = [], $upload: Boolean = false) {{ + project(name: $project, entityName: $entity) {{ + internalId + run(name: $name) {{ + fileCount + ...RunFilesFragment + }} + }} + }} + {} + """.format(FILE_FRAGMENT) + ) + + def __init__(self, client, run, names=None, per_page=50, upload=False): + self.run = run + variables = { + "project": run.project, + "entity": run.entity, + "name": run.id, + "fileNames": names or [], + "upload": upload, + } + super().__init__(client, variables, per_page) + + @property + def _length(self): + if not self.last_response: + self._load_page() + + return self.last_response["project"]["run"]["fileCount"] + + @property + def more(self): + if self.last_response: + return self.last_response["project"]["run"]["files"]["pageInfo"][ + "hasNextPage" + ] + else: + return True + + @property + def cursor(self): + if self.last_response: + return self.last_response["project"]["run"]["files"]["edges"][-1]["cursor"] + else: + return None + + def update_variables(self): + self.variables.update({"fileLimit": self.per_page, "fileCursor": self.cursor}) + + def convert_objects(self): + return [ + File(self.client, r["node"], self.run) + for r in self.last_response["project"]["run"]["files"]["edges"] + ] + + def __repr__(self): + return "".format("/".join(self.run.path), len(self)) + + +class File(Attrs): + """File is a class associated with a file saved by wandb. + + Attributes: + name (string): filename + url (string): path to file + direct_url (string): path to file in the bucket + md5 (string): md5 of file + mimetype (string): mimetype of file + updated_at (string): timestamp of last update + size (int): size of file in bytes + path_uri (str): path to file in the bucket, currently only available for files stored in S3 + """ + + def __init__(self, client, attrs, run=None): + self.client = client + self._attrs = attrs + self.run = run + self.server_supports_delete_file_with_project_id: Optional[bool] = None + super().__init__(dict(attrs)) + + @property + def size(self): + size_bytes = self._attrs["sizeBytes"] + if size_bytes is not None: + return int(size_bytes) + return 0 + + @property + def path_uri(self) -> str: + """ + Returns the uri path to the file in the storage bucket. + """ + path_uri = "" + try: + path_uri = utils.parse_s3_url_to_s3_uri(self._attrs["directUrl"]) + except ValueError: + wandb.termwarn("path_uri is only available for files stored in S3") + except LookupError: + wandb.termwarn("Unable to find direct_url of file") + return path_uri + + @normalize_exceptions + @retry.retriable( + retry_timedelta=RETRY_TIMEDELTA, + check_retry_fn=util.no_retry_auth, + retryable_exceptions=(RetryError, requests.RequestException), + ) + def download( + self, + root: str = ".", + replace: bool = False, + exist_ok: bool = False, + api: Optional[Api] = None, + ) -> io.TextIOWrapper: + """Downloads a file previously saved by a run from the wandb server. + + Args: + replace (boolean): If `True`, download will overwrite a local file + if it exists. Defaults to `False`. + root (str): Local directory to save the file. Defaults to ".". + exist_ok (boolean): If `True`, will not raise ValueError if file already + exists and will not re-download unless replace=True. Defaults to `False`. + api (Api, optional): If given, the `Api` instance used to download the file. + + Raises: + `ValueError` if file already exists, replace=False and exist_ok=False. + """ + if api is None: + api = wandb.Api() + + path = os.path.join(root, self.name) + if os.path.exists(path) and not replace: + if exist_ok: + return open(path) + else: + raise ValueError( + "File already exists, pass replace=True to overwrite or exist_ok=True to leave it as is and don't error." + ) + + util.download_file_from_url(path, self.url, api.api_key) + return open(path) + + @normalize_exceptions + def delete(self): + project_id_mutation_fragment = "" + project_id_variable_fragment = "" + variable_values = { + "files": [self.id], + } + + # Add projectId to mutation and variables if the server supports it. + # Otherwise, do not include projectId in mutation for older server versions which do not support it. + if self._server_accepts_project_id_for_delete_file(): + variable_values["projectId"] = self.run._project_internal_id + project_id_variable_fragment = ", $projectId: Int" + project_id_mutation_fragment = "projectId: $projectId" + + mutation_string = """ + mutation deleteFiles($files: [ID!]!{}) {{ + deleteFiles(input: {{ + files: $files + {} + }}) {{ + success + }} + }} + """.format(project_id_variable_fragment, project_id_mutation_fragment) + mutation = gql(mutation_string) + + self.client.execute( + mutation, + variable_values=variable_values, + ) + + def __repr__(self): + return "".format( + self.name, + self.mimetype, + util.to_human_size(self.size, units=util.POW_2_BYTES), + ) + + @normalize_exceptions + def _server_accepts_project_id_for_delete_file(self) -> bool: + """Returns True if the server supports deleting files with a projectId. + + This check is done by utilizing GraphQL introspection in the available fields on the DeleteFiles API. + """ + query_string = """ + query ProbeDeleteFilesProjectIdInput { + DeleteFilesProjectIdInputType: __type(name:"DeleteFilesInput") { + inputFields{ + name + } + } + } + """ + + # Only perform the query once to avoid extra network calls + if self.server_supports_delete_file_with_project_id is None: + query = gql(query_string) + res = self.client.execute(query) + + # If projectId is in the inputFields, the server supports deleting files with a projectId + self.server_supports_delete_file_with_project_id = "projectId" in [ + x["name"] + for x in ( + res.get("DeleteFilesProjectIdInputType", {}).get( + "inputFields", [{}] + ) + ) + ] + + return self.server_supports_delete_file_with_project_id diff --git a/lib/python3.12/site-packages/wandb/apis/public/history.py b/lib/python3.12/site-packages/wandb/apis/public/history.py new file mode 100644 index 0000000000000000000000000000000000000000..1e6ddd8a7d60ccc8269fda3857c942d89ef62d11 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/public/history.py @@ -0,0 +1,149 @@ +"""Public API: history.""" + +import json + +import requests +from wandb_gql import gql +from wandb_gql.client import RetryError + +from wandb import util +from wandb.apis.normalize import normalize_exceptions +from wandb.sdk.lib import retry + + +class HistoryScan: + QUERY = gql( + """ + query HistoryPage($entity: String!, $project: String!, $run: String!, $minStep: Int64!, $maxStep: Int64!, $pageSize: Int!) { + project(name: $project, entityName: $entity) { + run(name: $run) { + history(minStep: $minStep, maxStep: $maxStep, samples: $pageSize) + } + } + } + """ + ) + + def __init__(self, client, run, min_step, max_step, page_size=1000): + self.client = client + self.run = run + self.page_size = page_size + self.min_step = min_step + self.max_step = max_step + self.page_offset = min_step # minStep for next page + self.scan_offset = 0 # index within current page of rows + self.rows = [] # current page of rows + + def __iter__(self): + self.page_offset = self.min_step + self.scan_offset = 0 + self.rows = [] + return self + + def __next__(self): + while True: + if self.scan_offset < len(self.rows): + row = self.rows[self.scan_offset] + self.scan_offset += 1 + return row + if self.page_offset >= self.max_step: + raise StopIteration() + self._load_next() + + next = __next__ + + @normalize_exceptions + @retry.retriable( + check_retry_fn=util.no_retry_auth, + retryable_exceptions=(RetryError, requests.RequestException), + ) + def _load_next(self): + max_step = self.page_offset + self.page_size + if max_step > self.max_step: + max_step = self.max_step + variables = { + "entity": self.run.entity, + "project": self.run.project, + "run": self.run.id, + "minStep": int(self.page_offset), + "maxStep": int(max_step), + "pageSize": int(self.page_size), + } + + res = self.client.execute(self.QUERY, variable_values=variables) + res = res["project"]["run"]["history"] + self.rows = [json.loads(row) for row in res] + self.page_offset += self.page_size + self.scan_offset = 0 + + +class SampledHistoryScan: + QUERY = gql( + """ + query SampledHistoryPage($entity: String!, $project: String!, $run: String!, $spec: JSONString!) { + project(name: $project, entityName: $entity) { + run(name: $run) { + sampledHistory(specs: [$spec]) + } + } + } + """ + ) + + def __init__(self, client, run, keys, min_step, max_step, page_size=1000): + self.client = client + self.run = run + self.keys = keys + self.page_size = page_size + self.min_step = min_step + self.max_step = max_step + self.page_offset = min_step # minStep for next page + self.scan_offset = 0 # index within current page of rows + self.rows = [] # current page of rows + + def __iter__(self): + self.page_offset = self.min_step + self.scan_offset = 0 + self.rows = [] + return self + + def __next__(self): + while True: + if self.scan_offset < len(self.rows): + row = self.rows[self.scan_offset] + self.scan_offset += 1 + return row + if self.page_offset >= self.max_step: + raise StopIteration() + self._load_next() + + next = __next__ + + @normalize_exceptions + @retry.retriable( + check_retry_fn=util.no_retry_auth, + retryable_exceptions=(RetryError, requests.RequestException), + ) + def _load_next(self): + max_step = self.page_offset + self.page_size + if max_step > self.max_step: + max_step = self.max_step + variables = { + "entity": self.run.entity, + "project": self.run.project, + "run": self.run.id, + "spec": json.dumps( + { + "keys": self.keys, + "minStep": int(self.page_offset), + "maxStep": int(max_step), + "samples": int(self.page_size), + } + ), + } + + res = self.client.execute(self.QUERY, variable_values=variables) + res = res["project"]["run"]["sampledHistory"] + self.rows = res[0] + self.page_offset += self.page_size + self.scan_offset = 0 diff --git a/lib/python3.12/site-packages/wandb/apis/public/integrations.py b/lib/python3.12/site-packages/wandb/apis/public/integrations.py new file mode 100644 index 0000000000000000000000000000000000000000..c314e501341e9160cb6bbc9a361e6bb9865a3ce8 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/public/integrations.py @@ -0,0 +1,180 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Iterable + +from pydantic import ValidationError +from typing_extensions import override +from wandb_gql import gql +from wandb_graphql.language.ast import Document + +from wandb.apis.paginator import Paginator + +if TYPE_CHECKING: + from wandb.apis.paginator import _Client + from wandb.automations import Integration, SlackIntegration, WebhookIntegration + from wandb.automations._generated import ( + GenericWebhookIntegrationConnectionFields, + IntegrationConnectionFields, + SlackIntegrationConnectionFields, + ) + + +class Integrations(Paginator["Integration"]): + last_response: IntegrationConnectionFields | None + _query: Document + + def __init__(self, client: _Client, variables: dict[str, Any], per_page: int = 50): + from wandb.automations._generated import INTEGRATIONS_BY_ENTITY_GQL + + super().__init__(client, variables, per_page=per_page) + # All integrations for entity + self._query = gql(INTEGRATIONS_BY_ENTITY_GQL) + + @property + def more(self) -> bool: + """Whether there are more Integrations to fetch.""" + if self.last_response is None: + return True + return self.last_response.page_info.has_next_page + + @property + def cursor(self) -> str | None: + """The start cursor to use for the next page.""" + if self.last_response is None: + return None + return self.last_response.page_info.end_cursor + + @override + def _update_response(self) -> None: + """Fetch and parse the response data for the current page.""" + from wandb.automations._generated import IntegrationConnectionFields + + data: dict[str, Any] = self.client.execute( + self._query, variable_values=self.variables + ) + try: + page_data = data["entity"]["integrations"] + self.last_response = IntegrationConnectionFields.model_validate(page_data) + except (LookupError, AttributeError, ValidationError) as e: + raise ValueError("Unexpected response data") from e + + def convert_objects(self) -> Iterable[Integration]: + """Parse the page data into a list of integrations.""" + from wandb.automations.integrations import _IntegrationEdge + + page = self.last_response + return [_IntegrationEdge.model_validate(edge).node for edge in page.edges] + + +class WebhookIntegrations(Paginator["WebhookIntegration"]): + last_response: GenericWebhookIntegrationConnectionFields | None + _query: Document + + def __init__(self, client: _Client, variables: dict[str, Any], per_page: int = 50): + from wandb.automations._generated import ( + GENERIC_WEBHOOK_INTEGRATIONS_BY_ENTITY_GQL, + ) + + super().__init__(client, variables, per_page=per_page) + # Webhook integrations for entity + self._query = gql(GENERIC_WEBHOOK_INTEGRATIONS_BY_ENTITY_GQL) + + @property + def more(self) -> bool: + """Whether there are more webhook integrations to fetch.""" + if self.last_response is None: + return True + return self.last_response.page_info.has_next_page + + @property + def cursor(self) -> str | None: + """The start cursor to use for the next page.""" + if self.last_response is None: + return None + return self.last_response.page_info.end_cursor + + @override + def _update_response(self) -> None: + """Fetch and parse the response data for the current page.""" + from wandb.automations._generated import ( + GenericWebhookIntegrationConnectionFields, + ) + + data: dict[str, Any] = self.client.execute( + self._query, variable_values=self.variables + ) + try: + page_data = data["entity"]["integrations"] + self.last_response = ( + GenericWebhookIntegrationConnectionFields.model_validate(page_data) + ) + except (LookupError, AttributeError, ValidationError) as e: + raise ValueError("Unexpected response data") from e + + def convert_objects(self) -> Iterable[WebhookIntegration]: + """Parse the page data into a list of webhook integrations.""" + from wandb.automations import WebhookIntegration + + typename = "GenericWebhookIntegration" + return [ + # Filter on typename__ needed because the GQL response still + # includes all integration types + WebhookIntegration.model_validate(node) + for edge in self.last_response.edges + if (node := edge.node) and (node.typename__ == typename) + ] + + +class SlackIntegrations(Paginator["SlackIntegration"]): + last_response: SlackIntegrationConnectionFields | None + _query: Document + + def __init__(self, client: _Client, variables: dict[str, Any], per_page: int = 50): + from wandb.automations._generated import SLACK_INTEGRATIONS_BY_ENTITY_GQL + + super().__init__(client, variables, per_page=per_page) + # Slack integrations for entity + self._query = gql(SLACK_INTEGRATIONS_BY_ENTITY_GQL) + + @property + def more(self) -> bool: + """Whether there are more Slack integrations to fetch.""" + if self.last_response is None: + return True + return self.last_response.page_info.has_next_page + + @property + def cursor(self) -> str | None: + """The start cursor to use for the next page.""" + if self.last_response is None: + return None + return self.last_response.page_info.end_cursor + + @override + def _update_response(self) -> None: + """Fetch and parse the response data for the current page.""" + from wandb.automations._generated import SlackIntegrationConnectionFields + + data: dict[str, Any] = self.client.execute( + self._query, variable_values=self.variables + ) + try: + page_data = data["entity"]["integrations"] + self.last_response = SlackIntegrationConnectionFields.model_validate( + page_data + ) + except (LookupError, AttributeError, ValidationError) as e: + raise ValueError("Unexpected response data") from e + + def convert_objects(self) -> Iterable[SlackIntegration]: + """Parse the page data into a list of Slack integrations.""" + from wandb.automations import SlackIntegration + + typename = "SlackIntegration" + return [ + # Filter on typename__ needed because the GQL response still + # includes all integration types + SlackIntegration.model_validate(node) + for edge in self.last_response.edges + if (node := edge.node) and (node.typename__ == typename) + ] diff --git a/lib/python3.12/site-packages/wandb/apis/public/jobs.py b/lib/python3.12/site-packages/wandb/apis/public/jobs.py new file mode 100644 index 0000000000000000000000000000000000000000..2a4bfc3c1e03cfdf9013faf82446fa35c90ed559 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/public/jobs.py @@ -0,0 +1,654 @@ +"""Public API: jobs.""" + +import json +import os +import shutil +import time +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Mapping, Optional + +from wandb_gql import gql + +import wandb +from wandb import util +from wandb.apis import public +from wandb.apis.normalize import normalize_exceptions +from wandb.errors import CommError +from wandb.sdk.artifacts.artifact_state import ArtifactState +from wandb.sdk.data_types._dtypes import InvalidType, Type, TypeRegistry +from wandb.sdk.launch.errors import LaunchError +from wandb.sdk.launch.utils import ( + LAUNCH_DEFAULT_PROJECT, + _fetch_git_repo, + apply_patch, + convert_jupyter_notebook_to_script, +) + +if TYPE_CHECKING: + from wandb.apis.public import Api, RetryingClient + + +class Job: + _name: str + _input_types: Type + _output_types: Type + _entity: str + _project: str + _entrypoint: List[str] + _notebook_job: bool + _partial: bool + + def __init__(self, api: "Api", name, path: Optional[str] = None) -> None: + try: + self._job_artifact = api._artifact(name, type="job") + except CommError: + raise CommError(f"Job artifact {name} not found") + if path: + self._fpath = path + self._job_artifact.download(root=path) + else: + self._fpath = self._job_artifact.download() + self._name = name + self._api = api + self._entity = api.default_entity + + with open(os.path.join(self._fpath, "wandb-job.json")) as f: + self._job_info: Mapping[str, Any] = json.load(f) + source_info = self._job_info.get("source", {}) + # only use notebook job if entrypoint not set and notebook is set + self._notebook_job = source_info.get("notebook", False) + self._entrypoint = source_info.get("entrypoint") + self._dockerfile = source_info.get("dockerfile") + self._build_context = source_info.get("build_context") + self._base_image = source_info.get("base_image") + self._args = source_info.get("args") + self._partial = self._job_info.get("_partial", False) + self._requirements_file = os.path.join(self._fpath, "requirements.frozen.txt") + self._input_types = TypeRegistry.type_from_dict( + self._job_info.get("input_types") + ) + self._output_types = TypeRegistry.type_from_dict( + self._job_info.get("output_types") + ) + if self._job_info.get("source_type") == "artifact": + self._set_configure_launch_project(self._configure_launch_project_artifact) + if self._job_info.get("source_type") == "repo": + self._set_configure_launch_project(self._configure_launch_project_repo) + if self._job_info.get("source_type") == "image": + self._set_configure_launch_project(self._configure_launch_project_container) + + @property + def name(self): + return self._name + + def _set_configure_launch_project(self, func): + self.configure_launch_project = func + + def _get_code_artifact(self, artifact_string): + artifact_string, base_url, is_id = util.parse_artifact_string(artifact_string) + if is_id: + code_artifact = wandb.Artifact._from_id(artifact_string, self._api._client) + else: + code_artifact = self._api._artifact(name=artifact_string, type="code") + if code_artifact is None: + raise LaunchError("No code artifact found") + if code_artifact.state == ArtifactState.DELETED: + raise LaunchError( + f"Job {self.name} references deleted code artifact {code_artifact.name}" + ) + return code_artifact + + def _configure_launch_project_notebook(self, launch_project): + new_fname = convert_jupyter_notebook_to_script( + self._entrypoint[-1], launch_project.project_dir + ) + new_entrypoint = self._entrypoint + new_entrypoint[-1] = new_fname + launch_project.set_job_entry_point(new_entrypoint) + + def _configure_launch_project_repo(self, launch_project): + git_info = self._job_info.get("source", {}).get("git", {}) + _fetch_git_repo( + launch_project.project_dir, + git_info["remote"], + git_info["commit"], + ) + if os.path.exists(os.path.join(self._fpath, "diff.patch")): + with open(os.path.join(self._fpath, "diff.patch")) as f: + apply_patch(f.read(), launch_project.project_dir) + shutil.copy(self._requirements_file, launch_project.project_dir) + launch_project.python_version = self._job_info.get("runtime") + if self._notebook_job: + self._configure_launch_project_notebook(launch_project) + else: + launch_project.set_job_entry_point(self._entrypoint) + + if self._dockerfile: + launch_project.set_job_dockerfile(self._dockerfile) + if self._build_context: + launch_project.set_job_build_context(self._build_context) + if self._base_image: + launch_project.set_job_base_image(self._base_image) + + def _configure_launch_project_artifact(self, launch_project): + artifact_string = self._job_info.get("source", {}).get("artifact") + if artifact_string is None: + raise LaunchError(f"Job {self.name} had no source artifact") + + code_artifact = self._get_code_artifact(artifact_string) + launch_project.python_version = self._job_info.get("runtime") + shutil.copy(self._requirements_file, launch_project.project_dir) + + code_artifact.download(launch_project.project_dir) + + if self._notebook_job: + self._configure_launch_project_notebook(launch_project) + else: + launch_project.set_job_entry_point(self._entrypoint) + + if self._dockerfile: + launch_project.set_job_dockerfile(self._dockerfile) + if self._build_context: + launch_project.set_job_build_context(self._build_context) + if self._base_image: + launch_project.set_job_base_image(self._base_image) + + def _configure_launch_project_container(self, launch_project): + launch_project.docker_image = self._job_info.get("source", {}).get("image") + if launch_project.docker_image is None: + raise LaunchError( + "Job had malformed source dictionary without an image key" + ) + if self._entrypoint: + launch_project.set_job_entry_point(self._entrypoint) + + def set_entrypoint(self, entrypoint: List[str]): + self._entrypoint = entrypoint + + def call( + self, + config, + project=None, + entity=None, + queue=None, + resource="local-container", + resource_args=None, + template_variables=None, + project_queue=None, + priority=None, + ): + from wandb.sdk.launch import _launch_add + + run_config = {} + for key, item in config.items(): + if util._is_artifact_object(item): + if isinstance(item, wandb.Artifact) and item.is_draft(): + raise ValueError("Cannot queue jobs with unlogged artifacts") + run_config[key] = util.artifact_to_json(item) + + run_config.update(config) + + assigned_config_type = self._input_types.assign(run_config) + if self._partial: + wandb.termwarn( + "Launching manually created job for the first time, can't verify types" + ) + else: + if isinstance(assigned_config_type, InvalidType): + raise TypeError(self._input_types.explain(run_config)) + + queued_run = _launch_add.launch_add( + job=self._name, + config={"overrides": {"run_config": run_config}}, + template_variables=template_variables, + project=project or self._project, + entity=entity or self._entity, + queue_name=queue, + resource=resource, + project_queue=project_queue, + resource_args=resource_args, + priority=priority, + ) + return queued_run + + +class QueuedRun: + """A single queued run associated with an entity and project. Call `run = queued_run.wait_until_running()` or `run = queued_run.wait_until_finished()` to access the run.""" + + def __init__( + self, + client, + entity, + project, + queue_name, + run_queue_item_id, + project_queue=LAUNCH_DEFAULT_PROJECT, + priority=None, + ): + self.client = client + self._entity = entity + self._project = project + self._queue_name = queue_name + self._run_queue_item_id = run_queue_item_id + self.sweep = None + self._run = None + self.project_queue = project_queue + self.priority = priority + + @property + def queue_name(self): + return self._queue_name + + @property + def id(self): + return self._run_queue_item_id + + @property + def project(self): + return self._project + + @property + def entity(self): + return self._entity + + @property + def state(self): + item = self._get_item() + if item: + return item["state"].lower() + + raise ValueError( + f"Could not find QueuedRunItem associated with id: {self.id} on queue {self.queue_name} at itemId: {self.id}" + ) + + @normalize_exceptions + def _get_run_queue_item_legacy(self) -> Dict: + query = gql( + """ + query GetRunQueueItem($projectName: String!, $entityName: String!, $runQueue: String!) { + project(name: $projectName, entityName: $entityName) { + runQueue(name:$runQueue) { + runQueueItems { + edges { + node { + id + state + associatedRunId + } + } + } + } + } + } + """ + ) + variable_values = { + "projectName": self.project_queue, + "entityName": self._entity, + "runQueue": self.queue_name, + } + res = self.client.execute(query, variable_values) + + for item in res["project"]["runQueue"]["runQueueItems"]["edges"]: + if str(item["node"]["id"]) == str(self.id): + return item["node"] + + @normalize_exceptions + def _get_item(self): + query = gql( + """ + query GetRunQueueItem($projectName: String!, $entityName: String!, $runQueue: String!, $itemId: ID!) { + project(name: $projectName, entityName: $entityName) { + runQueue(name: $runQueue) { + runQueueItem(id: $itemId) { + id + state + associatedRunId + } + } + } + } + """ + ) + variable_values = { + "projectName": self.project_queue, + "entityName": self._entity, + "runQueue": self.queue_name, + "itemId": self.id, + } + try: + res = self.client.execute(query, variable_values) # exception w/ old server + if res["project"]["runQueue"].get("runQueueItem") is not None: + return res["project"]["runQueue"]["runQueueItem"] + except Exception as e: + if "Cannot query field" not in str(e): + raise LaunchError(f"Unknown exception: {e}") + + return self._get_run_queue_item_legacy() + + @normalize_exceptions + def wait_until_finished(self): + if not self._run: + self.wait_until_running() + + self._run.wait_until_finished() + # refetch run to get updated summary + self._run.load(force=True) + return self._run + + @normalize_exceptions + def delete(self, delete_artifacts=False): + """Delete the given queued run from the wandb backend.""" + query = gql( + """ + query fetchRunQueuesFromProject($entityName: String!, $projectName: String!, $runQueueName: String!) { + project(name: $projectName, entityName: $entityName) { + runQueue(name: $runQueueName) { + id + } + } + } + """ + ) + + res = self.client.execute( + query, + variable_values={ + "entityName": self.entity, + "projectName": self.project_queue, + "runQueueName": self.queue_name, + }, + ) + + if res["project"].get("runQueue") is not None: + queue_id = res["project"]["runQueue"]["id"] + + mutation = gql( + """ + mutation DeleteFromRunQueue( + $queueID: ID!, + $runQueueItemId: ID! + ) { + deleteFromRunQueue(input: { + queueID: $queueID + runQueueItemId: $runQueueItemId + }) { + success + clientMutationId + } + } + """ + ) + self.client.execute( + mutation, + variable_values={ + "queueID": queue_id, + "runQueueItemId": self._run_queue_item_id, + }, + ) + + @normalize_exceptions + def wait_until_running(self): + if self._run is not None: + return self._run + + while True: + # sleep here to hide an ugly warning + time.sleep(2) + item = self._get_item() + if item and item["associatedRunId"] is not None: + try: + self._run = public.Run( + self.client, + self._entity, + self.project, + item["associatedRunId"], + None, + ) + self._run_id = item["associatedRunId"] + except ValueError as e: + wandb.termwarn(str(e)) + else: + return self._run + elif item: + wandb.termlog("Waiting for run to start") + + time.sleep(3) + + def __repr__(self): + return f" None: + self._name: str = name + self._client = client + self._entity = entity + self._prioritization_mode = prioritization_mode + self._access = _access + self._default_resource_config_id = _default_resource_config_id + self._default_resource_config = _default_resource_config + self._template_variables = None + self._type = None + self._items = None + self._id = None + + @property + def name(self): + return self._name + + @property + def entity(self): + return self._entity + + @property + def prioritization_mode(self) -> RunQueuePrioritizationMode: + if self._prioritization_mode is None: + self._get_metadata() + return self._prioritization_mode + + @property + def access(self) -> RunQueueAccessType: + if self._access is None: + self._get_metadata() + return self._access + + @property + def external_links(self) -> Dict[str, str]: + if self._external_links is None: + self._get_metadata() + return self._external_links + + @property + def type(self) -> RunQueueResourceType: + if self._type is None: + if self._default_resource_config_id is None: + self._get_metadata() + self._get_default_resource_config() + return self._type + + @property + def default_resource_config(self): + if self._default_resource_config is None: + if self._default_resource_config_id is None: + self._get_metadata() + self._get_default_resource_config() + return self._default_resource_config + + @property + def template_variables(self): + if self._template_variables is None: + if self._default_resource_config_id is None: + self._get_metadata() + self._get_default_resource_config() + return self._template_variables + + @property + def id(self) -> str: + if self._id is None: + self._get_metadata() + return self._id + + @property + def items(self) -> List[QueuedRun]: + """Up to the first 100 queued runs. Modifying this list will not modify the queue or any enqueued items!""" + # TODO(np): Add a paginated interface + if self._items is None: + self._get_items() + return self._items + + @normalize_exceptions + def delete(self): + """Delete the run queue from the wandb backend.""" + query = gql( + """ + mutation DeleteRunQueue($id: ID!) { + deleteRunQueues(input: {queueIDs: [$id]}) { + success + clientMutationId + } + } + """ + ) + variable_values = {"id": self.id} + res = self._client.execute(query, variable_values) + if res["deleteRunQueues"]["success"]: + self._id = None + self._access = None + self._default_resource_config_id = None + self._default_resource_config = None + self._items = None + else: + raise CommError(f"Failed to delete run queue {self.name}") + + def __repr__(self): + return f"" + + @normalize_exceptions + def _get_metadata(self): + query = gql( + """ + query GetRunQueueMetadata($projectName: String!, $entityName: String!, $runQueue: String!) { + project(name: $projectName, entityName: $entityName) { + runQueue(name: $runQueue) { + id + access + defaultResourceConfigID + prioritizationMode + externalLinks + } + } + } + """ + ) + variable_values = { + "projectName": LAUNCH_DEFAULT_PROJECT, + "entityName": self._entity, + "runQueue": self._name, + } + res = self._client.execute(query, variable_values) + self._id = res["project"]["runQueue"]["id"] + self._access = res["project"]["runQueue"]["access"] + self._default_resource_config_id = res["project"]["runQueue"][ + "defaultResourceConfigID" + ] + self._external_links = res["project"]["runQueue"]["externalLinks"] + if self._default_resource_config_id is None: + self._default_resource_config = {} + self._prioritization_mode = res["project"]["runQueue"]["prioritizationMode"] + + @normalize_exceptions + def _get_default_resource_config(self): + query = gql( + """ + query GetDefaultResourceConfig($entityName: String!, $id: ID!) { + entity(name: $entityName) { + defaultResourceConfig(id: $id) { + config + resource + templateVariables { + name + schema + } + } + } + } + """ + ) + variable_values = { + "entityName": self._entity, + "id": self._default_resource_config_id, + } + res = self._client.execute(query, variable_values) + self._type = res["entity"]["defaultResourceConfig"]["resource"] + self._default_resource_config = res["entity"]["defaultResourceConfig"]["config"] + self._template_variables = res["entity"]["defaultResourceConfig"][ + "templateVariables" + ] + + @normalize_exceptions + def _get_items(self): + query = gql( + """ + query GetRunQueueItems($projectName: String!, $entityName: String!, $runQueue: String!) { + project(name: $projectName, entityName: $entityName) { + runQueue(name: $runQueue) { + runQueueItems(first: 100) { + edges { + node { + id + } + } + } + } + } + } + """ + ) + variable_values = { + "projectName": LAUNCH_DEFAULT_PROJECT, + "entityName": self._entity, + "runQueue": self._name, + } + res = self._client.execute(query, variable_values) + self._items = [] + for item in res["project"]["runQueue"]["runQueueItems"]["edges"]: + self._items.append( + QueuedRun( + self._client, + self._entity, + LAUNCH_DEFAULT_PROJECT, + self._name, + item["node"]["id"], + ) + ) + + @classmethod + def create( + cls, + name: str, + resource: "RunQueueResourceType", + entity: Optional[str] = None, + prioritization_mode: Optional["RunQueuePrioritizationMode"] = None, + config: Optional[dict] = None, + template_variables: Optional[dict] = None, + ) -> "RunQueue": + public_api = Api() + return public_api.create_run_queue( + name, resource, entity, prioritization_mode, config, template_variables + ) diff --git a/lib/python3.12/site-packages/wandb/apis/public/projects.py b/lib/python3.12/site-packages/wandb/apis/public/projects.py new file mode 100644 index 0000000000000000000000000000000000000000..ba54862ec1c18a3e9ce8b6c8f49ddea14f3f9de7 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/public/projects.py @@ -0,0 +1,225 @@ +"""Public API: projects.""" + +from contextlib import suppress + +from requests import HTTPError +from wandb_gql import gql + +from wandb.apis import public +from wandb.apis.attrs import Attrs +from wandb.apis.normalize import normalize_exceptions +from wandb.apis.paginator import Paginator +from wandb.apis.public.api import RetryingClient +from wandb.sdk.lib import ipython + +PROJECT_FRAGMENT = """fragment ProjectFragment on Project { + id + name + entityName + createdAt + isBenchmark +}""" + + +class Projects(Paginator["Project"]): + """An iterable collection of `Project` objects.""" + + QUERY = gql( + """ + query Projects($entity: String, $cursor: String, $perPage: Int = 50) {{ + models(entityName: $entity, after: $cursor, first: $perPage) {{ + edges {{ + node {{ + ...ProjectFragment + }} + cursor + }} + pageInfo {{ + endCursor + hasNextPage + }} + }} + }} + {} + """.format(PROJECT_FRAGMENT) + ) + + def __init__( + self, + client: RetryingClient, + entity: str, + per_page: int = 50, + ) -> "Projects": + """An iterable collection of `Project` objects. + + Args: + client: The API client used to query W&B. + entity: The entity which owns the projects. + per_page: The number of projects to fetch per request to the API. + """ + self.client = client + self.entity = entity + variables = { + "entity": self.entity, + } + super().__init__(client, variables, per_page) + + @property + def length(self) -> None: + # For backwards compatibility, even though this isn't a SizedPaginator + return None + + @property + def more(self): + if self.last_response: + return self.last_response["models"]["pageInfo"]["hasNextPage"] + else: + return True + + @property + def cursor(self): + if self.last_response: + return self.last_response["models"]["edges"][-1]["cursor"] + else: + return None + + def convert_objects(self): + return [ + Project(self.client, self.entity, p["node"]["name"], p["node"]) + for p in self.last_response["models"]["edges"] + ] + + def __repr__(self): + return f"" + + +class Project(Attrs): + """A project is a namespace for runs.""" + + QUERY = gql( + """ + query Project($project: String!, $entity: String!) { + project(name: $project, entityName: $entity) { + id + } + } + """ + ) + + def __init__( + self, + client: RetryingClient, + entity: str, + project: str, + attrs: dict, + ) -> "Project": + """A single project associated with an entity. + + Args: + client: The API client used to query W&B. + entity: The entity which owns the project. + project: The name of the project to query. + attrs: The attributes of the project. + """ + super().__init__(dict(attrs)) + self.client = client + self.name = project + self.entity = entity + + @property + def path(self): + return [self.entity, self.name] + + @property + def url(self): + return self.client.app_url + "/".join(self.path + ["workspace"]) + + def to_html(self, height=420, hidden=False): + """Generate HTML containing an iframe displaying this project.""" + url = self.url + "?jupyter=true" + style = f"border:none;width:100%;height:{height}px;" + prefix = "" + if hidden: + style += "display:none;" + prefix = ipython.toggle_button("project") + return prefix + f"" + + def _repr_html_(self) -> str: + return self.to_html() + + def __repr__(self): + return "".format("/".join(self.path)) + + @normalize_exceptions + def artifacts_types(self, per_page=50): + return public.ArtifactTypes(self.client, self.entity, self.name) + + @normalize_exceptions + def sweeps(self): + query = gql( + """ + query GetSweeps($project: String!, $entity: String!) {{ + project(name: $project, entityName: $entity) {{ + totalSweeps + sweeps {{ + edges {{ + node {{ + ...SweepFragment + }} + cursor + }} + pageInfo {{ + endCursor + hasNextPage + }} + }} + }} + }} + {} + """.format(public.SWEEP_FRAGMENT) + ) + variable_values = {"project": self.name, "entity": self.entity} + ret = self.client.execute(query, variable_values) + if not ret.get("project") or ret["project"]["totalSweeps"] < 1: + return [] + + return [ + # match format of existing public sweep apis + public.Sweep( + self.client, + self.entity, + self.name, + e["node"]["name"], + ) + for e in ret["project"]["sweeps"]["edges"] + ] + + _PROJECT_ID = gql( + """ + query ProjectID($projectName: String!, $entityName: String!) { + project(name: $projectName, entityName: $entityName) { + id + } + } + """ + ) + + @property + def id(self) -> str: + # This is a workaround to ensure that the project ID can be retrieved + # on demand, as it generally is not set or fetched on instantiation. + # This is necessary if using this project as the scope of a new Automation. + with suppress(LookupError): + return self._attrs["id"] + + variable_values = {"projectName": self.name, "entityName": self.entity} + try: + data = self.client.execute(self._PROJECT_ID, variable_values) + + if not data.get("project") or not data["project"].get("id"): + raise ValueError(f"Project {self.name} not found") + + self._attrs["id"] = data["project"]["id"] + return self._attrs["id"] + except (HTTPError, LookupError, TypeError) as e: + raise ValueError(f"Unable to fetch project ID: {variable_values!r}") from e diff --git a/lib/python3.12/site-packages/wandb/apis/public/query_generator.py b/lib/python3.12/site-packages/wandb/apis/public/query_generator.py new file mode 100644 index 0000000000000000000000000000000000000000..f1260d9127459b852d0dcc874fcc652b8e878452 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/public/query_generator.py @@ -0,0 +1,166 @@ +class QueryGenerator: + """QueryGenerator is a helper object to write filters for runs.""" + + INDIVIDUAL_OP_TO_MONGO = { + "!=": "$ne", + ">": "$gt", + ">=": "$gte", + "<": "$lt", + "<=": "$lte", + "IN": "$in", + "NIN": "$nin", + "REGEX": "$regex", + } + MONGO_TO_INDIVIDUAL_OP = {v: k for k, v in INDIVIDUAL_OP_TO_MONGO.items()} + + GROUP_OP_TO_MONGO = {"AND": "$and", "OR": "$or"} + MONGO_TO_GROUP_OP = {v: k for k, v in GROUP_OP_TO_MONGO.items()} + + def __init__(self): + pass + + @classmethod + def format_order_key(cls, key: str): + if key.startswith(("+", "-")): + direction = key[0] + key = key[1:] + else: + direction = "-" + parts = key.split(".") + if len(parts) == 1: + # Assume the user meant summary_metrics if not a run column + if parts[0] not in ["createdAt", "updatedAt", "name", "sweep"]: + return direction + "summary_metrics." + parts[0] + # Assume summary metrics if prefix isn't known + elif parts[0] not in ["config", "summary_metrics", "tags"]: + return direction + ".".join(["summary_metrics"] + parts) + else: + return direction + ".".join(parts) + + def _is_group(self, op): + return op.get("filters") is not None + + def _is_individual(self, op): + return op.get("key") is not None + + def _to_mongo_op_value(self, op, value): + if op == "=": + return value + else: + return {self.INDIVIDUAL_OP_TO_MONGO[op]: value} + + def key_to_server_path(self, key): + if key["section"] == "config": + return "config." + key["name"] + elif key["section"] == "summary": + return "summary_metrics." + key["name"] + elif key["section"] == "keys_info": + return "keys_info.keys." + key["name"] + elif key["section"] == "run": + return key["name"] + elif key["section"] == "tags": + return "tags." + key["name"] + raise ValueError("Invalid key: {}".format(key)) + + def server_path_to_key(self, path): + if path.startswith("config."): + return {"section": "config", "name": path.split("config.", 1)[1]} + elif path.startswith("summary_metrics."): + return {"section": "summary", "name": path.split("summary_metrics.", 1)[1]} + elif path.startswith("keys_info.keys."): + return {"section": "keys_info", "name": path.split("keys_info.keys.", 1)[1]} + elif path.startswith("tags."): + return {"section": "tags", "name": path.split("tags.", 1)[1]} + else: + return {"section": "run", "name": path} + + def keys_to_order(self, keys): + orders = [] + for key in keys["keys"]: + order = self.key_to_server_path(key["key"]) + if key.get("ascending"): + order = "+" + order + else: + order = "-" + order + orders.append(order) + # return ",".join(orders) + return orders + + def order_to_keys(self, order): + keys = [] + for k in order: # orderstr.split(","): + name = k[1:] + if k[0] == "+": + ascending = True + elif k[0] == "-": + ascending = False + else: + raise Exception("you must sort by ascending(+) or descending(-)") + + key = {"key": {"section": "run", "name": name}, "ascending": ascending} + keys.append(key) + + return {"keys": keys} + + def _to_mongo_individual(self, filter): + if filter["key"]["name"] == "": + return None + + if filter.get("value") is None and filter["op"] != "=" and filter["op"] != "!=": + return None + + if filter.get("disabled") is not None and filter["disabled"]: + return None + + if filter["key"]["section"] == "tags": + if filter["op"] == "IN": + return {"tags": {"$in": filter["value"]}} + if filter["value"] is False: + return { + "$or": [{"tags": None}, {"tags": {"$ne": filter["key"]["name"]}}] + } + else: + return {"tags": filter["key"]["name"]} + path = self.key_to_server_path(filter["key"]) + if path is None: + return path + return {path: self._to_mongo_op_value(filter["op"], filter["value"])} + + def filter_to_mongo(self, filter): + if self._is_individual(filter): + return self._to_mongo_individual(filter) + elif self._is_group(filter): + return { + self.GROUP_OP_TO_MONGO[filter["op"]]: [ + self.filter_to_mongo(f) for f in filter["filters"] + ] + } + + def mongo_to_filter(self, filter): + # Returns {"op": "OR", "filters": [{"op": "AND", "filters": []}]} + if filter is None: + return None # this covers the case where self.filter_to_mongo returns None. + + group_op = None + for key in filter.keys(): + # if self.MONGO_TO_GROUP_OP[key]: + if key in self.MONGO_TO_GROUP_OP: + group_op = key + break + if group_op is not None: + return { + "op": self.MONGO_TO_GROUP_OP[group_op], + "filters": [self.mongo_to_filter(f) for f in filter[group_op]], + } + else: + for k, v in filter.items(): + if isinstance(v, dict): + # TODO: do we always have one key in this case? + op = next(iter(v.keys())) + return { + "key": self.server_path_to_key(k), + "op": self.MONGO_TO_INDIVIDUAL_OP[op], + "value": v[op], + } + else: + return {"key": self.server_path_to_key(k), "op": "=", "value": v} diff --git a/lib/python3.12/site-packages/wandb/apis/public/registries/__init__.py b/lib/python3.12/site-packages/wandb/apis/public/registries/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/wandb/apis/public/registries/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/public/registries/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4beebae6b4738de1fab971698acdd16ca9db20cd Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/public/registries/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/public/registries/__pycache__/_freezable_list.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/public/registries/__pycache__/_freezable_list.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f264c4e472dbf2ff4f28cefd7b049a9e9bd59e72 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/public/registries/__pycache__/_freezable_list.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/public/registries/__pycache__/_utils.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/public/registries/__pycache__/_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c3be4b917f50b5d10f8bab3e7202ec766b9ab80b Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/public/registries/__pycache__/_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/public/registries/__pycache__/registries_search.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/public/registries/__pycache__/registries_search.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..51d4b5799fe6bf48b4140d2741af8ac45f96065a Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/public/registries/__pycache__/registries_search.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/public/registries/__pycache__/registry.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/public/registries/__pycache__/registry.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..babd3b325a0ec391f28160252543eb7b0f01fc20 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/public/registries/__pycache__/registry.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/public/registries/_freezable_list.py b/lib/python3.12/site-packages/wandb/apis/public/registries/_freezable_list.py new file mode 100644 index 0000000000000000000000000000000000000000..c3609fd5ad4e375404f7195f0dcdc565ef0c71a0 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/public/registries/_freezable_list.py @@ -0,0 +1,179 @@ +from itertools import chain +from typing import ( + Any, + Iterable, + Iterator, + List, + MutableSequence, + Sequence, + Tuple, + TypeVar, + Union, + final, + overload, +) + +T = TypeVar("T") + + +@final +class FreezableList(MutableSequence[T]): + """A list-like container type that only allows adding new items. + + It tracks "saved" (immutable) and "draft" (mutable) items. + Items can be added, inserted, and removed while in draft state, but once frozen, + they become immutable. Unlike a set, duplicate items are allowed in the draft + state but duplicates already present in the saved state cannot be added. + Any initial items passed to the constructor are saved. + """ + + def __init__(self, iterable: Union[Iterable[T], None] = None, /) -> None: + self._frozen: Tuple[T, ...] = tuple(iterable or ()) + self._draft: List[T] = [] + + def append(self, value: T) -> None: + """Append an item to the draft list. No duplicates are allowed.""" + if (value in self._frozen) or (value in self._draft): + return + self._draft.append(value) + + def remove(self, value: T) -> None: + """Remove the first occurrence of value from the draft list.""" + if value in self._frozen: + raise ValueError(f"Cannot remove item from frozen list: {value!r}") + self._draft.remove(value) + + def freeze(self) -> None: + """Freeze any draft items by adding them to the saved tuple.""" + # Filter out duplicates already in saved before extending + new_items = tuple(item for item in self._draft if item not in self._frozen) + self._frozen = self._frozen + new_items + self._draft.clear() + + def __eq__(self, value: object) -> bool: + if not isinstance(value, Sequence): + return NotImplemented + return list(self) == list(value) + + def __contains__(self, value: Any) -> bool: + return value in self._frozen or value in self._draft + + def __len__(self) -> int: + return len(self._frozen) + len(self._draft) + + def __iter__(self) -> Iterator[T]: + return iter(chain(self._frozen, self._draft)) + + @overload + def __getitem__(self, index: int) -> T: ... + + @overload + def __getitem__(self, index: slice) -> Sequence[T]: ... + + def __getitem__(self, index: Union[int, slice]) -> Union[T, Sequence[T]]: + return [*self._frozen, *self._draft][index] + + @overload + def __setitem__(self, index: int, value: T) -> None: ... + + @overload + def __setitem__(self, index: slice, value: Iterable[T]) -> None: ... + + def __setitem__( + self, index: Union[int, slice], value: Union[T, Iterable[T]] + ) -> None: + if isinstance(index, slice): + # Setting slices might affect saved items, disallow for simplicity + raise TypeError( + f"{type(self).__name__!r} does not support slice assignment" + ) + else: + if value in self._frozen or value in self._draft: + return + + # The frozen items are sequentially first and protected from changes + len_frozen = len(self._frozen) + size = len(self) + + if (index >= size) or (index < -size): + raise IndexError("Index out of range") + + draft_index = (index % size) - len_frozen + if draft_index < 0: + raise ValueError(f"Cannot assign to saved item at index {index!r}") + self._draft[draft_index] = value + + @overload + def __delitem__(self, index: int) -> None: ... + + @overload + def __delitem__(self, index: slice) -> None: ... + + def __delitem__(self, index: Union[int, slice]) -> None: + if isinstance(index, slice): + raise TypeError(f"{type(self).__name__!r} does not support slice deletion") + else: + # The frozen items are sequentially first and protected from changes + len_frozen = len(self._frozen) + size = len(self) + + if (index >= size) or (index < -size): + raise IndexError("Index out of range") + + draft_index = (index % size) - len_frozen + if draft_index < 0: + raise ValueError(f"Cannot delete saved item at index {index!r}") + del self._draft[draft_index] + + def insert(self, index: int, value: T) -> None: + """Insert item before index. + + Insertion is only allowed at indices corresponding to the draft portion + of the list (i.e., index >= len(frozen_items)). Negative indices are + interpreted relative to the combined length of frozen and draft items. + """ + if value in self._frozen or value in self._draft: + # Silently ignore duplicates, similar to append + return + + # The frozen items are sequentially first and protected from changes + len_frozen = len(self._frozen) + size = len(self) + + # Follow the behavior of `list.insert()` when the index is out of bounds. + # - negative out-of-bounds index: prepend. Will only work if the frozen items are empty. + if index < -size and not self._frozen: + return self._draft.insert(0, value) + + # - positive out-of-bounds index: append. + if index >= size: + return self._draft.append(value) + + # - in-bounds index: insert only if into the draft portion. + draft_index = (index % size) - len_frozen + if draft_index < 0: + raise IndexError( + f"Cannot insert into the frozen list (index < {len_frozen})" + ) + return self._draft.insert(draft_index, value) + + def __repr__(self) -> str: + return f"{type(self).__name__}(frozen={list(self._frozen)!r}, draft={list(self._draft)!r})" + + @property + def draft(self) -> Tuple[T, ...]: + """A read-only, tuple copy of the current draft items.""" + return tuple(self._draft) + + +class AddOnlyArtifactTypesList(FreezableList[str]): + def remove(self, value: str) -> None: + try: + super().remove(value) + except ValueError: + raise ValueError( + f"Cannot remove artifact type: {value!r} that has been saved to the registry" + ) + + def __repr__(self) -> str: + return f"{type(self).__name__}(saved={list(self._frozen)!r}, draft={list(self._draft)!r})" diff --git a/lib/python3.12/site-packages/wandb/apis/public/registries/_utils.py b/lib/python3.12/site-packages/wandb/apis/public/registries/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..2f0a3696823747b380a1e0f9b8379dc8dcdced05 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/public/registries/_utils.py @@ -0,0 +1,140 @@ +from enum import Enum +from functools import lru_cache +from typing import TYPE_CHECKING, Any, List, Literal, Mapping, Optional, Sequence + +from wandb.sdk.artifacts._validators import ( + REGISTRY_PREFIX, + validate_artifact_types_list, +) + +if TYPE_CHECKING: + from wandb_gql import Client + +from wandb_gql import gql + + +class Visibility(str, Enum): + # names are what users see/pass into Python methods + # values are what's expected by backend API + organization = "PRIVATE" + restricted = "RESTRICTED" + + @classmethod + def _missing_(cls, value: object) -> Any: + return next( + (e for e in cls if e.name == value), + None, + ) + + +def format_gql_artifact_types_input( + artifact_types: Optional[List[str]] = None, +): + """Format the artifact types for the GQL input. + + Args: + artifact_types: The artifact types to add to the registry. + + Returns: + The artifact types for the GQL input. + """ + if artifact_types is None: + return [] + new_types = validate_artifact_types_list(artifact_types) + return [{"name": type} for type in new_types] + + +def gql_to_registry_visibility( + visibility: str, +) -> Literal["organization", "restricted"]: + """Convert the GQL visibility to the registry visibility. + + Args: + visibility: The GQL visibility. + + Returns: + The registry visibility. + """ + try: + return Visibility(visibility).name + except ValueError: + raise ValueError(f"Invalid visibility: {visibility!r} from backend") + + +def registry_visibility_to_gql( + visibility: Literal["organization", "restricted"], +) -> str: + """Convert the registry visibility to the GQL visibility.""" + try: + return Visibility[visibility].value + except KeyError: + raise ValueError( + f"Invalid visibility: {visibility!r}. " + f"Must be one of: {', '.join(map(repr, (e.name for e in Visibility)))}" + ) + + +def ensure_registry_prefix_on_names(query, in_name=False): + """Traverse the filter to prepend the `name` key value with the registry prefix unless the value is a regex. + + - in_name: True if we are under a "name" key (or propagating from one). + + EX: {"name": "model"} -> {"name": "wandb-registry-model"} + """ + if isinstance((txt := query), str): + if in_name: + return txt if txt.startswith(REGISTRY_PREFIX) else f"{REGISTRY_PREFIX}{txt}" + return txt + if isinstance((dct := query), Mapping): + new_dict = {} + for key, obj in dct.items(): + if key == "name": + new_dict[key] = ensure_registry_prefix_on_names(obj, in_name=True) + elif key == "$regex": + # For regex operator, we skip transformation of its value. + new_dict[key] = obj + else: + # For any other key, propagate the in_name and skip_transform flags as-is. + new_dict[key] = ensure_registry_prefix_on_names(obj, in_name=in_name) + return new_dict + if isinstance((objs := query), Sequence): + return list( + map(lambda x: ensure_registry_prefix_on_names(x, in_name=in_name), objs) + ) + return query + + +@lru_cache(maxsize=10) +def fetch_org_entity_from_organization(client: "Client", organization: str) -> str: + """Fetch the org entity from the organization. + + Args: + client (Client): Graphql client. + organization (str): The organization to fetch the org entity for. + """ + query = gql( + """ + query FetchOrgEntityFromOrganization($organization: String!) { + organization(name: $organization) { + orgEntity { + name + } + } + } + """ + ) + try: + response = client.execute(query, variable_values={"organization": organization}) + except Exception as e: + raise ValueError( + f"Error fetching org entity for organization: {organization}" + ) from e + + if ( + not (org := response["organization"]) + or not (org_entity := org["orgEntity"]) + or not (org_name := org_entity["name"]) + ): + raise ValueError(f"Organization entity for {organization} not found.") + + return org_name diff --git a/lib/python3.12/site-packages/wandb/apis/public/registries/registries_search.py b/lib/python3.12/site-packages/wandb/apis/public/registries/registries_search.py new file mode 100644 index 0000000000000000000000000000000000000000..a94e4fd9fa927d440a8cfad791d1acc0d0a44bcb --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/public/registries/registries_search.py @@ -0,0 +1,468 @@ +"""Public API: registries search.""" + +import json +from typing import TYPE_CHECKING, Any, Dict, Optional + +if TYPE_CHECKING: + from wandb_gql import Client + +from wandb_gql import gql + +import wandb +from wandb.apis.paginator import Paginator +from wandb.sdk.artifacts._graphql_fragments import ( + _gql_artifact_fragment, + _gql_registry_fragment, +) + +from ._utils import ensure_registry_prefix_on_names + + +class Registries(Paginator): + """Iterator that returns Registries.""" + + QUERY = gql( + """ + query Registries($organization: String!, $filters: JSONString, $cursor: String, $perPage: Int) { + organization(name: $organization) { + orgEntity { + name + projects(filters: $filters, after: $cursor, first: $perPage) { + pageInfo { + endCursor + hasNextPage + } + edges { + node { + ...RegistryFragment + } + } + } + } + } + } + """ + + _gql_registry_fragment() + ) + + def __init__( + self, + client: "Client", + organization: str, + filter: Optional[Dict[str, Any]] = None, + per_page: Optional[int] = 100, + ): + self.client = client + self.organization = organization + self.filter = ensure_registry_prefix_on_names(filter or {}) + variables = { + "organization": organization, + "filters": json.dumps(self.filter), + } + + super().__init__(client, variables, per_page) + + def __bool__(self): + return bool(self.objects) + + def __next__(self): + # Implement custom next since its possible to load empty pages because of auth + self.index += 1 + while len(self.objects) <= self.index: + if not self._load_page(): + raise StopIteration + return self.objects[self.index] + + def collections(self, filter: Optional[Dict[str, Any]] = None) -> "Collections": + return Collections( + self.client, + self.organization, + registry_filter=self.filter, + collection_filter=filter, + ) + + def versions(self, filter: Optional[Dict[str, Any]] = None) -> "Versions": + return Versions( + self.client, + self.organization, + registry_filter=self.filter, + collection_filter=None, + artifact_filter=filter, + ) + + @property + def length(self): + if self.last_response: + return len( + self.last_response["organization"]["orgEntity"]["projects"]["edges"] + ) + else: + return None + + @property + def more(self): + if self.last_response: + return self.last_response["organization"]["orgEntity"]["projects"][ + "pageInfo" + ]["hasNextPage"] + else: + return True + + @property + def cursor(self): + if self.last_response: + return self.last_response["organization"]["orgEntity"]["projects"][ + "pageInfo" + ]["endCursor"] + else: + return None + + def convert_objects(self): + if not self.last_response: + return [] + if ( + not self.last_response["organization"] + or not self.last_response["organization"]["orgEntity"] + ): + raise ValueError( + f"Organization '{self.organization}' not found. Please verify the organization name is correct" + ) + + from wandb.apis.public.registries.registry import Registry + + return [ + Registry( + self.client, + self.organization, + self.last_response["organization"]["orgEntity"]["name"], + r["node"]["name"], + r["node"], + ) + for r in self.last_response["organization"]["orgEntity"]["projects"][ + "edges" + ] + ] + + +class Collections(Paginator): + """Iterator that returns Artifact collections in the Registry.""" + + QUERY = gql( + """ + query Collections( + $organization: String!, + $registryFilter: JSONString, + $collectionFilter: JSONString, + $collectionTypes: [ArtifactCollectionType!], + $cursor: String, + $perPage: Int + ) { + organization(name: $organization) { + orgEntity { + name + artifactCollections( + projectFilters: $registryFilter, + filters: $collectionFilter, + collectionTypes: $collectionTypes, + after: $cursor, + first: $perPage + ) { + totalCount + pageInfo { + endCursor + hasNextPage + } + edges { + cursor + node { + id + name + description + createdAt + tags { + edges { + node { + name + } + } + } + project { + name + entity { + name + } + } + defaultArtifactType { + name + } + aliases { + edges { + node { + alias + } + } + } + } + } + } + } + } + } + """ + ) + + def __init__( + self, + client: "Client", + organization: str, + registry_filter: Optional[Dict[str, Any]] = None, + collection_filter: Optional[Dict[str, Any]] = None, + per_page: Optional[int] = 100, + ): + self.client = client + self.organization = organization + self.registry_filter = registry_filter + self.collection_filter = collection_filter or {} + + variables = { + "registryFilter": ( + json.dumps(self.registry_filter) if self.registry_filter else None + ), + "collectionFilter": ( + json.dumps(self.collection_filter) if self.collection_filter else None + ), + "organization": self.organization, + "collectionTypes": ["PORTFOLIO"], + "perPage": per_page, + } + + super().__init__(client, variables, per_page) + + def __bool__(self): + return len(self) > 0 or len(self.objects) > 0 + + def __next__(self): + # Implement custom next since its possible to load empty pages because of auth + self.index += 1 + while len(self.objects) <= self.index: + if not self._load_page(): + raise StopIteration + return self.objects[self.index] + + def versions(self, filter: Optional[Dict[str, Any]] = None) -> "Versions": + return Versions( + self.client, + self.organization, + registry_filter=self.registry_filter, + collection_filter=self.collection_filter, + artifact_filter=filter, + ) + + @property + def length(self): + if self.last_response: + return self.last_response["organization"]["orgEntity"][ + "artifactCollections" + ]["totalCount"] + else: + return None + + @property + def more(self): + if self.last_response: + return self.last_response["organization"]["orgEntity"][ + "artifactCollections" + ]["pageInfo"]["hasNextPage"] + else: + return True + + @property + def cursor(self): + if self.last_response: + return self.last_response["organization"]["orgEntity"][ + "artifactCollections" + ]["pageInfo"]["endCursor"] + else: + return None + + def convert_objects(self): + from wandb.apis.public import ArtifactCollection + + if not self.last_response: + return [] + if ( + not self.last_response["organization"] + or not self.last_response["organization"]["orgEntity"] + ): + raise ValueError( + f"Organization '{self.organization}' not found. Please verify the organization name is correct" + ) + + return [ + ArtifactCollection( + self.client, + r["node"]["project"]["entity"]["name"], + r["node"]["project"]["name"], + r["node"]["name"], + r["node"]["defaultArtifactType"]["name"], + self.organization, + r["node"], + is_sequence=False, + ) + for r in self.last_response["organization"]["orgEntity"][ + "artifactCollections" + ]["edges"] + ] + + +class Versions(Paginator): + """Iterator that returns Artifact versions in the Registry.""" + + def __init__( + self, + client: "Client", + organization: str, + registry_filter: Optional[Dict[str, Any]] = None, + collection_filter: Optional[Dict[str, Any]] = None, + artifact_filter: Optional[Dict[str, Any]] = None, + per_page: int = 100, + ): + self.client = client + self.organization = organization + self.registry_filter = registry_filter + self.collection_filter = collection_filter + self.artifact_filter = artifact_filter or {} + self.QUERY = gql( + """ + query Versions( + $organization: String!, + $registryFilter: JSONString, + $collectionFilter: JSONString, + $artifactFilter: JSONString, + $cursor: String, + $perPage: Int + ) { + organization(name: $organization) { + orgEntity { + name + artifactMemberships( + projectFilters: $registryFilter, + collectionFilters: $collectionFilter, + filters: $artifactFilter, + after: $cursor, + first: $perPage + ) { + pageInfo { + endCursor + hasNextPage + } + edges { + node { + artifactCollection { + project { + name + entity { + name + } + } + name + } + versionIndex + artifact { + ...ArtifactFragment + } + aliases { + alias + } + } + } + } + } + } + } + """ + + _gql_artifact_fragment(include_aliases=False) + ) + + variables = { + "registryFilter": ( + json.dumps(self.registry_filter) if self.registry_filter else None + ), + "collectionFilter": ( + json.dumps(self.collection_filter) if self.collection_filter else None + ), + "artifactFilter": ( + json.dumps(self.artifact_filter) if self.artifact_filter else None + ), + "organization": self.organization, + } + + super().__init__(client, variables, per_page) + + def __next__(self): + # Implement custom next since its possible to load empty pages because of auth + self.index += 1 + while len(self.objects) <= self.index: + if not self._load_page(): + raise StopIteration + return self.objects[self.index] + + def __bool__(self): + return len(self) > 0 or len(self.objects) > 0 + + @property + def length(self): + if self.last_response: + return len( + self.last_response["organization"]["orgEntity"]["artifactMemberships"][ + "edges" + ] + ) + else: + return None + + @property + def more(self): + if self.last_response: + return self.last_response["organization"]["orgEntity"][ + "artifactMemberships" + ]["pageInfo"]["hasNextPage"] + else: + return True + + @property + def cursor(self): + if self.last_response: + return self.last_response["organization"]["orgEntity"][ + "artifactMemberships" + ]["pageInfo"]["endCursor"] + else: + return None + + def convert_objects(self): + if not self.last_response: + return [] + if ( + not self.last_response["organization"] + or not self.last_response["organization"]["orgEntity"] + ): + raise ValueError( + f"Organization '{self.organization}' not found. Please verify the organization name is correct" + ) + + artifacts = ( + wandb.Artifact._from_attrs( + a["node"]["artifactCollection"]["project"]["entity"]["name"], + a["node"]["artifactCollection"]["project"]["name"], + a["node"]["artifactCollection"]["name"] + + ":v" + + str(a["node"]["versionIndex"]), + a["node"]["artifact"], + self.client, + [alias["alias"] for alias in a["node"]["aliases"]], + ) + for a in self.last_response["organization"]["orgEntity"][ + "artifactMemberships" + ]["edges"] + ) + return artifacts diff --git a/lib/python3.12/site-packages/wandb/apis/public/registries/registry.py b/lib/python3.12/site-packages/wandb/apis/public/registries/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..faf80dafd4eafcd2afdfa246b1185c763ff51ccb --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/public/registries/registry.py @@ -0,0 +1,358 @@ +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional + +from wandb_gql import gql + +import wandb +from wandb.proto.wandb_internal_pb2 import ServerFeature +from wandb.sdk.artifacts._validators import REGISTRY_PREFIX, validate_project_name +from wandb.sdk.internal.internal_api import Api as InternalApi +from wandb.sdk.projects._generated import ( + DELETE_PROJECT_GQL, + FETCH_REGISTRY_GQL, + RENAME_PROJECT_GQL, + UPSERT_REGISTRY_PROJECT_GQL, + DeleteProject, + RenameProject, + UpsertRegistryProject, +) + +from ._freezable_list import AddOnlyArtifactTypesList +from ._utils import ( + fetch_org_entity_from_organization, + format_gql_artifact_types_input, + gql_to_registry_visibility, + registry_visibility_to_gql, +) +from .registries_search import Collections, Versions + +if TYPE_CHECKING: + from wandb_gql import Client + + +class Registry: + """A single registry in the Registry.""" + + def __init__( + self, + client: "Client", + organization: str, + entity: str, + name: str, + attrs: Optional[Dict[str, Any]] = None, + ): + self.client = client + self._name = name + self._saved_name = name + self._entity = entity + self._organization = organization + if attrs is not None: + self._update_attributes(attrs) + + def _update_attributes(self, attrs: Dict[str, Any]) -> None: + """Helper method to update instance attributes from a dictionary.""" + self._id = attrs.get("id", "") + if self._id is None: + raise ValueError(f"Registry {self.name}'s id is not found") + + self._description = attrs.get("description", "") + self._allow_all_artifact_types = attrs.get( + "allowAllArtifactTypesInRegistry", False + ) + self._artifact_types = AddOnlyArtifactTypesList( + t["node"]["name"] for t in attrs.get("artifactTypes", {}).get("edges", []) + ) + self._created_at = attrs.get("createdAt", "") + self._updated_at = attrs.get("updatedAt", "") + self._visibility = gql_to_registry_visibility(attrs.get("access", "")) + + @property + def full_name(self) -> str: + """Full name of the registry including the `wandb-registry-` prefix.""" + return f"wandb-registry-{self.name}" + + @property + def name(self) -> str: + """Name of the registry without the `wandb-registry-` prefix.""" + return self._name + + @name.setter + def name(self, value: str): + self._name = value + + @property + def entity(self) -> str: + """Organization entity of the registry.""" + return self._entity + + @property + def organization(self) -> str: + """Organization name of the registry.""" + return self._organization + + @property + def description(self) -> str: + """Description of the registry.""" + return self._description + + @description.setter + def description(self, value: str): + """Set the description of the registry.""" + self._description = value + + @property + def allow_all_artifact_types(self): + """Returns whether all artifact types are allowed in the registry. + + If `True` then artifacts of any type can be added to this registry. + If `False` then artifacts are restricted to the types in `artifact_types` for this registry. + """ + return self._allow_all_artifact_types + + @allow_all_artifact_types.setter + def allow_all_artifact_types(self, value: bool): + """Set whether all artifact types are allowed in the registry.""" + self._allow_all_artifact_types = value + + @property + def artifact_types(self) -> AddOnlyArtifactTypesList: + """Returns the artifact types allowed in the registry. + + If `allow_all_artifact_types` is `True` then `artifact_types` reflects the + types previously saved or currently used in the registry. + If `allow_all_artifact_types` is `False` then artifacts are restricted to the + types in `artifact_types`. + + Note: + Previously saved artifact types cannot be removed. + + Example: + ```python + registry.artifact_types.append("model") + registry.save() # once saved, the artifact type `model` cannot be removed + registry.artifact_types.append("accidentally_added") + registry.artifact_types.remove( + "accidentally_added" + ) # Types can only be removed if it has not been saved yet + ``` + """ + return self._artifact_types + + @property + def created_at(self) -> str: + """Timestamp of when the registry was created.""" + return self._created_at + + @property + def updated_at(self) -> str: + """Timestamp of when the registry was last updated.""" + return self._updated_at + + @property + def path(self): + return [self.entity, self.full_name] + + @property + def visibility(self) -> Literal["organization", "restricted"]: + """Visibility of the registry. + + Returns: + Literal["organization", "restricted"]: The visibility level. + - "organization": Anyone in the organization can view this registry. + You can edit their roles later from the settings in the UI. + - "restricted": Only invited members via the UI can access this registry. + Public sharing is disabled. + """ + return self._visibility + + @visibility.setter + def visibility(self, value: Literal["organization", "restricted"]): + """Set the visibility of the registry. + + Args: + value: The visibility level. Options are: + - "organization": Anyone in the organization can view this registry. + You can edit their roles later from the settings in the UI. + - "restricted": Only invited members via the UI can access this registry. + Public sharing is disabled. + """ + self._visibility = value + + def collections(self, filter: Optional[Dict[str, Any]] = None) -> Collections: + """Returns the collections belonging to the registry.""" + registry_filter = { + "name": self.full_name, + } + return Collections(self.client, self.organization, registry_filter, filter) + + def versions(self, filter: Optional[Dict[str, Any]] = None) -> Versions: + """Returns the versions belonging to the registry.""" + registry_filter = { + "name": self.full_name, + } + return Versions(self.client, self.organization, registry_filter, None, filter) + + @classmethod + def create( + cls, + client: "Client", + organization: str, + name: str, + visibility: Literal["organization", "restricted"], + description: Optional[str] = None, + artifact_types: Optional[List[str]] = None, + ): + """Create a new registry. + + The registry name must be unique within the organization. + This function should be called using `api.create_registry()` + + Args: + client: The GraphQL client. + organization: The name of the organization. + name: The name of the registry (without the `wandb-registry-` prefix). + visibility: The visibility level ('organization' or 'restricted'). + description: An optional description for the registry. + artifact_types: An optional list of allowed artifact types. + + Returns: + Registry: The newly created Registry object. + + Raises: + ValueError: If a registry with the same name already exists in the + organization or if the creation fails. + """ + org_entity = fetch_org_entity_from_organization(client, organization) + full_name = REGISTRY_PREFIX + name + validate_project_name(full_name) + accepted_artifact_types = [] + if artifact_types: + accepted_artifact_types = format_gql_artifact_types_input(artifact_types) + visibility_value = registry_visibility_to_gql(visibility) + registry_creation_error = ( + f"Failed to create registry {name!r} in organization {organization!r}." + ) + try: + response = client.execute( + gql(UPSERT_REGISTRY_PROJECT_GQL), + variable_values={ + "description": description, + "entityName": org_entity, + "name": full_name, + "access": visibility_value, + "allowAllArtifactTypesInRegistry": not accepted_artifact_types, + "artifactTypes": accepted_artifact_types, + }, + ) + except Exception: + raise ValueError(registry_creation_error) + if not response["upsertModel"]["inserted"]: + raise ValueError(registry_creation_error) + + return Registry( + client, + organization, + org_entity, + name, + response["upsertModel"]["project"], + ) + + def delete(self) -> None: + """Delete the registry. This is irreversible.""" + try: + response = self.client.execute( + gql(DELETE_PROJECT_GQL), variable_values={"id": self._id} + ) + result = DeleteProject.model_validate(response) + except Exception: + raise ValueError( + f"Failed to delete registry: {self.name!r} in organization: {self.organization!r}" + ) + if not result.delete_model.success: + raise ValueError( + f"Failed to delete registry: {self.name!r} in organization: {self.organization!r}" + ) + + def load(self) -> None: + """Load the registry attributes from the backend to reflect the latest saved state.""" + load_failure_message = ( + f"Failed to load registry {self.name!r} " + f"in organization {self.organization!r}." + ) + try: + response = self.client.execute( + gql(FETCH_REGISTRY_GQL), + variable_values={ + "name": self.full_name, + "entityName": self.entity, + }, + ) + except Exception: + raise ValueError(load_failure_message) + if response["entity"] is None: + raise ValueError(load_failure_message) + self.attrs = response["entity"]["project"] + if self.attrs is None: + raise ValueError(load_failure_message) + self._update_attributes(self.attrs) + + def save(self) -> None: + """Save registry attributes to the backend.""" + if not InternalApi()._server_supports( + ServerFeature.INCLUDE_ARTIFACT_TYPES_IN_REGISTRY_CREATION + ): + raise RuntimeError( + "saving the registry is not enabled on this wandb server version. " + "Please upgrade your server version or contact support at support@wandb.com." + ) + + if self._no_updating_registry_types(): + raise ValueError( + f"Cannot update artifact types when `allows_all_artifact_types` is {True!r}. Set it to {False!r} first." + ) + + validate_project_name(self.full_name) + visibility_value = registry_visibility_to_gql(self.visibility) + newly_added_types = format_gql_artifact_types_input(self.artifact_types.draft) + registry_save_error = f"Failed to save and update registry: {self.name} in organization: {self.organization}" + full_saved_name = f"{REGISTRY_PREFIX}{self._saved_name}" + try: + response = self.client.execute( + gql(UPSERT_REGISTRY_PROJECT_GQL), + variable_values={ + "description": self.description, + "entityName": self.entity, + "name": full_saved_name, # this makes it so we are updating the original registry in case the name has changed + "access": visibility_value, + "allowAllArtifactTypesInRegistry": self.allow_all_artifact_types, + "artifactTypes": newly_added_types, + }, + ) + result = UpsertRegistryProject.model_validate(response) + except Exception: + raise ValueError(registry_save_error) + if result.upsert_model.inserted: + # This is not suppose trigger unless the user has messed with the `_saved_name` variable + wandb.termlog( + f"Created registry {self.name!r} in organization {self.organization!r} on save" + ) + self._update_attributes(response["upsertModel"]["project"]) + + # Update the name of the registry if it has changed + if self._saved_name != self.name: + response = self.client.execute( + gql(RENAME_PROJECT_GQL), + variable_values={ + "entityName": self.entity, + "oldProjectName": full_saved_name, + "newProjectName": self.full_name, + }, + ) + result = RenameProject.model_validate(response) + self._saved_name = self.name + if result.rename_project.inserted: + # This is not suppose trigger unless the user has messed with the `_saved_name` variable + wandb.termlog(f"Created new registry {self.name!r} on save") + + def _no_updating_registry_types(self) -> bool: + # artifact types draft means user assigned types to add that are not yet saved + return len(self.artifact_types.draft) > 0 and self.allow_all_artifact_types diff --git a/lib/python3.12/site-packages/wandb/apis/public/reports.py b/lib/python3.12/site-packages/wandb/apis/public/reports.py new file mode 100644 index 0000000000000000000000000000000000000000..b91b63466732f9dc228a772f46eba010f4b72bb5 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/public/reports.py @@ -0,0 +1,514 @@ +"""Public API: reports.""" + +import ast +import json +import re +import urllib + +from wandb_gql import gql + +import wandb +from wandb.apis import public +from wandb.apis.attrs import Attrs +from wandb.apis.paginator import SizedPaginator +from wandb.sdk.lib import ipython + + +class Reports(SizedPaginator["BetaReport"]): + """Reports is an iterable collection of `BetaReport` objects.""" + + QUERY = gql( + """ + query ProjectViews($project: String!, $entity: String!, $reportCursor: String, + $reportLimit: Int!, $viewType: String = "runs", $viewName: String) { + project(name: $project, entityName: $entity) { + allViews(viewType: $viewType, viewName: $viewName, first: + $reportLimit, after: $reportCursor) { + edges { + node { + id + name + displayName + description + user { + username + photoUrl + email + } + spec + updatedAt + createdAt + } + cursor + } + pageInfo { + endCursor + hasNextPage + } + + } + } + } + """ + ) + + def __init__(self, client, project, name=None, entity=None, per_page=50): + self.project = project + self.name = name + variables = { + "project": project.name, + "entity": project.entity, + "viewName": self.name, + } + super().__init__(client, variables, per_page) + + @property + def _length(self): + # TODO: Add the count the backend + if self.last_response: + return len(self.objects) + else: + return None + + @property + def more(self): + if self.last_response: + return self.last_response["project"]["allViews"]["pageInfo"]["hasNextPage"] + else: + return True + + @property + def cursor(self): + if self.last_response: + return self.last_response["project"]["allViews"]["edges"][-1]["cursor"] + else: + return None + + def update_variables(self): + self.variables.update( + {"reportCursor": self.cursor, "reportLimit": self.per_page} + ) + + def convert_objects(self): + if self.last_response["project"] is None: + raise ValueError( + f"Project {self.variables['project']} does not exist under entity {self.variables['entity']}" + ) + return [ + BetaReport( + self.client, + r["node"], + entity=self.project.entity, + project=self.project.name, + ) + for r in self.last_response["project"]["allViews"]["edges"] + ] + + def __repr__(self): + return "".format("/".join(self.project.path)) + + +class BetaReport(Attrs): + """BetaReport is a class associated with reports created in wandb. + + WARNING: this API will likely change in a future release + + Attributes: + id (string): unique identifier of the report + name (string): report name + display_name (string): display name of the report + description (string): report description + user (User): the user that created the report (contains username and email) + spec (dict): the spec of the report + url (string): the url of the report + updated_at (string): timestamp of last update + created_at (string): timestamp when the report was created + """ + + def __init__(self, client, attrs, entity=None, project=None): + self.client = client + self.project = project + self.entity = entity + self.query_generator = public.QueryGenerator() + super().__init__(dict(attrs)) + + if "spec" in self._attrs: + if isinstance(self._attrs["spec"], str): + self._attrs["spec"] = json.loads(self._attrs["spec"]) + else: + self._attrs["spec"] = {} + + @property + def spec(self): + return self._attrs["spec"] + + @property + def sections(self): + return self.spec["panelGroups"] + + def runs(self, section, per_page=50, only_selected=True): + run_set_idx = section.get("openRunSet", 0) + run_set = section["runSets"][run_set_idx] + order = self.query_generator.key_to_server_path(run_set["sort"]["key"]) + if run_set["sort"].get("ascending"): + order = "+" + order + else: + order = "-" + order + filters = self.query_generator.filter_to_mongo(run_set["filters"]) + if only_selected: + # TODO: handle this not always existing + filters["$or"][0]["$and"].append( + {"name": {"$in": run_set["selections"]["tree"]}} + ) + return public.Runs( + self.client, + self.entity, + self.project, + filters=filters, + order=order, + per_page=per_page, + ) + + @property + def id(self): + return self._attrs.get("id") + + @property + def name(self): + return self._attrs.get("name") + + @property + def display_name(self): + return self._attrs.get("displayName") + + @property + def description(self): + return self._attrs.get("description") + + @property + def user(self): + return self._attrs.get("user") + + @property + def updated_at(self): + return self._attrs.get("updatedAt") + + @property + def created_at(self): + return self._attrs.get("createdAt") + + @property + def url(self): + if ( + not self.client + or not self.entity + or not self.project + or not self.display_name + or not self.id + ): + return None + return self.client.app_url + "/".join( + [ + self.entity, + self.project, + "reports", + "--".join( + [ + # made this more closely match the url creation in the frontend (https://github.com/wandb/core/blob/76943979c8e967f7a62dae8bef0a001a2672584c/frontends/app/src/util/report/urls.ts#L19) + urllib.parse.quote( + re.sub( + r"-+", "-", re.sub(r"\W", "-", self.display_name) + ).strip("-") + ), + self.id.replace("=", ""), + ] + ), + ] + ) + + def to_html(self, height=1024, hidden=False): + """Generate HTML containing an iframe displaying this report.""" + url = self.url + if url is None: + return "
Report URL not available
" + url = url + "?jupyter=true" + style = f"border:none;width:100%;height:{height}px;" + prefix = "" + if hidden: + style += "display:none;" + prefix = ipython.toggle_button("report") + return prefix + f"" + + def _repr_html_(self) -> str: + return self.to_html() + + +class PythonMongoishQueryGenerator: + SPACER = "----------" + DECIMAL_SPACER = ";;;" + FRONTEND_NAME_MAPPING = { + "ID": "name", + "Name": "displayName", + "Tags": "tags", + "State": "state", + "CreatedTimestamp": "createdAt", + "Runtime": "duration", + "User": "username", + "Sweep": "sweep", + "Group": "group", + "JobType": "jobType", + "Hostname": "host", + "UsingArtifact": "inputArtifacts", + "OutputtingArtifact": "outputArtifacts", + "Step": "_step", + "Relative Time (Wall)": "_absolute_runtime", + "Relative Time (Process)": "_runtime", + "Wall Time": "_timestamp", + # "GroupedRuns": "__wb_group_by_all" + } + FRONTEND_NAME_MAPPING_REVERSED = {v: k for k, v in FRONTEND_NAME_MAPPING.items()} + AST_OPERATORS = { + ast.Lt: "$lt", + ast.LtE: "$lte", + ast.Gt: "$gt", + ast.GtE: "$gte", + ast.Eq: "=", + ast.Is: "=", + ast.NotEq: "$ne", + ast.IsNot: "$ne", + ast.In: "$in", + ast.NotIn: "$nin", + ast.And: "$and", + ast.Or: "$or", + ast.Not: "$not", + } + + AST_FIELDS = { + ast.Constant: "value", + ast.Name: "id", + ast.List: "elts", + ast.Tuple: "elts", + } + + def __init__(self, run_set): + self.run_set = run_set + self.panel_metrics_helper = PanelMetricsHelper() + + def _handle_compare(self, node): + # only left side can be a col + left = self.front_to_back(self._handle_fields(node.left)) + op = self._handle_ops(node.ops[0]) + right = self._handle_fields(node.comparators[0]) + + # Eq has no op for some reason + if op == "=": + return {left: right} + else: + return {left: {op: right}} + + def _handle_fields(self, node): + result = getattr(node, self.AST_FIELDS.get(type(node))) + if isinstance(result, list): + return [self._handle_fields(node) for node in result] + elif isinstance(result, str): + return self._unconvert(result) + return result + + def _handle_ops(self, node): + return self.AST_OPERATORS.get(type(node)) + + def _replace_numeric_dots(self, s): + numeric_dots = [] + for i, (left, mid, right) in enumerate(zip(s, s[1:], s[2:]), 1): + if mid == ".": + if ( + left.isdigit() + and right.isdigit() # 1.2 + or left.isdigit() + and right == " " # 1. + or left == " " + and right.isdigit() # .2 + ): + numeric_dots.append(i) + # Edge: Catch number ending in dot at end of string + if s[-2].isdigit() and s[-1] == ".": + numeric_dots.append(len(s) - 1) + numeric_dots = [-1] + numeric_dots + [len(s)] + + substrs = [] + for start, stop in zip(numeric_dots, numeric_dots[1:]): + substrs.append(s[start + 1 : stop]) + substrs.append(self.DECIMAL_SPACER) + substrs = substrs[:-1] + return "".join(substrs) + + def _convert(self, filterstr): + _conversion = ( + self._replace_numeric_dots(filterstr) # temporarily sub numeric dots + .replace(".", self.SPACER) # Allow dotted fields + .replace(self.DECIMAL_SPACER, ".") # add them back + ) + return "(" + _conversion + ")" + + def _unconvert(self, field_name): + return field_name.replace(self.SPACER, ".") # Allow dotted fields + + def python_to_mongo(self, filterstr): + try: + tree = ast.parse(self._convert(filterstr), mode="eval") + except SyntaxError as e: + raise ValueError( + "Invalid python comparison expression; form something like `my_col == 123`" + ) from e + + multiple_filters = hasattr(tree.body, "op") + + if multiple_filters: + op = self.AST_OPERATORS.get(type(tree.body.op)) + values = [self._handle_compare(v) for v in tree.body.values] + else: + op = "$and" + values = [self._handle_compare(tree.body)] + return {"$or": [{op: values}]} + + def front_to_back(self, name): + name, *rest = name.split(".") + rest = "." + ".".join(rest) if rest else "" + + if name in self.FRONTEND_NAME_MAPPING: + return self.FRONTEND_NAME_MAPPING[name] + elif name in self.FRONTEND_NAME_MAPPING_REVERSED: + return name + elif name in self.run_set._runs_config: + return f"config.{name}.value{rest}" + else: # assume summary metrics + return f"summary_metrics.{name}{rest}" + + def back_to_front(self, name): + if name in self.FRONTEND_NAME_MAPPING_REVERSED: + return self.FRONTEND_NAME_MAPPING_REVERSED[name] + elif name in self.FRONTEND_NAME_MAPPING: + return name + elif ( + name.startswith("config.") and ".value" in name + ): # may be brittle: originally "endswith", but that doesn't work with nested keys... + # strip is weird sometimes (??) + return name.replace("config.", "").replace(".value", "") + elif name.startswith("summary_metrics."): + return name.replace("summary_metrics.", "") + wandb.termerror(f"Unknown token: {name}") + return name + + # These are only used for ParallelCoordinatesPlot because it has weird backend names... + def pc_front_to_back(self, name): + name, *rest = name.split(".") + rest = "." + ".".join(rest) if rest else "" + if name is None: + return None + elif name in self.panel_metrics_helper.FRONTEND_NAME_MAPPING: + return "summary:" + self.panel_metrics_helper.FRONTEND_NAME_MAPPING[name] + elif name in self.FRONTEND_NAME_MAPPING: + return self.FRONTEND_NAME_MAPPING[name] + elif name in self.FRONTEND_NAME_MAPPING_REVERSED: + return name + elif name in self.run_set._runs_config: + return f"config:{name}.value{rest}" + else: # assume summary metrics + return f"summary:{name}{rest}" + + def pc_back_to_front(self, name): + if name is None: + return None + elif "summary:" in name: + name = name.replace("summary:", "") + return self.panel_metrics_helper.FRONTEND_NAME_MAPPING_REVERSED.get( + name, name + ) + elif name in self.FRONTEND_NAME_MAPPING_REVERSED: + return self.FRONTEND_NAME_MAPPING_REVERSED[name] + elif name in self.FRONTEND_NAME_MAPPING: + return name + elif name.startswith("config:") and ".value" in name: + return name.replace("config:", "").replace(".value", "") + elif name.startswith("summary_metrics."): + return name.replace("summary_metrics.", "") + return name + + +class PanelMetricsHelper: + FRONTEND_NAME_MAPPING = { + "Step": "_step", + "Relative Time (Wall)": "_absolute_runtime", + "Relative Time (Process)": "_runtime", + "Wall Time": "_timestamp", + } + FRONTEND_NAME_MAPPING_REVERSED = {v: k for k, v in FRONTEND_NAME_MAPPING.items()} + + RUN_MAPPING = {"Created Timestamp": "createdAt", "Latest Timestamp": "heartbeatAt"} + RUN_MAPPING_REVERSED = {v: k for k, v in RUN_MAPPING.items()} + + def front_to_back(self, name): + if name in self.FRONTEND_NAME_MAPPING: + return self.FRONTEND_NAME_MAPPING[name] + return name + + def back_to_front(self, name): + if name in self.FRONTEND_NAME_MAPPING_REVERSED: + return self.FRONTEND_NAME_MAPPING_REVERSED[name] + return name + + # ScatterPlot and ParallelCoords have weird conventions + def special_front_to_back(self, name): + if name is None: + return name + + name, *rest = name.split(".") + rest = "." + ".".join(rest) if rest else "" + + # special case for config + if name.startswith("c::"): + name = name[3:] + return f"config:{name}.value{rest}" + + # special case for summary + if name.startswith("s::"): + name = name[3:] + rest + return f"summary:{name}" + + name = name + rest + if name in self.RUN_MAPPING: + return "run:" + self.RUN_MAPPING[name] + if name in self.FRONTEND_NAME_MAPPING: + return "summary:" + self.FRONTEND_NAME_MAPPING[name] + if name == "Index": + return name + return "summary:" + name + + def special_back_to_front(self, name): + if name is not None: + kind, rest = name.split(":", 1) + + if kind == "config": + pieces = rest.split(".") + if len(pieces) <= 1: + raise ValueError(f"Invalid name: {name}") + elif len(pieces) == 2: + name = pieces[0] + elif len(pieces) >= 3: + name = pieces[:1] + pieces[2:] + name = ".".join(name) + return f"c::{name}" + + elif kind == "summary": + name = rest + return f"s::{name}" + + if name is None: + return name + elif "summary:" in name: + name = name.replace("summary:", "") + return self.FRONTEND_NAME_MAPPING_REVERSED.get(name, name) + elif "run:" in name: + name = name.replace("run:", "") + return self.RUN_MAPPING_REVERSED[name] + return name diff --git a/lib/python3.12/site-packages/wandb/apis/public/runs.py b/lib/python3.12/site-packages/wandb/apis/public/runs.py new file mode 100644 index 0000000000000000000000000000000000000000..bfe3ec05a44f3486932f3526c36fecae3a98850f --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/public/runs.py @@ -0,0 +1,1022 @@ +"""Public API: runs.""" + +import json +import os +import tempfile +import time +import urllib +from typing import ( + TYPE_CHECKING, + Any, + Collection, + Dict, + List, + Literal, + Mapping, + Optional, +) + +from wandb_gql import gql + +import wandb +from wandb import env, util +from wandb.apis import public +from wandb.apis.attrs import Attrs +from wandb.apis.internal import Api as InternalApi +from wandb.apis.normalize import normalize_exceptions +from wandb.apis.paginator import SizedPaginator +from wandb.apis.public.const import RETRY_TIMEDELTA +from wandb.sdk.lib import ipython, json_util, runid +from wandb.sdk.lib.paths import LogicalPath + +if TYPE_CHECKING: + from wandb.apis.public import RetryingClient + +WANDB_INTERNAL_KEYS = {"_wandb", "wandb_version"} + +RUN_FRAGMENT = """fragment RunFragment on Run { + id + tags + name + displayName + sweepName + state + config + group + jobType + commit + readOnly + createdAt + heartbeatAt + description + notes + systemMetrics + summaryMetrics + historyLineCount + user { + name + username + } + historyKeys +}""" + + +@normalize_exceptions +def _server_provides_internal_id_for_project(client) -> bool: + """Returns True if the server allows us to query the internalId field for a project. + + This check is done by utilizing GraphQL introspection in the available fields on the Project type. + """ + query_string = """ + query ProbeRunInput { + RunType: __type(name:"Run") { + fields { + name + } + } + } + """ + + # Only perform the query once to avoid extra network calls + query = gql(query_string) + res = client.execute(query) + return "projectId" in [ + x["name"] for x in (res.get("RunType", {}).get("fields", [{}])) + ] + + +class Runs(SizedPaginator["Run"]): + """An iterable collection of runs associated with a project and optional filter. + + This is generally used indirectly via the `Api`.runs method. + """ + + def __init__( + self, + client: "RetryingClient", + entity: str, + project: str, + filters: Optional[Dict[str, Any]] = None, + order: Optional[str] = None, + per_page: int = 50, + include_sweeps: bool = True, + ): + self.QUERY = gql( + f"""#graphql + query Runs($project: String!, $entity: String!, $cursor: String, $perPage: Int = 50, $order: String, $filters: JSONString) {{ + project(name: $project, entityName: $entity) {{ + internalId + runCount(filters: $filters) + readOnly + runs(filters: $filters, after: $cursor, first: $perPage, order: $order) {{ + edges {{ + node {{ + {"projectId" if _server_provides_internal_id_for_project(client) else ""} + ...RunFragment + }} + cursor + }} + pageInfo {{ + endCursor + hasNextPage + }} + }} + }} + }} + {RUN_FRAGMENT} + """ + ) + + self.entity = entity + self.project = project + self._project_internal_id = None + self.filters = filters or {} + self.order = order + self._sweeps = {} + self._include_sweeps = include_sweeps + variables = { + "project": self.project, + "entity": self.entity, + "order": self.order, + "filters": json.dumps(self.filters), + } + super().__init__(client, variables, per_page) + + @property + def _length(self): + if not self.last_response: + self._load_page() + return self.last_response["project"]["runCount"] + + @property + def more(self): + if self.last_response: + return self.last_response["project"]["runs"]["pageInfo"]["hasNextPage"] + else: + return True + + @property + def cursor(self): + if self.last_response: + return self.last_response["project"]["runs"]["edges"][-1]["cursor"] + else: + return None + + def convert_objects(self): + objs = [] + if self.last_response is None or self.last_response.get("project") is None: + raise ValueError("Could not find project {}".format(self.project)) + for run_response in self.last_response["project"]["runs"]["edges"]: + run = Run( + self.client, + self.entity, + self.project, + run_response["node"]["name"], + run_response["node"], + include_sweeps=self._include_sweeps, + ) + objs.append(run) + + if self._include_sweeps and run.sweep_name: + if run.sweep_name in self._sweeps: + sweep = self._sweeps[run.sweep_name] + else: + sweep = public.Sweep.get( + self.client, + self.entity, + self.project, + run.sweep_name, + withRuns=False, + ) + self._sweeps[run.sweep_name] = sweep + + if sweep is None: + continue + run.sweep = sweep + + return objs + + @normalize_exceptions + def histories( + self, + samples: int = 500, + keys: Optional[List[str]] = None, + x_axis: str = "_step", + format: Literal["default", "pandas", "polars"] = "default", + stream: Literal["default", "system"] = "default", + ): + """Return sampled history metrics for all runs that fit the filters conditions. + + Args: + samples : (int, optional) The number of samples to return per run + keys : (list[str], optional) Only return metrics for specific keys + x_axis : (str, optional) Use this metric as the xAxis defaults to _step + format : (Literal, optional) Format to return data in, options are "default", "pandas", "polars" + stream : (Literal, optional) "default" for metrics, "system" for machine metrics + Returns: + pandas.DataFrame: If format="pandas", returns a `pandas.DataFrame` of history metrics. + polars.DataFrame: If format="polars", returns a `polars.DataFrame` of history metrics. + list of dicts: If format="default", returns a list of dicts containing history metrics with a run_id key. + """ + if format not in ("default", "pandas", "polars"): + raise ValueError( + f"Invalid format: {format}. Must be one of 'default', 'pandas', 'polars'" + ) + + histories = [] + + if format == "default": + for run in self: + history_data = run.history( + samples=samples, + keys=keys, + x_axis=x_axis, + pandas=False, + stream=stream, + ) + if not history_data: + continue + for entry in history_data: + entry["run_id"] = run.id + histories.extend(history_data) + + return histories + + if format == "pandas": + pd = util.get_module( + "pandas", required="Exporting pandas DataFrame requires pandas" + ) + for run in self: + history_data = run.history( + samples=samples, + keys=keys, + x_axis=x_axis, + pandas=False, + stream=stream, + ) + if not history_data: + continue + df = pd.DataFrame.from_records(history_data) + df["run_id"] = run.id + histories.append(df) + if not histories: + return pd.DataFrame() + combined_df = pd.concat(histories) + combined_df.reset_index(drop=True, inplace=True) + # sort columns for consistency + combined_df = combined_df[(sorted(combined_df.columns))] + + return combined_df + + if format == "polars": + pl = util.get_module( + "polars", required="Exporting polars DataFrame requires polars" + ) + for run in self: + history_data = run.history( + samples=samples, + keys=keys, + x_axis=x_axis, + pandas=False, + stream=stream, + ) + if not history_data: + continue + df = pl.from_records(history_data) + df = df.with_columns(pl.lit(run.id).alias("run_id")) + histories.append(df) + if not histories: + return pl.DataFrame() + combined_df = pl.concat(histories, how="vertical") + # sort columns for consistency + combined_df = combined_df.select(sorted(combined_df.columns)) + + return combined_df + + def __repr__(self): + return f"" + + +class Run(Attrs): + """A single run associated with an entity and project. + + Attributes: + tags ([str]): a list of tags associated with the run + url (str): the url of this run + id (str): unique identifier for the run (defaults to eight characters) + name (str): the name of the run + state (str): one of: running, finished, crashed, killed, preempting, preempted + config (dict): a dict of hyperparameters associated with the run + created_at (str): ISO timestamp when the run was started + system_metrics (dict): the latest system metrics recorded for the run + summary (dict): A mutable dict-like property that holds the current summary. + Calling update will persist any changes. + project (str): the project associated with the run + entity (str): the name of the entity associated with the run + project_internal_id (int): the internal id of the project + user (str): the name of the user who created the run + path (str): Unique identifier [entity]/[project]/[run_id] + notes (str): Notes about the run + read_only (boolean): Whether the run is editable + history_keys (str): Keys of the history metrics that have been logged + with `wandb.log({key: value})` + metadata (str): Metadata about the run from wandb-metadata.json + """ + + def __init__( + self, + client: "RetryingClient", + entity: str, + project: str, + run_id: str, + attrs: Optional[Mapping] = None, + include_sweeps: bool = True, + ): + """Initialize a Run object. + + Run is always initialized by calling api.runs() where api is an instance of + wandb.Api. + """ + _attrs = attrs or {} + super().__init__(dict(_attrs)) + self.client = client + self._entity = entity + self.project = project + self._files = {} + self._base_dir = env.get_dir(tempfile.gettempdir()) + self.id = run_id + self.sweep = None + self._include_sweeps = include_sweeps + self.dir = os.path.join(self._base_dir, *self.path) + try: + os.makedirs(self.dir) + except OSError: + pass + self._summary = None + self._metadata: Optional[Dict[str, Any]] = None + self._state = _attrs.get("state", "not found") + self.server_provides_internal_id_field: Optional[bool] = None + + self.load(force=not _attrs) + + @property + def state(self): + return self._state + + @property + def entity(self): + return self._entity + + @property + def username(self): + wandb.termwarn("Run.username is deprecated. Please use Run.entity instead.") + return self._entity + + @property + def storage_id(self): + # For compatibility with wandb.Run, which has storage IDs + # in self.storage_id and names in self.id. + + return self._attrs.get("id") + + @property + def id(self): + return self._attrs.get("name") + + @id.setter + def id(self, new_id): + attrs = self._attrs + attrs["name"] = new_id + return new_id + + @property + def name(self): + return self._attrs.get("displayName") + + @name.setter + def name(self, new_name): + self._attrs["displayName"] = new_name + return new_name + + @classmethod + def create( + cls, + api, + run_id=None, + project=None, + entity=None, + state: Literal["running", "pending"] = "running", + ): + """Create a run for the given project.""" + run_id = run_id or runid.generate_id() + project = project or api.settings.get("project") or "uncategorized" + mutation = gql( + """ + mutation UpsertBucket($project: String, $entity: String, $name: String!, $state: String) { + upsertBucket(input: {modelName: $project, entityName: $entity, name: $name, state: $state}) { + bucket { + project { + name + entity { name } + } + id + name + } + inserted + } + } + """ + ) + variables = { + "entity": entity, + "project": project, + "name": run_id, + "state": state, + } + res = api.client.execute(mutation, variable_values=variables) + res = res["upsertBucket"]["bucket"] + return Run( + api.client, + res["project"]["entity"]["name"], + res["project"]["name"], + res["name"], + { + "id": res["id"], + "config": "{}", + "systemMetrics": "{}", + "summaryMetrics": "{}", + "tags": [], + "description": None, + "notes": None, + "state": state, + }, + ) + + def load(self, force=False): + if force or not self._attrs: + query = gql(f"""#graphql + query Run($project: String!, $entity: String!, $name: String!) {{ + project(name: $project, entityName: $entity) {{ + run(name: $name) {{ + {"projectId" if _server_provides_internal_id_for_project(self.client) else ""} + ...RunFragment + }} + }} + }} + {RUN_FRAGMENT} + """) + + response = self._exec(query) + if ( + response is None + or response.get("project") is None + or response["project"].get("run") is None + ): + raise ValueError("Could not find run {}".format(self)) + self._attrs = response["project"]["run"] + self._state = self._attrs["state"] + + if self._include_sweeps and self.sweep_name and not self.sweep: + # There may be a lot of runs. Don't bother pulling them all + # just for the sake of this one. + self.sweep = public.Sweep.get( + self.client, + self.entity, + self.project, + self.sweep_name, + withRuns=False, + ) + + if "projectId" in self._attrs: + self._project_internal_id = int(self._attrs["projectId"]) + else: + self._project_internal_id = None + + try: + self._attrs["summaryMetrics"] = ( + json.loads(self._attrs["summaryMetrics"]) + if self._attrs.get("summaryMetrics") + else {} + ) + except json.decoder.JSONDecodeError: + # ignore invalid utf-8 or control characters + self._attrs["summaryMetrics"] = json.loads( + self._attrs["summaryMetrics"], + strict=False, + ) + self._attrs["systemMetrics"] = ( + json.loads(self._attrs["systemMetrics"]) + if self._attrs.get("systemMetrics") + else {} + ) + if self._attrs.get("user"): + self.user = public.User(self.client, self._attrs["user"]) + config_user, config_raw = {}, {} + for key, value in json.loads(self._attrs.get("config") or "{}").items(): + config = config_raw if key in WANDB_INTERNAL_KEYS else config_user + if isinstance(value, dict) and "value" in value: + config[key] = value["value"] + else: + config[key] = value + config_raw.update(config_user) + self._attrs["config"] = config_user + self._attrs["rawconfig"] = config_raw + return self._attrs + + @normalize_exceptions + def wait_until_finished(self): + query = gql( + """ + query RunState($project: String!, $entity: String!, $name: String!) { + project(name: $project, entityName: $entity) { + run(name: $name) { + state + } + } + } + """ + ) + while True: + res = self._exec(query) + state = res["project"]["run"]["state"] + if state in ["finished", "crashed", "failed"]: + self._attrs["state"] = state + self._state = state + return + time.sleep(5) + + @normalize_exceptions + def update(self): + """Persist changes to the run object to the wandb backend.""" + mutation = gql( + """ + mutation UpsertBucket($id: String!, $description: String, $display_name: String, $notes: String, $tags: [String!], $config: JSONString!, $groupName: String, $jobType: String) {{ + upsertBucket(input: {{id: $id, description: $description, displayName: $display_name, notes: $notes, tags: $tags, config: $config, groupName: $groupName, jobType: $jobType}}) {{ + bucket {{ + ...RunFragment + }} + }} + }} + {} + """.format(RUN_FRAGMENT) + ) + _ = self._exec( + mutation, + id=self.storage_id, + tags=self.tags, + description=self.description, + notes=self.notes, + display_name=self.display_name, + config=self.json_config, + groupName=self.group, + jobType=self.job_type, + ) + self.summary.update() + + @normalize_exceptions + def delete(self, delete_artifacts=False): + """Delete the given run from the wandb backend.""" + mutation = gql( + """ + mutation DeleteRun( + $id: ID!, + {} + ) {{ + deleteRun(input: {{ + id: $id, + {} + }}) {{ + clientMutationId + }} + }} + """.format( + "$deleteArtifacts: Boolean" if delete_artifacts else "", + "deleteArtifacts: $deleteArtifacts" if delete_artifacts else "", + ) + ) + + self.client.execute( + mutation, + variable_values={ + "id": self.storage_id, + "deleteArtifacts": delete_artifacts, + }, + ) + + def save(self): + self.update() + + @property + def json_config(self): + config = {} + if "_wandb" in self.rawconfig: + config["_wandb"] = {"value": self.rawconfig["_wandb"], "desc": None} + for k, v in self.config.items(): + config[k] = {"value": v, "desc": None} + return json.dumps(config) + + def _exec(self, query, **kwargs): + """Execute a query against the cloud backend.""" + variables = {"entity": self.entity, "project": self.project, "name": self.id} + variables.update(kwargs) + return self.client.execute(query, variable_values=variables) + + def _sampled_history(self, keys, x_axis="_step", samples=500): + spec = {"keys": [x_axis] + keys, "samples": samples} + query = gql( + """ + query RunSampledHistory($project: String!, $entity: String!, $name: String!, $specs: [JSONString!]!) { + project(name: $project, entityName: $entity) { + run(name: $name) { sampledHistory(specs: $specs) } + } + } + """ + ) + + response = self._exec(query, specs=[json.dumps(spec)]) + # sampledHistory returns one list per spec, we only send one spec + return response["project"]["run"]["sampledHistory"][0] + + def _full_history(self, samples=500, stream="default"): + node = "history" if stream == "default" else "events" + query = gql( + """ + query RunFullHistory($project: String!, $entity: String!, $name: String!, $samples: Int) {{ + project(name: $project, entityName: $entity) {{ + run(name: $name) {{ {}(samples: $samples) }} + }} + }} + """.format(node) + ) + + response = self._exec(query, samples=samples) + return [json.loads(line) for line in response["project"]["run"][node]] + + @normalize_exceptions + def files(self, names=None, per_page=50): + """Return a file path for each file named. + + Args: + names (list): names of the requested files, if empty returns all files + per_page (int): number of results per page. + + Returns: + A `Files` object, which is an iterator over `File` objects. + """ + return public.Files(self.client, self, names or [], per_page) + + @normalize_exceptions + def file(self, name): + """Return the path of a file with a given name in the artifact. + + Args: + name (str): name of requested file. + + Returns: + A `File` matching the name argument. + """ + return public.Files(self.client, self, [name])[0] + + @normalize_exceptions + def upload_file(self, path, root="."): + """Upload a file. + + Args: + path (str): name of file to upload. + root (str): the root path to save the file relative to. i.e. + If you want to have the file saved in the run as "my_dir/file.txt" + and you're currently in "my_dir" you would set root to "../". + + Returns: + A `File` matching the name argument. + """ + api = InternalApi( + default_settings={"entity": self.entity, "project": self.project}, + retry_timedelta=RETRY_TIMEDELTA, + ) + api.set_current_run_id(self.id) + root = os.path.abspath(root) + name = os.path.relpath(path, root) + with open(os.path.join(root, name), "rb") as f: + api.push({LogicalPath(name): f}) + return public.Files(self.client, self, [name])[0] + + @normalize_exceptions + def history( + self, samples=500, keys=None, x_axis="_step", pandas=True, stream="default" + ): + """Return sampled history metrics for a run. + + This is simpler and faster if you are ok with the history records being sampled. + + Args: + samples : (int, optional) The number of samples to return + pandas : (bool, optional) Return a pandas dataframe + keys : (list, optional) Only return metrics for specific keys + x_axis : (str, optional) Use this metric as the xAxis defaults to _step + stream : (str, optional) "default" for metrics, "system" for machine metrics + + Returns: + pandas.DataFrame: If pandas=True returns a `pandas.DataFrame` of history + metrics. + list of dicts: If pandas=False returns a list of dicts of history metrics. + """ + if keys is not None and not isinstance(keys, list): + wandb.termerror("keys must be specified in a list") + return [] + if keys is not None and len(keys) > 0 and not isinstance(keys[0], str): + wandb.termerror("keys argument must be a list of strings") + return [] + + if keys and stream != "default": + wandb.termerror("stream must be default when specifying keys") + return [] + elif keys: + lines = self._sampled_history(keys=keys, x_axis=x_axis, samples=samples) + else: + lines = self._full_history(samples=samples, stream=stream) + if pandas: + pd = util.get_module("pandas") + if pd: + lines = pd.DataFrame.from_records(lines) + else: + wandb.termwarn("Unable to load pandas, call history with pandas=False") + return lines + + @normalize_exceptions + def scan_history(self, keys=None, page_size=1000, min_step=None, max_step=None): + """Returns an iterable collection of all history records for a run. + + Example: + Export all the loss values for an example run + + ```python + run = api.run("l2k2/examples-numpy-boston/i0wt6xua") + history = run.scan_history(keys=["Loss"]) + losses = [row["Loss"] for row in history] + ``` + + Args: + keys ([str], optional): only fetch these keys, and only fetch rows that have all of keys defined. + page_size (int, optional): size of pages to fetch from the api. + min_step (int, optional): the minimum number of pages to scan at a time. + max_step (int, optional): the maximum number of pages to scan at a time. + + Returns: + An iterable collection over history records (dict). + """ + if keys is not None and not isinstance(keys, list): + wandb.termerror("keys must be specified in a list") + return [] + if keys is not None and len(keys) > 0 and not isinstance(keys[0], str): + wandb.termerror("keys argument must be a list of strings") + return [] + + last_step = self.lastHistoryStep + # set defaults for min/max step + if min_step is None: + min_step = 0 + if max_step is None: + max_step = last_step + 1 + # if the max step is past the actual last step, clamp it down + if max_step > last_step: + max_step = last_step + 1 + if keys is None: + return public.HistoryScan( + run=self, + client=self.client, + page_size=page_size, + min_step=min_step, + max_step=max_step, + ) + else: + return public.SampledHistoryScan( + run=self, + client=self.client, + keys=keys, + page_size=page_size, + min_step=min_step, + max_step=max_step, + ) + + @normalize_exceptions + def logged_artifacts(self, per_page: int = 100) -> public.RunArtifacts: + """Fetches all artifacts logged by this run. + + Retrieves all output artifacts that were logged during the run. Returns a + paginated result that can be iterated over or collected into a single list. + + Args: + per_page: Number of artifacts to fetch per API request. + + Returns: + An iterable collection of all Artifact objects logged as outputs during this run. + + Example: + >>> import wandb + >>> import tempfile + >>> with tempfile.NamedTemporaryFile( + ... mode="w", delete=False, suffix=".txt" + ... ) as tmp: + ... tmp.write("This is a test artifact") + ... tmp_path = tmp.name + >>> run = wandb.init(project="artifact-example") + >>> artifact = wandb.Artifact("test_artifact", type="dataset") + >>> artifact.add_file(tmp_path) + >>> run.log_artifact(artifact) + >>> run.finish() + >>> api = wandb.Api() + >>> finished_run = api.run(f"{run.entity}/{run.project}/{run.id}") + >>> for logged_artifact in finished_run.logged_artifacts(): + ... print(logged_artifact.name) + test_artifact + """ + return public.RunArtifacts(self.client, self, mode="logged", per_page=per_page) + + @normalize_exceptions + def used_artifacts(self, per_page: int = 100) -> public.RunArtifacts: + """Fetches artifacts explicitly used by this run. + + Retrieves only the input artifacts that were explicitly declared as used + during the run, typically via `run.use_artifact()`. Returns a paginated + result that can be iterated over or collected into a single list. + + Args: + per_page: Number of artifacts to fetch per API request. + + Returns: + An iterable collection of Artifact objects explicitly used as inputs in this run. + + Example: + >>> import wandb + >>> run = wandb.init(project="artifact-example") + >>> run.use_artifact("test_artifact:latest") + >>> run.finish() + >>> api = wandb.Api() + >>> finished_run = api.run(f"{run.entity}/{run.project}/{run.id}") + >>> for used_artifact in finished_run.used_artifacts(): + ... print(used_artifact.name) + test_artifact + """ + return public.RunArtifacts(self.client, self, mode="used", per_page=per_page) + + @normalize_exceptions + def use_artifact(self, artifact, use_as=None): + """Declare an artifact as an input to a run. + + Args: + artifact (`Artifact`): An artifact returned from + `wandb.Api().artifact(name)` + use_as (string, optional): A string identifying + how the artifact is used in the script. Used + to easily differentiate artifacts used in a + run, when using the beta wandb launch + feature's artifact swapping functionality. + + Returns: + A `Artifact` object. + """ + api = InternalApi( + default_settings={"entity": self.entity, "project": self.project}, + retry_timedelta=RETRY_TIMEDELTA, + ) + api.set_current_run_id(self.id) + + if isinstance(artifact, wandb.Artifact) and not artifact.is_draft(): + api.use_artifact( + artifact.id, + use_as=use_as or artifact.name, + artifact_entity_name=artifact.entity, + artifact_project_name=artifact.project, + ) + return artifact + elif isinstance(artifact, wandb.Artifact) and artifact.is_draft(): + raise ValueError( + "Only existing artifacts are accepted by this api. " + "Manually create one with `wandb artifact put`" + ) + else: + raise ValueError("You must pass a wandb.Api().artifact() to use_artifact") + + @normalize_exceptions + def log_artifact( + self, + artifact: "wandb.Artifact", + aliases: Optional[Collection[str]] = None, + tags: Optional[Collection[str]] = None, + ): + """Declare an artifact as output of a run. + + Args: + artifact (`Artifact`): An artifact returned from + `wandb.Api().artifact(name)`. + aliases (list, optional): Aliases to apply to this artifact. + tags: (list, optional) Tags to apply to this artifact, if any. + + Returns: + A `Artifact` object. + """ + api = InternalApi( + default_settings={"entity": self.entity, "project": self.project}, + retry_timedelta=RETRY_TIMEDELTA, + ) + api.set_current_run_id(self.id) + + if not isinstance(artifact, wandb.Artifact): + raise TypeError("You must pass a wandb.Api().artifact() to use_artifact") + if artifact.is_draft(): + raise ValueError( + "Only existing artifacts are accepted by this api. " + "Manually create one with `wandb artifact put`" + ) + if ( + self.entity != artifact.source_entity + or self.project != artifact.source_project + ): + raise ValueError("A run can't log an artifact to a different project.") + + artifact_collection_name = artifact.source_name.split(":")[0] + api.create_artifact( + artifact.type, + artifact_collection_name, + artifact.digest, + aliases=aliases, + tags=tags, + ) + return artifact + + @property + def summary(self): + if self._summary is None: + from wandb.old.summary import HTTPSummary + + # TODO: fix the outdir issue + self._summary = HTTPSummary(self, self.client, summary=self.summary_metrics) + return self._summary + + @property + def path(self): + return [ + urllib.parse.quote_plus(str(self.entity)), + urllib.parse.quote_plus(str(self.project)), + urllib.parse.quote_plus(str(self.id)), + ] + + @property + def url(self): + path = self.path + path.insert(2, "runs") + return self.client.app_url + "/".join(path) + + @property + def metadata(self): + if self._metadata is None: + try: + f = self.file("wandb-metadata.json") + session = self.client._client.transport.session + response = session.get(f.url, timeout=5) + response.raise_for_status() + contents = response.content + self._metadata = json_util.loads(contents) + except: # noqa: E722 + # file doesn't exist, or can't be downloaded, or can't be parsed + pass + return self._metadata + + @property + def lastHistoryStep(self): # noqa: N802 + query = gql( + """ + query RunHistoryKeys($project: String!, $entity: String!, $name: String!) { + project(name: $project, entityName: $entity) { + run(name: $name) { historyKeys } + } + } + """ + ) + response = self._exec(query) + if ( + response is None + or response.get("project") is None + or response["project"].get("run") is None + or response["project"]["run"].get("historyKeys") is None + ): + return -1 + history_keys = response["project"]["run"]["historyKeys"] + return history_keys["lastStep"] if "lastStep" in history_keys else -1 + + def to_html(self, height=420, hidden=False): + """Generate HTML containing an iframe displaying this run.""" + url = self.url + "?jupyter=true" + style = f"border:none;width:100%;height:{height}px;" + prefix = "" + if hidden: + style += "display:none;" + prefix = ipython.toggle_button() + return prefix + f"" + + def _repr_html_(self) -> str: + return self.to_html() + + def __repr__(self): + return "".format("/".join(self.path), self.state) diff --git a/lib/python3.12/site-packages/wandb/apis/public/sweeps.py b/lib/python3.12/site-packages/wandb/apis/public/sweeps.py new file mode 100644 index 0000000000000000000000000000000000000000..3384a5193667c6d6b716a6e1a57b341de3046805 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/public/sweeps.py @@ -0,0 +1,240 @@ +"""Public API: sweeps.""" + +import urllib +from typing import Optional + +from wandb_gql import gql + +import wandb +from wandb import util +from wandb.apis import public +from wandb.apis.attrs import Attrs +from wandb.sdk.lib import ipython + +SWEEP_FRAGMENT = """fragment SweepFragment on Sweep { + id + name + method + state + description + displayName + bestLoss + config + createdAt + updatedAt + runCount +} +""" + + +class Sweep(Attrs): + """A set of runs associated with a sweep. + + Examples: + Instantiate with: + ``` + api = wandb.Api() + sweep = api.sweep(path / to / sweep) + ``` + + Attributes: + runs: (`Runs`) list of runs + id: (str) sweep id + project: (str) name of project + config: (str) dictionary of sweep configuration + state: (str) the state of the sweep + expected_run_count: (int) number of expected runs for the sweep + """ + + QUERY = gql( + """ + query Sweep($project: String, $entity: String, $name: String!) { + project(name: $project, entityName: $entity) { + sweep(sweepName: $name) { + id + name + state + runCountExpected + bestLoss + config + } + } + } + """ + ) + + LEGACY_QUERY = gql( + """ + query Sweep($project: String, $entity: String, $name: String!) { + project(name: $project, entityName: $entity) { + sweep(sweepName: $name) { + id + name + state + bestLoss + config + } + } + } + """ + ) + + def __init__(self, client, entity, project, sweep_id, attrs=None): + # TODO: Add agents / flesh this out. + super().__init__(dict(attrs or {})) + self.client = client + self._entity = entity + self.project = project + self.id = sweep_id + self.runs = [] + + self.load(force=not attrs) + + @property + def entity(self): + return self._entity + + @property + def username(self): + wandb.termwarn("Sweep.username is deprecated. please use Sweep.entity instead.") + return self._entity + + @property + def config(self): + return util.load_yaml(self._attrs["config"]) + + def load(self, force: bool = False): + if force or not self._attrs: + sweep = self.get(self.client, self.entity, self.project, self.id) + if sweep is None: + raise ValueError("Could not find sweep {}".format(self)) + self._attrs = sweep._attrs + self.runs = sweep.runs + + return self._attrs + + @property + def order(self): + if self._attrs.get("config") and self.config.get("metric"): + sort_order = self.config["metric"].get("goal", "minimize") + prefix = "+" if sort_order == "minimize" else "-" + return public.QueryGenerator.format_order_key( + prefix + self.config["metric"]["name"] + ) + + def best_run(self, order=None): + """Return the best run sorted by the metric defined in config or the order passed in.""" + if order is None: + order = self.order + else: + order = public.QueryGenerator.format_order_key(order) + if order is None: + wandb.termwarn( + "No order specified and couldn't find metric in sweep config, returning most recent run" + ) + else: + wandb.termlog("Sorting runs by {}".format(order)) + filters = {"$and": [{"sweep": self.id}]} + try: + return public.Runs( + self.client, + self.entity, + self.project, + order=order, + filters=filters, + per_page=1, + )[0] + except IndexError: + return None + + @property + def expected_run_count(self) -> Optional[int]: + """Return the number of expected runs in the sweep or None for infinite runs.""" + return self._attrs.get("runCountExpected") + + @property + def path(self): + return [ + urllib.parse.quote_plus(str(self.entity)), + urllib.parse.quote_plus(str(self.project)), + urllib.parse.quote_plus(str(self.id)), + ] + + @property + def url(self): + path = self.path + path.insert(2, "sweeps") + return self.client.app_url + "/".join(path) + + @property + def name(self): + return self.config.get("name") or self.id + + @classmethod + def get( + cls, + client, + entity=None, + project=None, + sid=None, + order=None, + query=None, + **kwargs, + ): + """Execute a query against the cloud backend.""" + if query is None: + query = cls.QUERY + + variables = { + "entity": entity, + "project": project, + "name": sid, + } + variables.update(kwargs) + + response = None + try: + response = client.execute(query, variable_values=variables) + except Exception: + # Don't handle exception, rely on legacy query + # TODO(gst): Implement updated introspection workaround + query = cls.LEGACY_QUERY + response = client.execute(query, variable_values=variables) + + if ( + not response + or not response.get("project") + or not response["project"].get("sweep") + ): + return None + + sweep_response = response["project"]["sweep"] + sweep = cls(client, entity, project, sid, attrs=sweep_response) + sweep.runs = public.Runs( + client, + entity, + project, + order=order, + per_page=10, + filters={"$and": [{"sweep": sweep.id}]}, + ) + + return sweep + + def to_html(self, height=420, hidden=False): + """Generate HTML containing an iframe displaying this sweep.""" + url = self.url + "?jupyter=true" + style = f"border:none;width:100%;height:{height}px;" + prefix = "" + if hidden: + style += "display:none;" + prefix = ipython.toggle_button("sweep") + return prefix + f"" + + def _repr_html_(self) -> str: + return self.to_html() + + def __repr__(self): + return "".format( + "/".join(self.path), self._attrs.get("state", "Unknown State") + ) diff --git a/lib/python3.12/site-packages/wandb/apis/public/teams.py b/lib/python3.12/site-packages/wandb/apis/public/teams.py new file mode 100644 index 0000000000000000000000000000000000000000..bb55a1e694c54d5b9aa980422dd11e2533fbb997 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/public/teams.py @@ -0,0 +1,198 @@ +"""Public API: teams.""" + +import requests +from wandb_gql import gql + +from wandb.apis.attrs import Attrs + + +class Member(Attrs): + DELETE_MEMBER_MUTATION = gql( + """ + mutation DeleteInvite($id: String, $entityName: String) { + deleteInvite(input: {id: $id, entityName: $entityName}) { + success + } + } + """ + ) + + def __init__(self, client, team, attrs): + super().__init__(attrs) + self._client = client + self.team = team + + def delete(self): + """Remove a member from a team. + + Returns: + Boolean indicating success + """ + try: + return self._client.execute( + self.DELETE_MEMBER_MUTATION, {"id": self.id, "entityName": self.team} + )["deleteInvite"]["success"] + except requests.exceptions.HTTPError: + return False + + def __repr__(self): + return f"" + + +class Team(Attrs): + CREATE_TEAM_MUTATION = gql( + """ + mutation CreateTeam($teamName: String!, $teamAdminUserName: String) { + createTeam(input: {teamName: $teamName, teamAdminUserName: $teamAdminUserName}) { + entity { + id + name + available + photoUrl + limits + } + } + } + """ + ) + CREATE_INVITE_MUTATION = gql( + """ + mutation CreateInvite($entityName: String!, $email: String, $username: String, $admin: Boolean) { + createInvite(input: {entityName: $entityName, email: $email, username: $username, admin: $admin}) { + invite { + id + name + email + createdAt + toUser { + name + } + } + } + } + """ + ) + TEAM_QUERY = gql( + """ + query Entity($name: String!) { + entity(name: $name) { + id + name + available + photoUrl + readOnly + readOnlyAdmin + isTeam + privateOnly + storageBytes + codeSavingEnabled + defaultAccess + isPaid + members { + id + admin + pending + email + username + name + photoUrl + accountType + apiKey + } + } + } + """ + ) + CREATE_SERVICE_ACCOUNT_MUTATION = gql( + """ + mutation CreateServiceAccount($entityName: String!, $description: String!) { + createServiceAccount( + input: {description: $description, entityName: $entityName} + ) { + user { + id + } + } + } + """ + ) + + def __init__(self, client, name, attrs=None): + super().__init__(attrs or {}) + self._client = client + self.name = name + self.load() + + @classmethod + def create(cls, api, team, admin_username=None): + """Create a new team. + + Args: + api: (`Api`) The api instance to use + team: (str) The name of the team + admin_username: (str) optional username of the admin user of the team, defaults to the current user. + + Returns: + A `Team` object + """ + try: + api.client.execute( + cls.CREATE_TEAM_MUTATION, + {"teamName": team, "teamAdminUserName": admin_username}, + ) + except requests.exceptions.HTTPError: + pass + return Team(api.client, team) + + def invite(self, username_or_email, admin=False): + """Invite a user to a team. + + Args: + username_or_email: (str) The username or email address of the user you want to invite + admin: (bool) Whether to make this user a team admin, defaults to False + + Returns: + True on success, False if user was already invited or didn't exist + """ + variables = {"entityName": self.name, "admin": admin} + if "@" in username_or_email: + variables["email"] = username_or_email + else: + variables["username"] = username_or_email + try: + self._client.execute(self.CREATE_INVITE_MUTATION, variables) + except requests.exceptions.HTTPError: + return False + return True + + def create_service_account(self, description): + """Create a service account for the team. + + Args: + description: (str) A description for this service account + + Returns: + The service account `Member` object, or None on failure + """ + try: + self._client.execute( + self.CREATE_SERVICE_ACCOUNT_MUTATION, + {"description": description, "entityName": self.name}, + ) + self.load(True) + return self.members[-1] + except requests.exceptions.HTTPError: + return None + + def load(self, force=False): + if force or not self._attrs: + response = self._client.execute(self.TEAM_QUERY, {"name": self.name}) + self._attrs = response["entity"] + self._attrs["members"] = [ + Member(self._client, self.name, member) + for member in self._attrs["members"] + ] + return self._attrs + + def __repr__(self): + return f"" diff --git a/lib/python3.12/site-packages/wandb/apis/public/users.py b/lib/python3.12/site-packages/wandb/apis/public/users.py new file mode 100644 index 0000000000000000000000000000000000000000..eda80581a64e761cfdf2211fdd2f7fac987f16ff --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/public/users.py @@ -0,0 +1,136 @@ +"""Public API: users.""" + +import requests +from wandb_gql import gql + +import wandb +from wandb.apis.attrs import Attrs + + +class User(Attrs): + CREATE_USER_MUTATION = gql( + """ + mutation CreateUserFromAdmin($email: String!, $admin: Boolean) { + createUser(input: {email: $email, admin: $admin}) { + user { + id + name + username + email + admin + } + } + } + """ + ) + + DELETE_API_KEY_MUTATION = gql( + """ + mutation DeleteApiKey($id: String!) { + deleteApiKey(input: {id: $id}) { + success + } + } + """ + ) + GENERATE_API_KEY_MUTATION = gql( + """ + mutation GenerateApiKey($description: String) { + generateApiKey(input: {description: $description}) { + apiKey { + id + name + } + } + } + """ + ) + + def __init__(self, client, attrs): + super().__init__(attrs) + self._client = client + self._user_api = None + + @property + def user_api(self): + """An instance of the api using credentials from the user.""" + if self._user_api is None and len(self.api_keys) > 0: + self._user_api = wandb.Api(api_key=self.api_keys[0]) + return self._user_api + + @classmethod + def create(cls, api, email, admin=False): + """Create a new user. + + Args: + api: (`Api`) The api instance to use + email: (str) The name of the team + admin: (bool) Whether this user should be a global instance admin + + Returns: + A `User` object + """ + res = api.client.execute( + cls.CREATE_USER_MUTATION, + {"email": email, "admin": admin}, + ) + return User(api.client, res["createUser"]["user"]) + + @property + def api_keys(self): + if self._attrs.get("apiKeys") is None: + return [] + return [k["node"]["name"] for k in self._attrs["apiKeys"]["edges"]] + + @property + def teams(self): + if self._attrs.get("teams") is None: + return [] + return [k["node"]["name"] for k in self._attrs["teams"]["edges"]] + + def delete_api_key(self, api_key): + """Delete a user's api key. + + Returns: + Boolean indicating success + + Raises: + ValueError if the api_key couldn't be found + """ + idx = self.api_keys.index(api_key) + try: + self._client.execute( + self.DELETE_API_KEY_MUTATION, + {"id": self._attrs["apiKeys"]["edges"][idx]["node"]["id"]}, + ) + except requests.exceptions.HTTPError: + return False + return True + + def generate_api_key(self, description=None): + """Generate a new api key. + + Returns: + The new api key, or None on failure + """ + try: + # We must make this call using credentials from the original user + key = self.user_api.client.execute( + self.GENERATE_API_KEY_MUTATION, {"description": description} + )["generateApiKey"]["apiKey"] + self._attrs["apiKeys"]["edges"].append({"node": key}) + return key["name"] + except (requests.exceptions.HTTPError, AttributeError): + return None + + def __repr__(self): + if "email" in self._attrs: + return f"" + elif "username" in self._attrs: + return f"" + elif "id" in self._attrs: + return f"" + elif "name" in self._attrs: + return f"" + else: + return "" diff --git a/lib/python3.12/site-packages/wandb/apis/public/utils.py b/lib/python3.12/site-packages/wandb/apis/public/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c868b9304110f7a3195ab45564f3bc04b209e321 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/public/utils.py @@ -0,0 +1,210 @@ +import re +from enum import Enum +from typing import Any, Dict, Iterable, Mapping, Optional, Set +from urllib.parse import urlparse + +from wandb_gql import gql +from wandb_graphql.language import ast, visitor + +from wandb._iterutils import one +from wandb.sdk.artifacts._validators import is_artifact_registry_project +from wandb.sdk.internal.internal_api import Api as InternalApi + + +def parse_s3_url_to_s3_uri(url) -> str: + """Convert an S3 HTTP(S) URL to an S3 URI. + + Arguments: + url (str): The S3 URL to convert, in the format + 'http(s)://.s3..amazonaws.com/'. + or 'http(s)://.s3.amazonaws.com/' + + Returns: + str: The corresponding S3 URI in the format 's3:///'. + + Raises: + ValueError: If the provided URL is not a valid S3 URL. + """ + # Regular expression to match S3 URL pattern + s3_pattern = r"^https?://.*s3.*amazonaws\.com.*" + parsed_url = urlparse(url) + + # Check if it's an S3 URL + match = re.match(s3_pattern, parsed_url.geturl()) + if not match: + raise ValueError("Invalid S3 URL") + + # Extract bucket name and key + bucket_name, *_ = parsed_url.netloc.split(".") + key = parsed_url.path.lstrip("/") + + # Construct the S3 URI + s3_uri = f"s3://{bucket_name}/{key}" + + return s3_uri + + +class PathType(Enum): + """We have lots of different paths users pass in to fetch artifacts, projects, etc. + + This enum is used for specifying what format the path is in given a string path. + """ + + PROJECT = "PROJECT" + ARTIFACT = "ARTIFACT" + + +def parse_org_from_registry_path(path: str, path_type: PathType) -> str: + """Parse the org from a registry path. + + Essentially fetching the "entity" from the path but for Registries the entity is actually the org. + + Args: + path (str): The path to parse. Can be a project path / or or an + artifact path like // or / or + path_type (PathType): The type of path to parse. + """ + parts = path.split("/") + expected_parts = 3 if path_type == PathType.ARTIFACT else 2 + + if len(parts) >= expected_parts: + org, project = parts[:2] + if is_artifact_registry_project(project): + return org + return "" + + +def fetch_org_from_settings_or_entity( + settings: dict, default_entity: Optional[str] = None +) -> str: + """Fetch the org from either the settings or deriving it from the entity. + + Returns the org from the settings if available. If no org is passed in or set, the entity is used to fetch the org. + + Args: + organization (str | None): The organization to fetch the org for. + settings (dict): The settings to fetch the org for. + default_entity (str | None): The default entity to fetch the org for. + """ + if (organization := settings.get("organization")) is None: + # Fetch the org via the Entity. Won't work if default entity is a personal entity and belongs to multiple orgs + entity = settings.get("entity") or default_entity + if entity is None: + raise ValueError( + "No entity specified and can't fetch organization from the entity" + ) + entity_orgs = InternalApi()._fetch_orgs_and_org_entities_from_entity(entity) + entity_org = one( + entity_orgs, + too_short=ValueError( + "No organizations found for entity. Please specify an organization in the settings." + ), + too_long=ValueError( + "Multiple organizations found for entity. Please specify an organization in the settings." + ), + ) + organization = entity_org.display_name + return organization + + +class _GQLCompatRewriter(visitor.Visitor): + """GraphQL AST visitor to rewrite queries/mutations to be compatible with older server versions.""" + + omit_variables: Set[str] + omit_fragments: Set[str] + omit_fields: Set[str] + rename_fields: Dict[str, str] + + def __init__( + self, + omit_variables: Optional[Iterable[str]] = None, + omit_fragments: Optional[Iterable[str]] = None, + omit_fields: Optional[Iterable[str]] = None, + rename_fields: Optional[Mapping[str, str]] = None, + ): + self.omit_variables = set(omit_variables or ()) + self.omit_fragments = set(omit_fragments or ()) + self.omit_fields = set(omit_fields or ()) + self.rename_fields = dict(rename_fields or {}) + + def enter_VariableDefinition(self, node: ast.VariableDefinition, *_, **__) -> Any: # noqa: N802 + if node.variable.name.value in self.omit_variables: + return visitor.REMOVE + # return node + + def enter_ObjectField(self, node: ast.ObjectField, *_, **__) -> Any: # noqa: N802 + # For context, note that e.g.: + # + # {description: $description + # ...} + # + # Is parsed as: + # + # ObjectValue(fields=[ + # ObjectField(name=Name(value='description'), value=Variable(name=Name(value='description'))), + # ...]) + if ( + isinstance(var := node.value, ast.Variable) + and var.name.value in self.omit_variables + ): + return visitor.REMOVE + + def enter_Argument(self, node: ast.Argument, *_, **__) -> Any: # noqa: N802 + if node.name.value in self.omit_variables: + return visitor.REMOVE + + def enter_FragmentDefinition(self, node: ast.FragmentDefinition, *_, **__) -> Any: # noqa: N802 + if node.name.value in self.omit_fragments: + return visitor.REMOVE + + def enter_FragmentSpread(self, node: ast.FragmentSpread, *_, **__) -> Any: # noqa: N802 + if node.name.value in self.omit_fragments: + return visitor.REMOVE + + def enter_Field(self, node: ast.Field, *_, **__) -> Any: # noqa: N802 + if node.name.value in self.omit_fields: + return visitor.REMOVE + if new_name := self.rename_fields.get(node.name.value): + node.name.value = new_name + return node + + def leave_Field(self, node: ast.Field, *_, **__) -> Any: # noqa: N802 + # If the field had a selection set, but now it's empty, remove the field entirely + if (node.selection_set is not None) and (not node.selection_set.selections): + return visitor.REMOVE + + +def gql_compat( + request_string: str, + omit_variables: Optional[Iterable[str]] = None, + omit_fragments: Optional[Iterable[str]] = None, + omit_fields: Optional[Iterable[str]] = None, + rename_fields: Optional[Mapping[str, str]] = None, +) -> ast.Document: + """Rewrite a GraphQL request string to ensure compatibility with older server versions. + + Args: + request_string (str): The GraphQL request string to rewrite. + omit_variables (Iterable[str] | None): Names of variables to remove from the request string. + omit_fragments (Iterable[str] | None): Names of fragments to remove from the request string. + omit_fields (Iterable[str] | None): Names of fields to remove from the request string. + rename_fields (Mapping[str, str] | None): + A mapping of fields to rename in the request string, given as `{old_name -> new_name}`. + + Returns: + str: Modified GraphQL request string with fragments on omitted types removed. + """ + # Parse the request into a GraphQL AST + doc = gql(request_string) + + if not (omit_variables or omit_fragments or omit_fields or rename_fields): + return doc + + # Visit the AST with our visitor to filter out unwanted fragments + rewriter = _GQLCompatRewriter( + omit_variables=omit_variables, + omit_fragments=omit_fragments, + omit_fields=omit_fields, + rename_fields=rename_fields, + ) + return visitor.visit(doc, rewriter) diff --git a/lib/python3.12/site-packages/wandb/apis/reports/__init__.py b/lib/python3.12/site-packages/wandb/apis/reports/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6d2b9ad9385524f2b868e498d3030de028fac546 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/reports/__init__.py @@ -0,0 +1 @@ +from .v2 import * # noqa: F403 diff --git a/lib/python3.12/site-packages/wandb/apis/reports/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/reports/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..83854514833363b3411fcfb71819d2a4e9d16583 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/reports/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/reports/v1/__init__.py b/lib/python3.12/site-packages/wandb/apis/reports/v1/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0bd41baac9f44dd1f7dc5145779d33fe95d10763 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/reports/v1/__init__.py @@ -0,0 +1,8 @@ +import wandb + +try: + from wandb_workspaces.reports.v1 import * # noqa: F403 +except ImportError: + wandb.termerror( + "Failed to import wandb_workspaces. To edit reports programmatically, please install it using `pip install wandb[workspaces]`." + ) diff --git a/lib/python3.12/site-packages/wandb/apis/reports/v1/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/reports/v1/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b90934129b9673ed71bfae021ffa3544bdfb9ae3 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/reports/v1/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/reports/v2/__init__.py b/lib/python3.12/site-packages/wandb/apis/reports/v2/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..532129620940c7deb87ae483a08059cb792d15f4 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/reports/v2/__init__.py @@ -0,0 +1,8 @@ +import wandb + +try: + from wandb_workspaces.reports.v2 import * # noqa: F403 +except ImportError: + wandb.termerror( + "Failed to import wandb_workspaces. To edit reports programmatically, please install it using `pip install wandb[workspaces]`." + ) diff --git a/lib/python3.12/site-packages/wandb/apis/reports/v2/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/reports/v2/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d33fe0340315b5846b588026d2d701cfa467883a Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/reports/v2/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/apis/workspaces/__init__.py b/lib/python3.12/site-packages/wandb/apis/workspaces/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..269a35473d08b01200c593d2c8958c85349a0749 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/apis/workspaces/__init__.py @@ -0,0 +1,8 @@ +import wandb + +try: + from wandb_workspaces.workspaces import * # noqa: F403 +except ImportError: + wandb.termerror( + "Failed to import wandb_workspaces. To edit workspaces programmatically, please install it using `pip install wandb[workspaces]`." + ) diff --git a/lib/python3.12/site-packages/wandb/apis/workspaces/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/apis/workspaces/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c53a50f659fc71f622eea42b281908b2cbd9f3a3 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/apis/workspaces/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/automations/__init__.py b/lib/python3.12/site-packages/wandb/automations/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7a95cdd8ed5027012d90e141b2343fab6de2f811 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/automations/__init__.py @@ -0,0 +1,73 @@ +import wandb +from wandb._pydantic import IS_PYDANTIC_V2 + +from .actions import ActionType, DoNothing, SendNotification, SendWebhook +from .automations import Automation, NewAutomation +from .events import ( + ArtifactEvent, + EventType, + MetricChangeFilter, + MetricThresholdFilter, + OnAddArtifactAlias, + OnCreateArtifact, + OnLinkArtifact, + OnRunMetric, + RunEvent, +) +from .integrations import Integration, SlackIntegration, WebhookIntegration +from .scopes import ArtifactCollectionScope, ProjectScope, ScopeType + +# ---------------------------------------------------------------------------- +# WARNINGS on import +if not IS_PYDANTIC_V2: + # Raises an error in Pydantic v1 environments, where the Automations API + # has not been tested and is unlikely to work as expected. + # + # Remove this when we either: + # - Drop support for Pydantic v1 + # - Are able to implement (limited) Pydantic v1 support + raise ImportError( + "The W&B Automations API requires Pydantic v2. " + "We recommend upgrading `pydantic` to use this feature." + ) + +else: + # If Pydantic v2 is available, we can use the full Automations API + # but communicate to users that the API is still experimental and + # may change rapidly. + wandb.termwarn( + "The W&B Automations API is experimental and the implementation is subject to change." + "Review the release notes before upgrading. We recommend pinning your " + f"package version to `{wandb.__package__}=={wandb.__version__}` to reduce the risk of disruption.", + repeat=False, + ) +# ---------------------------------------------------------------------------- + +__all__ = [ + # Scopes + "ScopeType", # doc:exclude + "ArtifactCollectionScope", # doc:exclude + "ProjectScope", # doc:exclude + # Events + "EventType", # doc:exclude + "OnAddArtifactAlias", + "OnCreateArtifact", + "OnLinkArtifact", + "OnRunMetric", + "ArtifactEvent", # doc:exclude + "RunEvent", # doc:exclude + "MetricThresholdFilter", + "MetricChangeFilter", + # Actions + "ActionType", # doc:exclude + "SendNotification", + "SendWebhook", + "DoNothing", + # Automations + "Automation", + "NewAutomation", + # Integrations + "Integration", # doc:exclude + "SlackIntegration", # doc:exclude + "WebhookIntegration", # doc:exclude +] diff --git a/lib/python3.12/site-packages/wandb/automations/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/automations/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..023744a78bbd05a05fbc2ba603b2c1067457f5da Binary files /dev/null and b/lib/python3.12/site-packages/wandb/automations/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/automations/__pycache__/_utils.cpython-312.pyc b/lib/python3.12/site-packages/wandb/automations/__pycache__/_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..54ebe04393db42bce53c4d085770921020df34a9 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/automations/__pycache__/_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/automations/__pycache__/_validators.cpython-312.pyc b/lib/python3.12/site-packages/wandb/automations/__pycache__/_validators.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8341ea5fc6de2a71417b4af7a8b140b4b1aee0ea Binary files /dev/null and b/lib/python3.12/site-packages/wandb/automations/__pycache__/_validators.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/automations/__pycache__/actions.cpython-312.pyc b/lib/python3.12/site-packages/wandb/automations/__pycache__/actions.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ddae88d5c71f7c27a2df2e6dfd254aa31ccfe736 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/automations/__pycache__/actions.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/automations/__pycache__/automations.cpython-312.pyc b/lib/python3.12/site-packages/wandb/automations/__pycache__/automations.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c13eb3f4404435eb3ed5704e73838f7466cb9b67 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/automations/__pycache__/automations.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/automations/__pycache__/events.cpython-312.pyc b/lib/python3.12/site-packages/wandb/automations/__pycache__/events.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1c7813db0a5ea88e3fc1ce75ffefd73b5278ca07 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/automations/__pycache__/events.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/automations/__pycache__/integrations.cpython-312.pyc b/lib/python3.12/site-packages/wandb/automations/__pycache__/integrations.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e6ff21c6ac0f853352c48df0ddee66af7178e165 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/automations/__pycache__/integrations.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/automations/__pycache__/scopes.cpython-312.pyc b/lib/python3.12/site-packages/wandb/automations/__pycache__/scopes.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a09c768b5703ace13aa4578624f35af9b2afdcff Binary files /dev/null and b/lib/python3.12/site-packages/wandb/automations/__pycache__/scopes.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/automations/_filters/__init__.py b/lib/python3.12/site-packages/wandb/automations/_filters/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..56d551d4ed507da432151e8eb90167495fdb5acf --- /dev/null +++ b/lib/python3.12/site-packages/wandb/automations/_filters/__init__.py @@ -0,0 +1,40 @@ +from .expressions import FilterExpr, MongoLikeFilter +from .operators import ( + And, + Contains, + Eq, + Exists, + Gt, + Gte, + In, + Lt, + Lte, + Ne, + Nor, + Not, + NotIn, + Op, + Or, + Regex, +) + +__all__ = [ + "And", + "Or", + "Nor", + "Not", + "Op", + "Gt", + "Lt", + "Gte", + "Lte", + "Eq", + "Ne", + "In", + "NotIn", + "Contains", + "Exists", + "Regex", + "FilterExpr", + "MongoLikeFilter", +] diff --git a/lib/python3.12/site-packages/wandb/automations/_filters/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/automations/_filters/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fdb49c3267f67d5785c6c2fc0cb697c2225129bb Binary files /dev/null and b/lib/python3.12/site-packages/wandb/automations/_filters/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/automations/_filters/__pycache__/expressions.cpython-312.pyc b/lib/python3.12/site-packages/wandb/automations/_filters/__pycache__/expressions.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cfe55f2eb417caf14db73b88e9ca3414f16229d3 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/automations/_filters/__pycache__/expressions.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/automations/_filters/__pycache__/operators.cpython-312.pyc b/lib/python3.12/site-packages/wandb/automations/_filters/__pycache__/operators.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8c814bac72a03e03f1ae1d64d73517b9d3af6ff6 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/automations/_filters/__pycache__/operators.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/automations/_filters/__pycache__/run_metrics.cpython-312.pyc b/lib/python3.12/site-packages/wandb/automations/_filters/__pycache__/run_metrics.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5d17af15fe950708adff3ca397290c91a22c9f5d Binary files /dev/null and b/lib/python3.12/site-packages/wandb/automations/_filters/__pycache__/run_metrics.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/automations/_filters/expressions.py b/lib/python3.12/site-packages/wandb/automations/_filters/expressions.py new file mode 100644 index 0000000000000000000000000000000000000000..689cb7acf8c7b07945b0846d90681019bb32731a --- /dev/null +++ b/lib/python3.12/site-packages/wandb/automations/_filters/expressions.py @@ -0,0 +1,181 @@ +"""Pydantic-compatible representations of MongoDB expressions.""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any, Union + +from pydantic import ConfigDict, model_serializer +from typing_extensions import Self, TypeAlias, get_args + +from wandb._pydantic import CompatBaseModel, model_validator + +from .operators import ( + Contains, + Eq, + Exists, + Gt, + Gte, + In, + Lt, + Lte, + Ne, + NotIn, + Op, + Regex, + RichReprResult, + Scalar, + ScalarTypes, + SupportsLogicalOpSyntax, +) + + +class FilterableField: + """A descriptor that can be used to define a "filterable" field on a class. + + Internal helper to support syntactic sugar for defining event filters. + """ + + _python_name: str #: The name of the field this descriptor was assigned to in the Python class. + _server_name: str | None #: If set, the actual server-side field name to filter on. + + def __init__(self, server_name: str | None = None): + self._server_name = server_name + + def __set_name__(self, owner: type, name: str) -> None: + self._python_name = name + + def __get__(self, obj: Any, objtype: type) -> Self: + # By default, if we didn't explicitly provide a backend name for + # filtering, assume the field has the same name in the backend as + # the python attribute. + return self + + @property + def _name(self) -> str: + return self._server_name or self._python_name + + def __str__(self) -> str: + return self._name + + def __repr__(self) -> str: + return f"{type(self).__name__}({self._name!r})" + + # Methods to define filter expressions through chaining + def matches_regex(self, pattern: str) -> FilterExpr: + return FilterExpr(field=self._name, op=Regex(regex_=pattern)) + + def contains(self, text: str) -> FilterExpr: + return FilterExpr(field=self._name, op=Contains(contains_=text)) + + def exists(self, exists: bool = True) -> FilterExpr: + return FilterExpr(field=self._name, op=Exists(exists_=exists)) + + def lt(self, value: Scalar) -> FilterExpr: + return FilterExpr(field=self._name, op=Lt(lt_=value)) + + def gt(self, value: Scalar) -> FilterExpr: + return FilterExpr(field=self._name, op=Gt(gt_=value)) + + def lte(self, value: Scalar) -> FilterExpr: + return FilterExpr(field=self._name, op=Lte(lte_=value)) + + def gte(self, value: Scalar) -> FilterExpr: + return FilterExpr(field=self._name, op=Gte(gte_=value)) + + def ne(self, value: Scalar) -> FilterExpr: + return FilterExpr(field=self._name, op=Ne(ne_=value)) + + def eq(self, value: Scalar) -> FilterExpr: + return FilterExpr(field=self._name, op=Eq(eq_=value)) + + def in_(self, values: Iterable[Scalar]) -> FilterExpr: + return FilterExpr(field=self._name, op=In(in_=values)) + + def not_in(self, values: Iterable[Scalar]) -> FilterExpr: + return FilterExpr(field=self._name, op=NotIn(nin_=values)) + + # Override the default behavior of comparison operators: <, >=, ==, etc + def __lt__(self, other: Any) -> FilterExpr: + if isinstance(other, ScalarTypes): + return self.lt(other) # type: ignore[arg-type] + raise TypeError(f"Invalid operand type in filter expression: {type(other)!r}") + + def __gt__(self, other: Any) -> FilterExpr: + if isinstance(other, ScalarTypes): + return self.gt(other) # type: ignore[arg-type] + raise TypeError(f"Invalid operand type in filter expression: {type(other)!r}") + + def __le__(self, other: Any) -> FilterExpr: + if isinstance(other, ScalarTypes): + return self.lte(other) # type: ignore[arg-type] + raise TypeError(f"Invalid operand type in filter expression: {type(other)!r}") + + def __ge__(self, other: Any) -> FilterExpr: + if isinstance(other, ScalarTypes): + return self.gte(other) # type: ignore[arg-type] + raise TypeError(f"Invalid operand type in filter expression: {type(other)!r}") + + # Operator behavior is intentionally overridden to allow defining + # filter expressions like `field == "value"`. See similar overrides + # of built-in dunder methods in sqlalchemy, polars, pandas, numpy, etc. + # + # sqlalchemy example for illustrative purposes: + # https://github.com/sqlalchemy/sqlalchemy/blob/f21ae633486380a26dc0b67b70ae1c0efc6b4dc4/lib/sqlalchemy/orm/descriptor_props.py#L808-L812 + def __eq__(self, other: Any) -> FilterExpr: + if isinstance(other, ScalarTypes): + return self.eq(other) # type: ignore[arg-type] + raise TypeError(f"Invalid operand type in filter expression: {type(other)!r}") + + def __ne__(self, other: Any) -> FilterExpr: + if isinstance(other, ScalarTypes): + return self.ne(other) # type: ignore[arg-type] + raise TypeError(f"Invalid operand type in filter expression: {type(other)!r}") + + +# ------------------------------------------------------------------------------ +class FilterExpr(CompatBaseModel, SupportsLogicalOpSyntax): + """A MongoDB filter expression on a specific field.""" + + model_config = ConfigDict( + arbitrary_types_allowed=True, + ) + + field: str + op: Op + + def __repr__(self) -> str: + return f"{type(self).__name__}({self.field!s}: {self.op!r})" + + def __rich_repr__(self) -> RichReprResult: + # https://rich.readthedocs.io/en/stable/pretty.html + yield self.field, self.op + + @model_validator(mode="before") + @classmethod + def _validate(cls, data: Any) -> Any: + """Parse a MongoDB dict representation of the filter expression.""" + if ( + isinstance(data, dict) + and len(data) == 1 + and not any(key.startswith("$") for key in data) + ): + # This looks like a MongoDB filter dict. E.g.: + # - in: `{"display_name": {"$contains": "my-run"}}` + # - out: `FilterExpr(field="display_name", op=Contains(contains_="my-run"))` + ((field, op),) = data.items() + return {"field": field, "op": op} + return data + + @model_serializer(mode="plain") + def _serialize(self) -> dict[str, Any]: + """Return a MongoDB dict representation of the expression.""" + from pydantic_core import to_jsonable_python # Only valid in pydantic v2 + + return {self.field: to_jsonable_python(self.op, by_alias=True, round_trip=True)} + + +# for type annotations +MongoLikeFilter: TypeAlias = Union[Op, FilterExpr] +# for runtime type checks +MongoLikeFilterTypes: tuple[type, ...] = get_args(MongoLikeFilter) diff --git a/lib/python3.12/site-packages/wandb/automations/_filters/operators.py b/lib/python3.12/site-packages/wandb/automations/_filters/operators.py new file mode 100644 index 0000000000000000000000000000000000000000..1c2543421a6fa0a9f3d5e0afff283f04f69ae169 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/automations/_filters/operators.py @@ -0,0 +1,258 @@ +"""Types that represent operators in MongoDB filter expressions.""" + +from __future__ import annotations + +from typing import Any, Dict, Iterable, Tuple, TypeVar, Union + +from pydantic import ConfigDict, Field, StrictBool, StrictFloat, StrictInt, StrictStr +from typing_extensions import TypeAlias, get_args + +from wandb._pydantic import GQLBase + +# for type annotations +Scalar = Union[StrictStr, StrictInt, StrictFloat, StrictBool] +# for runtime type checks +ScalarTypes: tuple[type, ...] = tuple(t.__origin__ for t in get_args(Scalar)) + +# See: https://rich.readthedocs.io/en/stable/pretty.html#rich-repr-protocol +RichReprResult: TypeAlias = Iterable[ + Union[ + Any, + Tuple[Any], + Tuple[str, Any], + Tuple[str, Any, Any], + ] +] + +T = TypeVar("T") +TupleOf: TypeAlias = Tuple[T, ...] + + +# NOTE: Wherever class descriptions that are not docstrings, this is deliberate. +# This is done to ensure the descriptions are omitted from generated API docs. + + +# Mixin to support syntactic sugar for MongoDB expressions with (bitwise) logical operators, +# e.g. `a | b` -> `{"$or": [a, b]}` or `~a` -> `{"$not": a}`. +class SupportsLogicalOpSyntax: + def __or__(self, other: Any) -> Or: + """Syntactic sugar for: `a | b` -> `Or(a, b)`.""" + return Or(or_=[self, other]) + + def __and__(self, other: Any) -> And: + """Syntactic sugar for: `a & b` -> `And(a, b)`.""" + from .expressions import FilterExpr + + if isinstance(other, (BaseOp, FilterExpr)): + return And(and_=[self, other]) + return NotImplemented + + def __invert__(self) -> Not: + """Syntactic sugar for: `~a` -> `Not(a)`.""" + return Not(not_=self) + + +# Base class for parsed MongoDB filter/query operators, e.g. `{"$and": [...]}`. +class BaseOp(GQLBase, SupportsLogicalOpSyntax): + model_config = ConfigDict( + extra="forbid", + frozen=True, # Make pseudo-immutable for easier comparison and hashing + ) + + def __repr__(self) -> str: + # Display operand as a positional arg + values_repr = ", ".join(map(repr, self.model_dump().values())) + return f"{type(self).__name__}({values_repr})" + + def __rich_repr__(self) -> RichReprResult: + # Display field values as positional args: + # https://rich.readthedocs.io/en/stable/pretty.html + yield from ((None, v) for v in self.model_dump().values()) + + +# Logical operator(s) +# https://www.mongodb.com/docs/manual/reference/operator/query/and/ +# https://www.mongodb.com/docs/manual/reference/operator/query/or/ +# https://www.mongodb.com/docs/manual/reference/operator/query/nor/ +# https://www.mongodb.com/docs/manual/reference/operator/query/not/ +class And(BaseOp): + and_: TupleOf[Any] = Field(default=(), alias="$and") + + +class Or(BaseOp): + or_: TupleOf[Any] = Field(default=(), alias="$or") + + def __invert__(self) -> Nor: + """Syntactic sugar for: `~Or(a, b)` -> `Nor(a, b)`.""" + return Nor(nor_=self.or_) + + +class Nor(BaseOp): + nor_: TupleOf[Any] = Field(default=(), alias="$nor") + + def __invert__(self) -> Or: + """Syntactic sugar for: `~Nor(a, b)` -> `Or(a, b)`.""" + return Or(or_=self.nor_) + + +class Not(BaseOp): + not_: Any = Field(alias="$not") + + def __invert__(self) -> Any: + """Syntactic sugar for: `~Not(a)` -> `a`.""" + return self.not_ + + +# Comparison operator(s) +# https://www.mongodb.com/docs/manual/reference/operator/query/lt/ +# https://www.mongodb.com/docs/manual/reference/operator/query/gt/ +# https://www.mongodb.com/docs/manual/reference/operator/query/lte/ +# https://www.mongodb.com/docs/manual/reference/operator/query/gte/ +# https://www.mongodb.com/docs/manual/reference/operator/query/eq/ +# https://www.mongodb.com/docs/manual/reference/operator/query/ne/ +# https://www.mongodb.com/docs/manual/reference/operator/query/in/ +# https://www.mongodb.com/docs/manual/reference/operator/query/nin/ +class Lt(BaseOp): + lt_: Scalar = Field(alias="$lt") + + def __invert__(self) -> Gte: + """Syntactic sugar for: `~Lt(a)` -> `Gte(a)`.""" + return Gte(gte_=self.lt_) + + +class Gt(BaseOp): + gt_: Scalar = Field(alias="$gt") + + def __invert__(self) -> Lte: + """Syntactic sugar for: `~Gt(a)` -> `Lte(a)`.""" + return Lte(lte_=self.gt_) + + +class Lte(BaseOp): + lte_: Scalar = Field(alias="$lte") + + def __invert__(self) -> Gt: + """Syntactic sugar for: `~Lte(a)` -> `Gt(a)`.""" + return Gt(gt_=self.lte_) + + +class Gte(BaseOp): + gte_: Scalar = Field(alias="$gte") + + def __invert__(self) -> Lt: + """Syntactic sugar for: `~Gte(a)` -> `Lt(a)`.""" + return Lt(lt_=self.gte_) + + +class Eq(BaseOp): + eq_: Scalar = Field(alias="$eq") + + def __invert__(self) -> Ne: + """Syntactic sugar for: `~Eq(a)` -> `Ne(a)`.""" + return Ne(ne_=self.eq_) + + +class Ne(BaseOp): + ne_: Scalar = Field(alias="$ne") + + def __invert__(self) -> Eq: + """Syntactic sugar for: `~Ne(a)` -> `Eq(a)`.""" + return Eq(eq_=self.ne_) + + +class In(BaseOp): + in_: TupleOf[Scalar] = Field(default=(), alias="$in") + + def __invert__(self) -> NotIn: + """Syntactic sugar for: `~In(a)` -> `NotIn(a)`.""" + return NotIn(nin_=self.in_) + + +class NotIn(BaseOp): + nin_: TupleOf[Scalar] = Field(default=(), alias="$nin") + + def __invert__(self) -> In: + """Syntactic sugar for: `~NotIn(a)` -> `In(a)`.""" + return In(in_=self.nin_) + + +# Element operator(s) +# https://www.mongodb.com/docs/manual/reference/operator/query/exists/ +class Exists(BaseOp): + exists_: bool = Field(alias="$exists") + + +# Evaluation operator(s) +# https://www.mongodb.com/docs/manual/reference/operator/query/regex/ +# +# Note: "$contains" is NOT a formal MongoDB operator, but the backend recognizes and +# executes it as a substring-match filter. +class Regex(BaseOp): + regex_: str = Field(alias="$regex") #: The regex expression to match against. + + +class Contains(BaseOp): + contains_: str = Field(alias="$contains") #: The substring to match against. + + +And.model_rebuild() +Or.model_rebuild() +Not.model_rebuild() +Lt.model_rebuild() +Gt.model_rebuild() +Lte.model_rebuild() +Gte.model_rebuild() +Eq.model_rebuild() +Ne.model_rebuild() +In.model_rebuild() +NotIn.model_rebuild() +Exists.model_rebuild() +Regex.model_rebuild() +Contains.model_rebuild() + + +# ------------------------------------------------------------------------------ +# Convenience helpers, constants, and utils for supported MongoDB operators +# ------------------------------------------------------------------------------ +KEY_TO_OP: dict[str, type[BaseOp]] = { + "$and": And, + "$or": Or, + "$nor": Nor, + "$not": Not, + "$lt": Lt, + "$gt": Gt, + "$lte": Lte, + "$gte": Gte, + "$eq": Eq, + "$ne": Ne, + "$in": In, + "$nin": NotIn, + "$exists": Exists, + "$regex": Regex, + "$contains": Contains, +} + + +KnownOp = Union[ + And, + Or, + Nor, + Not, + Lt, + Gt, + Lte, + Gte, + Eq, + Ne, + In, + NotIn, + Exists, + Regex, + Contains, +] +UnknownOp = Dict[str, Any] + +# for type annotations +Op = Union[KnownOp, UnknownOp] +# for runtime type checks +OpTypes: tuple[type, ...] = (*get_args(KnownOp), dict) diff --git a/lib/python3.12/site-packages/wandb/automations/_filters/run_metrics.py b/lib/python3.12/site-packages/wandb/automations/_filters/run_metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..4ddeb46bb7bf418138107e03a9c39051c471650c --- /dev/null +++ b/lib/python3.12/site-packages/wandb/automations/_filters/run_metrics.py @@ -0,0 +1,330 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, overload + +from pydantic import ( + Field, + PositiveFloat, + PositiveInt, + StrictFloat, + StrictInt, + field_validator, +) +from typing_extensions import Annotated, TypeAlias, override + +from wandb._pydantic import GQLBase +from wandb.automations._validators import LenientStrEnum + +from .expressions import FilterExpr +from .operators import BaseOp, RichReprResult + +if TYPE_CHECKING: + from wandb.automations.events import RunMetricFilter + +# Maps MongoDB comparison operators -> Python literal (str) representations +MONGO2PY_OPS: Final[dict[str, str]] = { + "$eq": "==", + "$ne": "!=", + "$gt": ">", + "$lt": "<", + "$gte": ">=", + "$lte": "<=", +} +# Reverse mapping from Python literal (str) -> MongoDB operator key +PY2MONGO_OPS: Final[dict[str, str]] = {v: k for k, v in MONGO2PY_OPS.items()} + +# Type hint for positive numbers (int or float) +PosNum: TypeAlias = Union[PositiveInt, PositiveFloat] + + +class Agg(LenientStrEnum): # from: Aggregation + """Supported run metric aggregation operations.""" + + MAX = "MAX" + MIN = "MIN" + AVERAGE = "AVERAGE" + + # Shorter aliases for convenience + AVG = AVERAGE + + +class ChangeType(LenientStrEnum): # from: RunMetricChangeType + """Describes the metric change as absolute (arithmetic difference) or relative (decimal percentage).""" + + ABSOLUTE = "ABSOLUTE" + RELATIVE = "RELATIVE" + + # Shorter aliases for convenience + ABS = ABSOLUTE + REL = RELATIVE + + +class ChangeDir(LenientStrEnum): # from: RunMetricChangeDirection + """Describes the direction of the metric change.""" + + INCREASE = "INCREASE" + DECREASE = "DECREASE" + ANY = "ANY" + + # Shorter aliases for convenience + INC = INCREASE + DEC = DECREASE + + +class BaseMetricFilter(GQLBase, ABC, extra="forbid"): + name: str + """Name of the observed metric.""" + + agg: Optional[Agg] + """Aggregate operation, if any, to apply over the window size.""" + + window: PositiveInt + """Size of the window over which the metric is aggregated (ignored if `agg is None`).""" + + # ------------------------------------------------------------------------------ + cmp: Optional[str] + """Comparison between the metric expression (left) vs. the threshold or target value (right).""" + + # ------------------------------------------------------------------------------ + threshold: Union[StrictInt, StrictFloat] + """Threshold value to compare against.""" + + def __and__(self, other: Any) -> RunMetricFilter: + """Implements `(metric_filter & run_filter) -> RunMetricFilter`.""" + from wandb.automations.events import RunMetricFilter + + if isinstance(run_filter := other, (BaseOp, FilterExpr)): + # Assume `other` is a run filter, and we are building a RunMetricEvent. + # For the metric filter, delegate to the inner validator(s) to further wrap/nest as appropriate. + return RunMetricFilter(run=run_filter, metric=self) + return NotImplemented + + def __rand__(self, other: BaseOp | FilterExpr) -> RunMetricFilter: + """Ensures `&` is commutative: `(run_filter & metric_filter) == (metric_filter & run_filter)`.""" + return self.__and__(other) + + @abstractmethod + def __repr__(self) -> str: + """The text representation of the metric filter.""" + raise NotImplementedError + + @override + def __rich_repr__(self) -> RichReprResult: + """The representation of the metric filter when using `rich` for pretty-printing.""" + # See: https://rich.readthedocs.io/en/stable/pretty.html#rich-repr-protocol + yield None, repr(self) + + +class MetricThresholdFilter(BaseMetricFilter): # from: RunMetricThresholdFilter + """Defines a filter that compares a run metric against a user-defined threshold value.""" + + name: str + agg: Annotated[Optional[Agg], Field(alias="agg_op")] = None + window: Annotated[PositiveInt, Field(alias="window_size")] = 1 + + cmp: Annotated[Literal["$gte", "$gt", "$lt", "$lte"], Field(alias="cmp_op")] + """Comparison operator used to compare the metric value (left) vs. the threshold value (right).""" + + threshold: Union[StrictInt, StrictFloat] + + @field_validator("cmp", mode="before") + def _validate_cmp(cls, v: Any) -> Any: + # Be helpful: e.g. ">" -> "$gt" + return PY2MONGO_OPS.get(v.strip(), v) if isinstance(v, str) else v + + def __repr__(self) -> str: + metric = f"{self.agg.value}({self.name})" if self.agg else self.name + op = MONGO2PY_OPS.get(self.cmp, self.cmp) + return repr(rf"{metric} {op} {self.threshold}") + + +class MetricChangeFilter(BaseMetricFilter): # from: RunMetricChangeFilter + """Defines a filter that compares a change in a run metric against a user-defined threshold. + + The change is calculated over "tumbling" windows, i.e. the difference + between the current window and the non-overlapping prior window. + """ + + name: str + agg: Annotated[Optional[Agg], Field(alias="agg_op")] = None + window: Annotated[PositiveInt, Field(alias="current_window_size")] = 1 + + # `prior_window` is only for `RUN_METRIC_CHANGE` events + prior_window: Annotated[ + PositiveInt, + # By default, set `window -> prior_window` if the latter wasn't provided. + Field(alias="prior_window_size", default_factory=lambda data: data["window"]), + ] + """Size of the prior window over which the metric is aggregated (ignored if `agg is None`). + + If omitted, defaults to the size of the current window. + """ + + # ------------------------------------------------------------------------------ + # NOTE: + # - The "comparison" operator isn't actually part of the backend schema, + # but it's defined here for consistency -- and ignored otherwise. + # - In the backend, it's effectively "$gte" or "$lte", depending on the sign + # (change_dir), though again, this is not explicit in the schema. + cmp: Annotated[None, Field(frozen=True, exclude=True, repr=False)] = None + """Ignored.""" + + # ------------------------------------------------------------------------------ + change_type: Annotated[ChangeType, Field(alias="change_type")] + change_dir: Annotated[ChangeDir, Field(alias="change_dir")] + threshold: Annotated[PosNum, Field(alias="change_amount")] + + def __repr__(self) -> str: + metric = f"{self.agg.value}({self.name})" if self.agg else self.name + verb = ( + "changes" + if (self.change_dir is ChangeDir.ANY) + else f"{self.change_dir.value.lower()}s" + ) + + fmt_spec = ".2%" if (self.change_type is ChangeType.REL) else "" + amt = f"{self.threshold:{fmt_spec}}" + return repr(rf"{metric} {verb} {amt}") + + +class BaseMetricOperand(GQLBase, extra="forbid"): + def gt(self, value: int | float, /) -> MetricThresholdFilter: + """Defines a `MetricThresholdFilter` that observes for `metric_expr > threshold`.""" + return self > value + + def lt(self, value: int | float, /) -> MetricThresholdFilter: + """Defines a `MetricThresholdFilter` that observes for `metric_expr < threshold`.""" + return self < value + + def gte(self, value: int | float, /) -> MetricThresholdFilter: + """Defines a `MetricThresholdFilter` that observes for `metric_expr >= threshold`.""" + return self >= value + + def lte(self, value: int | float, /) -> MetricThresholdFilter: + """Defines a `MetricThresholdFilter` that observes for `metric_expr <= threshold`.""" + return self <= value + + # Overloads to implement: + # - `(metric_operand > threshold) -> MetricThresholdFilter` + # - `(metric_operand < threshold) -> MetricThresholdFilter` + # - `(metric_operand >= threshold) -> MetricThresholdFilter` + # - `(metric_operand <= threshold) -> MetricThresholdFilter` + def __gt__(self, other: Any) -> MetricThresholdFilter: + if isinstance(other, (int, float)): + return MetricThresholdFilter(**dict(self), cmp="$gt", threshold=other) + return NotImplemented + + def __lt__(self, other: Any) -> MetricThresholdFilter: + if isinstance(other, (int, float)): + return MetricThresholdFilter(**dict(self), cmp="$lt", threshold=other) + return NotImplemented + + def __ge__(self, other: Any) -> MetricThresholdFilter: + if isinstance(other, (int, float)): + return MetricThresholdFilter(**dict(self), cmp="$gte", threshold=other) + return NotImplemented + + def __le__(self, other: Any) -> MetricThresholdFilter: + if isinstance(other, (int, float)): + return MetricThresholdFilter(**dict(self), cmp="$lte", threshold=other) + return NotImplemented + + @overload + def changes_by(self, *, diff: PosNum, frac: None) -> MetricChangeFilter: ... + @overload + def changes_by(self, *, diff: None, frac: PosNum) -> MetricChangeFilter: ... + @overload # NOTE: This overload is for internal use only. + def changes_by( + self, *, diff: PosNum | None, frac: PosNum | None, _dir: ChangeDir + ) -> MetricChangeFilter: ... + def changes_by( + self, + *, + diff: PosNum | None = None, + frac: PosNum | None = None, + _dir: ChangeDir = ChangeDir.ANY, + ) -> MetricChangeFilter: + """Defines a filter that observes for any change (increase OR decrease) in a run metric. + + Exactly one of the keyword arguments `frac` or `diff` must be provided. + + Args: + diff: + If given, the arithmetic difference that must be observed + in the metric. Must be a positive number. + frac: + If given, the fractional (relative) change that must be observed + in the metric. Must be a positive number. E.g. `frac=0.1` + denotes a 10% relative increase OR decrease. + """ + # Enforce mutually exclusive keyword args + if (frac is None) is (diff is None): + raise ValueError("Must provide exactly one of `frac` or `diff`") + + # Enforce positive values + if (frac is not None) and (frac <= 0): + raise ValueError(f"Expected positive quantity, got: {frac=}") + if (diff is not None) and (diff <= 0): + raise ValueError(f"Expected positive quantity, got: {diff=}") + + if diff is None: + change_kws = dict(change_type=ChangeType.REL, threshold=frac) + return MetricChangeFilter(**dict(self), change_dir=_dir, **change_kws) + else: + change_kws = dict(change_type=ChangeType.ABS, threshold=diff) + return MetricChangeFilter(**dict(self), change_dir=_dir, **change_kws) + + @overload + def increases_by(self, *, diff: PosNum, frac: None) -> MetricChangeFilter: ... + @overload + def increases_by(self, *, diff: None, frac: PosNum) -> MetricChangeFilter: ... + def increases_by( + self, *, diff: PosNum | None = None, frac: PosNum | None = None + ) -> MetricChangeFilter: + """Defines a filter that observes for an increase in the numerical value of a run metric. + + Arguments are the same as for `.changes_by()`. + """ + return self.changes_by(diff=diff, frac=frac, _dir=ChangeDir.INC) + + @overload + def decreases_by(self, *, diff: PosNum, frac: None) -> MetricChangeFilter: ... + @overload + def decreases_by(self, *, diff: None, frac: PosNum) -> MetricChangeFilter: ... + def decreases_by( + self, *, diff: PosNum | None = None, frac: PosNum | None = None + ) -> MetricChangeFilter: + """Defines a filter that observes for a decrease in the numerical value of a run metric. + + Arguments are the same as for `.changes_by()`. + """ + return self.changes_by(diff=diff, frac=frac, _dir=ChangeDir.DEC) + + +class MetricVal(BaseMetricOperand): + """Represents a single, unaggregated metric value when defining a metric filter.""" + + name: str + + # Allow users to convert this single-value metric into an aggregated metric expression. + def max(self, window: int) -> MetricAgg: + return MetricAgg(name=self.name, agg=Agg.MAX, window=window) + + def min(self, window: int) -> MetricAgg: + return MetricAgg(name=self.name, agg=Agg.MIN, window=window) + + def avg(self, window: int) -> MetricAgg: + return MetricAgg(name=self.name, agg=Agg.AVG, window=window) + + # Aliased method for users familiar with e.g. torch/tf/numpy/pandas/polars/etc. + def mean(self, window: int) -> MetricAgg: + return self.avg(window=window) + + +class MetricAgg(BaseMetricOperand): + """Represents an aggregated metric value when defining a metric filter.""" + + name: str + agg: Annotated[Agg, Field(alias="agg_op")] + window: Annotated[PositiveInt, Field(alias="window_size")] diff --git a/lib/python3.12/site-packages/wandb/automations/_generated/__init__.py b/lib/python3.12/site-packages/wandb/automations/_generated/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..717c170e4a34333ee08c27b14b7ced31b91bf4ff --- /dev/null +++ b/lib/python3.12/site-packages/wandb/automations/_generated/__init__.py @@ -0,0 +1,177 @@ +# Generated by ariadne-codegen + +from .create_automation import CreateAutomation +from .create_generic_webhook_integration import ( + CreateGenericWebhookIntegration, + CreateGenericWebhookIntegrationCreateGenericWebhookIntegration, + CreateGenericWebhookIntegrationCreateGenericWebhookIntegrationIntegrationGenericWebhookIntegration, + CreateGenericWebhookIntegrationCreateGenericWebhookIntegrationIntegrationIntegration, +) +from .delete_automation import DeleteAutomation +from .enums import ( + AlertSeverity, + EventTriggeringConditionType, + TriggeredActionType, + TriggerScopeType, +) +from .fragments import ( + ArtifactPortfolioScopeFields, + ArtifactSequenceScopeFields, + FilterEventFields, + GenericWebhookActionFields, + GenericWebhookActionFieldsIntegrationGenericWebhookIntegration, + GenericWebhookActionFieldsIntegrationIntegration, + GenericWebhookIntegrationConnectionFields, + GenericWebhookIntegrationConnectionFieldsEdges, + GenericWebhookIntegrationConnectionFieldsEdgesNodeGenericWebhookIntegration, + GenericWebhookIntegrationConnectionFieldsEdgesNodeIntegration, + GenericWebhookIntegrationFields, + IntegrationConnectionFields, + IntegrationConnectionFieldsEdges, + IntegrationConnectionFieldsEdgesNodeGenericWebhookIntegration, + IntegrationConnectionFieldsEdgesNodeIntegration, + IntegrationConnectionFieldsEdgesNodeSlackIntegration, + NoOpActionFields, + NotificationActionFields, + NotificationActionFieldsIntegrationIntegration, + NotificationActionFieldsIntegrationSlackIntegration, + PageInfoFields, + ProjectConnectionFields, + ProjectConnectionFieldsEdges, + ProjectConnectionFieldsEdgesNode, + ProjectScopeFields, + QueueJobActionFields, + QueueJobActionFieldsQueue, + SlackIntegrationConnectionFields, + SlackIntegrationConnectionFieldsEdges, + SlackIntegrationConnectionFieldsEdgesNodeIntegration, + SlackIntegrationConnectionFieldsEdgesNodeSlackIntegration, + SlackIntegrationFields, + TriggerFields, + TriggerFieldsActionGenericWebhookTriggeredAction, + TriggerFieldsActionNoOpTriggeredAction, + TriggerFieldsActionNotificationTriggeredAction, + TriggerFieldsActionQueueJobTriggeredAction, + TriggerFieldsEventFilterEventTriggeringCondition, + TriggerFieldsScopeArtifactPortfolio, + TriggerFieldsScopeArtifactSequence, + TriggerFieldsScopeProject, +) +from .generic_webhook_integrations_by_entity import ( + GenericWebhookIntegrationsByEntity, + GenericWebhookIntegrationsByEntityEntity, +) +from .get_automations import GetAutomations, GetAutomationsSearchScope +from .get_automations_by_entity import ( + GetAutomationsByEntity, + GetAutomationsByEntitySearchScope, +) +from .input_types import ( + CreateFilterTriggerInput, + CreateGenericWebhookIntegrationInput, + GenericWebhookActionInput, + NoOpTriggeredActionInput, + NotificationActionInput, + QueueJobActionInput, + TriggeredActionConfig, + UpdateFilterTriggerInput, +) +from .integrations_by_entity import IntegrationsByEntity, IntegrationsByEntityEntity +from .operations import ( + CREATE_AUTOMATION_GQL, + CREATE_GENERIC_WEBHOOK_INTEGRATION_GQL, + DELETE_AUTOMATION_GQL, + GENERIC_WEBHOOK_INTEGRATIONS_BY_ENTITY_GQL, + GET_AUTOMATIONS_BY_ENTITY_GQL, + GET_AUTOMATIONS_GQL, + INTEGRATIONS_BY_ENTITY_GQL, + SLACK_INTEGRATIONS_BY_ENTITY_GQL, + UPDATE_AUTOMATION_GQL, +) +from .slack_integrations_by_entity import ( + SlackIntegrationsByEntity, + SlackIntegrationsByEntityEntity, +) +from .update_automation import UpdateAutomation + +__all__ = [ + "CREATE_AUTOMATION_GQL", + "CREATE_GENERIC_WEBHOOK_INTEGRATION_GQL", + "DELETE_AUTOMATION_GQL", + "GENERIC_WEBHOOK_INTEGRATIONS_BY_ENTITY_GQL", + "GET_AUTOMATIONS_BY_ENTITY_GQL", + "GET_AUTOMATIONS_GQL", + "INTEGRATIONS_BY_ENTITY_GQL", + "SLACK_INTEGRATIONS_BY_ENTITY_GQL", + "UPDATE_AUTOMATION_GQL", + "GetAutomations", + "GetAutomationsSearchScope", + "GetAutomationsByEntity", + "GetAutomationsByEntitySearchScope", + "CreateAutomation", + "UpdateAutomation", + "DeleteAutomation", + "IntegrationsByEntity", + "IntegrationsByEntityEntity", + "SlackIntegrationsByEntity", + "SlackIntegrationsByEntityEntity", + "GenericWebhookIntegrationsByEntity", + "GenericWebhookIntegrationsByEntityEntity", + "CreateGenericWebhookIntegration", + "CreateGenericWebhookIntegrationCreateGenericWebhookIntegration", + "CreateGenericWebhookIntegrationCreateGenericWebhookIntegrationIntegrationGenericWebhookIntegration", + "CreateGenericWebhookIntegrationCreateGenericWebhookIntegrationIntegrationIntegration", + "CreateFilterTriggerInput", + "CreateGenericWebhookIntegrationInput", + "GenericWebhookActionInput", + "NoOpTriggeredActionInput", + "NotificationActionInput", + "QueueJobActionInput", + "TriggeredActionConfig", + "UpdateFilterTriggerInput", + "ArtifactPortfolioScopeFields", + "ArtifactSequenceScopeFields", + "FilterEventFields", + "GenericWebhookActionFields", + "GenericWebhookActionFieldsIntegrationGenericWebhookIntegration", + "GenericWebhookActionFieldsIntegrationIntegration", + "GenericWebhookIntegrationConnectionFields", + "GenericWebhookIntegrationConnectionFieldsEdges", + "GenericWebhookIntegrationConnectionFieldsEdgesNodeGenericWebhookIntegration", + "GenericWebhookIntegrationConnectionFieldsEdgesNodeIntegration", + "GenericWebhookIntegrationFields", + "IntegrationConnectionFields", + "IntegrationConnectionFieldsEdges", + "IntegrationConnectionFieldsEdgesNodeGenericWebhookIntegration", + "IntegrationConnectionFieldsEdgesNodeIntegration", + "IntegrationConnectionFieldsEdgesNodeSlackIntegration", + "NoOpActionFields", + "NotificationActionFields", + "NotificationActionFieldsIntegrationIntegration", + "NotificationActionFieldsIntegrationSlackIntegration", + "PageInfoFields", + "ProjectConnectionFields", + "ProjectConnectionFieldsEdges", + "ProjectConnectionFieldsEdgesNode", + "ProjectScopeFields", + "QueueJobActionFields", + "QueueJobActionFieldsQueue", + "SlackIntegrationConnectionFields", + "SlackIntegrationConnectionFieldsEdges", + "SlackIntegrationConnectionFieldsEdgesNodeIntegration", + "SlackIntegrationConnectionFieldsEdgesNodeSlackIntegration", + "SlackIntegrationFields", + "TriggerFields", + "TriggerFieldsActionGenericWebhookTriggeredAction", + "TriggerFieldsActionNoOpTriggeredAction", + "TriggerFieldsActionNotificationTriggeredAction", + "TriggerFieldsActionQueueJobTriggeredAction", + "TriggerFieldsEventFilterEventTriggeringCondition", + "TriggerFieldsScopeArtifactPortfolio", + "TriggerFieldsScopeArtifactSequence", + "TriggerFieldsScopeProject", + "AlertSeverity", + "EventTriggeringConditionType", + "TriggerScopeType", + "TriggeredActionType", +] diff --git a/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9c642ecc57d31fc207779407dd8502b503a11267 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/create_automation.cpython-312.pyc b/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/create_automation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..67fff79dc7171ecb4c74e8679a4152d6ae7a7810 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/create_automation.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/create_generic_webhook_integration.cpython-312.pyc b/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/create_generic_webhook_integration.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..580908bb517cd70797eac0de94952459b1fd12ec Binary files /dev/null and b/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/create_generic_webhook_integration.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/delete_automation.cpython-312.pyc b/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/delete_automation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4ea24da09cc9b8b14f0436402f2d9a5cac202d99 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/delete_automation.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/enums.cpython-312.pyc b/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/enums.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..76d4ad8c508123e79f5aa8b1b1dbae3b4b333c12 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/enums.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/fragments.cpython-312.pyc b/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/fragments.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f26e5b550c4f8ebc278a63328084c11726b981ab Binary files /dev/null and b/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/fragments.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/generic_webhook_integrations_by_entity.cpython-312.pyc b/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/generic_webhook_integrations_by_entity.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9d345a6325d91d7be432de794d4c13dfdd67b5d9 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/generic_webhook_integrations_by_entity.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/get_automations.cpython-312.pyc b/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/get_automations.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9b47c6aac90d71a078406b353de8a6468e999734 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/get_automations.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/get_automations_by_entity.cpython-312.pyc b/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/get_automations_by_entity.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9e5f8e3e9bda27038b77daf414da970c445c0fe7 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/get_automations_by_entity.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/input_types.cpython-312.pyc b/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/input_types.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d7768fb99c34ccde33450787ed4787f267708f3a Binary files /dev/null and b/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/input_types.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/integrations_by_entity.cpython-312.pyc b/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/integrations_by_entity.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bb946b35c86bd9ce983a651dd45126f886a2584e Binary files /dev/null and b/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/integrations_by_entity.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/operations.cpython-312.pyc b/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/operations.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2f406f26058fa546f01b5ebe06cc25f297f5d814 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/operations.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/slack_integrations_by_entity.cpython-312.pyc b/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/slack_integrations_by_entity.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6264726f590cd3fd85c3b413abac7a2bafa222c3 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/slack_integrations_by_entity.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/update_automation.cpython-312.pyc b/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/update_automation.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fe2cc722d84282a2512425b772e570c2f63edbc4 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/automations/_generated/__pycache__/update_automation.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/automations/_generated/create_automation.py b/lib/python3.12/site-packages/wandb/automations/_generated/create_automation.py new file mode 100644 index 0000000000000000000000000000000000000000..63c551908b961291bd0ef87e79e90bae24e66686 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/automations/_generated/create_automation.py @@ -0,0 +1,17 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/automations/ + +from __future__ import annotations + +from typing import Optional + +from wandb._pydantic import GQLBase + +from .fragments import CreateAutomationResult + + +class CreateAutomation(GQLBase): + result: Optional[CreateAutomationResult] + + +CreateAutomation.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/automations/_generated/create_generic_webhook_integration.py b/lib/python3.12/site-packages/wandb/automations/_generated/create_generic_webhook_integration.py new file mode 100644 index 0000000000000000000000000000000000000000..cc215f452cd3ed0e97bc28996787a2b739216a7e --- /dev/null +++ b/lib/python3.12/site-packages/wandb/automations/_generated/create_generic_webhook_integration.py @@ -0,0 +1,43 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/automations/ + +from __future__ import annotations + +from typing import Literal, Optional, Union + +from pydantic import Field + +from wandb._pydantic import GQLBase, Typename + +from .fragments import GenericWebhookIntegrationFields + + +class CreateGenericWebhookIntegration(GQLBase): + create_generic_webhook_integration: Optional[ + CreateGenericWebhookIntegrationCreateGenericWebhookIntegration + ] = Field(alias="createGenericWebhookIntegration") + + +class CreateGenericWebhookIntegrationCreateGenericWebhookIntegration(GQLBase): + integration: Union[ + CreateGenericWebhookIntegrationCreateGenericWebhookIntegrationIntegrationIntegration, + CreateGenericWebhookIntegrationCreateGenericWebhookIntegrationIntegrationGenericWebhookIntegration, + ] = Field(discriminator="typename__") + + +class CreateGenericWebhookIntegrationCreateGenericWebhookIntegrationIntegrationIntegration( + GQLBase +): + typename__: Typename[ + Literal["GitHubOAuthIntegration", "Integration", "SlackIntegration"] + ] + + +class CreateGenericWebhookIntegrationCreateGenericWebhookIntegrationIntegrationGenericWebhookIntegration( + GenericWebhookIntegrationFields +): + typename__: Typename[Literal["GenericWebhookIntegration"]] + + +CreateGenericWebhookIntegration.model_rebuild() +CreateGenericWebhookIntegrationCreateGenericWebhookIntegration.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/automations/_generated/delete_automation.py b/lib/python3.12/site-packages/wandb/automations/_generated/delete_automation.py new file mode 100644 index 0000000000000000000000000000000000000000..9d889f69b6df5c7c158ba26a22cea58dfc7ef338 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/automations/_generated/delete_automation.py @@ -0,0 +1,17 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/automations/ + +from __future__ import annotations + +from typing import Optional + +from wandb._pydantic import GQLBase + +from .fragments import DeleteAutomationResult + + +class DeleteAutomation(GQLBase): + result: Optional[DeleteAutomationResult] + + +DeleteAutomation.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/automations/_generated/enums.py b/lib/python3.12/site-packages/wandb/automations/_generated/enums.py new file mode 100644 index 0000000000000000000000000000000000000000..6f69e5ebf8be598b4a205ce611a78545a8dbc548 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/automations/_generated/enums.py @@ -0,0 +1,33 @@ +# Generated by ariadne-codegen +# Source: core/api/graphql/schemas/schema-latest.graphql + +from __future__ import annotations + +from enum import Enum + + +class EventTriggeringConditionType(str, Enum): + CREATE_ARTIFACT = "CREATE_ARTIFACT" + UPDATE_ARTIFACT_ALIAS = "UPDATE_ARTIFACT_ALIAS" + ADD_ARTIFACT_ALIAS = "ADD_ARTIFACT_ALIAS" + LINK_MODEL = "LINK_MODEL" + RUN_METRIC = "RUN_METRIC" + RUN_METRIC_CHANGE = "RUN_METRIC_CHANGE" + + +class TriggerScopeType(str, Enum): + PROJECT = "PROJECT" + ARTIFACT_COLLECTION = "ARTIFACT_COLLECTION" + + +class AlertSeverity(str, Enum): + INFO = "INFO" + WARN = "WARN" + ERROR = "ERROR" + + +class TriggeredActionType(str, Enum): + QUEUE_JOB = "QUEUE_JOB" + NOTIFICATION = "NOTIFICATION" + GENERIC_WEBHOOK = "GENERIC_WEBHOOK" + NO_OP = "NO_OP" diff --git a/lib/python3.12/site-packages/wandb/automations/_generated/fragments.py b/lib/python3.12/site-packages/wandb/automations/_generated/fragments.py new file mode 100644 index 0000000000000000000000000000000000000000..2572295b179af07a8c785dee3e45317e6e509d5d --- /dev/null +++ b/lib/python3.12/site-packages/wandb/automations/_generated/fragments.py @@ -0,0 +1,358 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/automations/ + +from __future__ import annotations + +from datetime import datetime +from typing import List, Literal, Optional, Union + +from pydantic import Field +from typing_extensions import Annotated + +from wandb._pydantic import GQLBase, GQLId, Typename + +from .enums import AlertSeverity, EventTriggeringConditionType + + +class ArtifactPortfolioScopeFields(GQLBase): + typename__: Typename[Literal["ArtifactPortfolio"]] = "ArtifactPortfolio" + id: GQLId + name: str + + +class ArtifactSequenceScopeFields(GQLBase): + typename__: Typename[Literal["ArtifactSequence"]] = "ArtifactSequence" + id: GQLId + name: str + + +class CreateAutomationResult(GQLBase): + typename__: Typename[Literal["CreateFilterTriggerPayload"]] = ( + "CreateFilterTriggerPayload" + ) + trigger: Optional[TriggerFields] + + +class DeleteAutomationResult(GQLBase): + typename__: Typename[Literal["DeleteTriggerPayload"]] = "DeleteTriggerPayload" + success: bool + + +class FilterEventFields(GQLBase): + typename__: Typename[Literal["FilterEventTriggeringCondition"]] = ( + "FilterEventTriggeringCondition" + ) + event_type: EventTriggeringConditionType = Field(alias="eventType") + filter: str + + +class GenericWebhookActionFields(GQLBase): + typename__: Typename[Literal["GenericWebhookTriggeredAction"]] = ( + "GenericWebhookTriggeredAction" + ) + integration: Union[ + GenericWebhookActionFieldsIntegrationIntegration, + GenericWebhookActionFieldsIntegrationGenericWebhookIntegration, + ] = Field(discriminator="typename__") + request_payload: Optional[str] = Field(alias="requestPayload") + + +class GenericWebhookActionFieldsIntegrationIntegration(GQLBase): + typename__: Typename[ + Literal["GitHubOAuthIntegration", "Integration", "SlackIntegration"] + ] + + +class GenericWebhookIntegrationConnectionFields(GQLBase): + typename__: Typename[Literal["IntegrationConnection"]] = "IntegrationConnection" + page_info: PageInfoFields = Field(alias="pageInfo") + edges: List[GenericWebhookIntegrationConnectionFieldsEdges] + + +class GenericWebhookIntegrationConnectionFieldsEdges(GQLBase): + cursor: str + node: Optional[ + Annotated[ + Union[ + GenericWebhookIntegrationConnectionFieldsEdgesNodeIntegration, + GenericWebhookIntegrationConnectionFieldsEdgesNodeGenericWebhookIntegration, + ], + Field(discriminator="typename__"), + ] + ] + + +class GenericWebhookIntegrationConnectionFieldsEdgesNodeIntegration(GQLBase): + typename__: Typename[ + Literal["GitHubOAuthIntegration", "Integration", "SlackIntegration"] + ] + + +class GenericWebhookIntegrationFields(GQLBase): + typename__: Typename[Literal["GenericWebhookIntegration"]] = ( + "GenericWebhookIntegration" + ) + id: GQLId + name: str + url_endpoint: str = Field(alias="urlEndpoint") + + +class IntegrationConnectionFields(GQLBase): + typename__: Typename[Literal["IntegrationConnection"]] = "IntegrationConnection" + page_info: PageInfoFields = Field(alias="pageInfo") + edges: List[IntegrationConnectionFieldsEdges] + + +class IntegrationConnectionFieldsEdges(GQLBase): + cursor: str + node: Optional[ + Annotated[ + Union[ + IntegrationConnectionFieldsEdgesNodeIntegration, + IntegrationConnectionFieldsEdgesNodeGenericWebhookIntegration, + IntegrationConnectionFieldsEdgesNodeSlackIntegration, + ], + Field(discriminator="typename__"), + ] + ] + + +class IntegrationConnectionFieldsEdgesNodeIntegration(GQLBase): + typename__: Typename[Literal["GitHubOAuthIntegration", "Integration"]] + + +class NoOpActionFields(GQLBase): + typename__: Typename[Literal["NoOpTriggeredAction"]] = "NoOpTriggeredAction" + no_op: Optional[bool] = Field(alias="noOp") + + +class NotificationActionFields(GQLBase): + typename__: Typename[Literal["NotificationTriggeredAction"]] = ( + "NotificationTriggeredAction" + ) + integration: Union[ + NotificationActionFieldsIntegrationIntegration, + NotificationActionFieldsIntegrationSlackIntegration, + ] = Field(discriminator="typename__") + title: Optional[str] + message: Optional[str] + severity: Optional[AlertSeverity] + + +class NotificationActionFieldsIntegrationIntegration(GQLBase): + typename__: Typename[ + Literal["GenericWebhookIntegration", "GitHubOAuthIntegration", "Integration"] + ] + + +class PageInfoFields(GQLBase): + end_cursor: Optional[str] = Field(alias="endCursor") + has_next_page: bool = Field(alias="hasNextPage") + + +class ProjectConnectionFields(GQLBase): + typename__: Typename[Literal["ProjectConnection"]] = "ProjectConnection" + page_info: PageInfoFields = Field(alias="pageInfo") + edges: List[ProjectConnectionFieldsEdges] + + +class ProjectConnectionFieldsEdges(GQLBase): + cursor: str + node: Optional[ProjectConnectionFieldsEdgesNode] + + +class ProjectConnectionFieldsEdgesNode(GQLBase): + triggers: List[TriggerFields] + + +class ProjectScopeFields(GQLBase): + typename__: Typename[Literal["Project"]] = "Project" + id: GQLId + name: str + + +class QueueJobActionFields(GQLBase): + typename__: Typename[Literal["QueueJobTriggeredAction"]] = "QueueJobTriggeredAction" + queue: Optional[QueueJobActionFieldsQueue] + template: str + + +class QueueJobActionFieldsQueue(GQLBase): + id: GQLId + name: str + + +class SlackIntegrationConnectionFields(GQLBase): + typename__: Typename[Literal["IntegrationConnection"]] = "IntegrationConnection" + page_info: PageInfoFields = Field(alias="pageInfo") + edges: List[SlackIntegrationConnectionFieldsEdges] + + +class SlackIntegrationConnectionFieldsEdges(GQLBase): + cursor: str + node: Optional[ + Annotated[ + Union[ + SlackIntegrationConnectionFieldsEdgesNodeIntegration, + SlackIntegrationConnectionFieldsEdgesNodeSlackIntegration, + ], + Field(discriminator="typename__"), + ] + ] + + +class SlackIntegrationConnectionFieldsEdgesNodeIntegration(GQLBase): + typename__: Typename[ + Literal["GenericWebhookIntegration", "GitHubOAuthIntegration", "Integration"] + ] + + +class SlackIntegrationFields(GQLBase): + typename__: Typename[Literal["SlackIntegration"]] = "SlackIntegration" + id: GQLId + team_name: str = Field(alias="teamName") + channel_name: str = Field(alias="channelName") + + +class TriggerFields(GQLBase): + typename__: Typename[Literal["Trigger"]] = "Trigger" + id: GQLId + created_at: datetime = Field(alias="createdAt") + updated_at: Optional[datetime] = Field(alias="updatedAt") + name: str + description: Optional[str] + enabled: bool + scope: Union[ + TriggerFieldsScopeProject, + TriggerFieldsScopeArtifactSequence, + TriggerFieldsScopeArtifactPortfolio, + ] = Field(discriminator="typename__") + event: TriggerFieldsEventFilterEventTriggeringCondition + action: Union[ + TriggerFieldsActionQueueJobTriggeredAction, + TriggerFieldsActionNotificationTriggeredAction, + TriggerFieldsActionGenericWebhookTriggeredAction, + TriggerFieldsActionNoOpTriggeredAction, + ] = Field(discriminator="typename__") + + +class UpdateAutomationResult(GQLBase): + typename__: Typename[Literal["UpdateFilterTriggerPayload"]] = ( + "UpdateFilterTriggerPayload" + ) + trigger: Optional[TriggerFields] + + +class TriggerFieldsScopeArtifactPortfolio(ArtifactPortfolioScopeFields): + typename__: Typename[Literal["ArtifactPortfolio"]] + + +class TriggerFieldsScopeArtifactSequence(ArtifactSequenceScopeFields): + typename__: Typename[Literal["ArtifactSequence"]] + + +class TriggerFieldsEventFilterEventTriggeringCondition(FilterEventFields): + typename__: Typename[Literal["FilterEventTriggeringCondition"]] + + +class TriggerFieldsActionGenericWebhookTriggeredAction(GenericWebhookActionFields): + typename__: Typename[Literal["GenericWebhookTriggeredAction"]] + + +class GenericWebhookActionFieldsIntegrationGenericWebhookIntegration( + GenericWebhookIntegrationFields +): + typename__: Typename[Literal["GenericWebhookIntegration"]] + + +class GenericWebhookIntegrationConnectionFieldsEdgesNodeGenericWebhookIntegration( + GenericWebhookIntegrationFields +): + typename__: Typename[Literal["GenericWebhookIntegration"]] + + +class IntegrationConnectionFieldsEdgesNodeGenericWebhookIntegration( + GenericWebhookIntegrationFields +): + typename__: Typename[Literal["GenericWebhookIntegration"]] + + +class TriggerFieldsActionNoOpTriggeredAction(NoOpActionFields): + typename__: Typename[Literal["NoOpTriggeredAction"]] + + +class TriggerFieldsActionNotificationTriggeredAction(NotificationActionFields): + typename__: Typename[Literal["NotificationTriggeredAction"]] + + +class TriggerFieldsScopeProject(ProjectScopeFields): + typename__: Typename[Literal["Project"]] + + +class TriggerFieldsActionQueueJobTriggeredAction(QueueJobActionFields): + typename__: Typename[Literal["QueueJobTriggeredAction"]] + + +class IntegrationConnectionFieldsEdgesNodeSlackIntegration(SlackIntegrationFields): + typename__: Typename[Literal["SlackIntegration"]] + + +class NotificationActionFieldsIntegrationSlackIntegration(SlackIntegrationFields): + typename__: Typename[Literal["SlackIntegration"]] + + +class SlackIntegrationConnectionFieldsEdgesNodeSlackIntegration(SlackIntegrationFields): + typename__: Typename[Literal["SlackIntegration"]] + + +ArtifactPortfolioScopeFields.model_rebuild() +ArtifactSequenceScopeFields.model_rebuild() +CreateAutomationResult.model_rebuild() +DeleteAutomationResult.model_rebuild() +FilterEventFields.model_rebuild() +GenericWebhookActionFields.model_rebuild() +GenericWebhookActionFieldsIntegrationIntegration.model_rebuild() +GenericWebhookIntegrationConnectionFields.model_rebuild() +GenericWebhookIntegrationConnectionFieldsEdges.model_rebuild() +GenericWebhookIntegrationConnectionFieldsEdgesNodeIntegration.model_rebuild() +GenericWebhookIntegrationFields.model_rebuild() +IntegrationConnectionFields.model_rebuild() +IntegrationConnectionFieldsEdges.model_rebuild() +IntegrationConnectionFieldsEdgesNodeIntegration.model_rebuild() +NoOpActionFields.model_rebuild() +NotificationActionFields.model_rebuild() +NotificationActionFieldsIntegrationIntegration.model_rebuild() +PageInfoFields.model_rebuild() +ProjectConnectionFields.model_rebuild() +ProjectConnectionFieldsEdges.model_rebuild() +ProjectConnectionFieldsEdgesNode.model_rebuild() +ProjectScopeFields.model_rebuild() +QueueJobActionFields.model_rebuild() +QueueJobActionFieldsQueue.model_rebuild() +SlackIntegrationConnectionFields.model_rebuild() +SlackIntegrationConnectionFieldsEdges.model_rebuild() +SlackIntegrationConnectionFieldsEdgesNodeIntegration.model_rebuild() +SlackIntegrationFields.model_rebuild() +TriggerFields.model_rebuild() +UpdateAutomationResult.model_rebuild() +TriggerFieldsScopeArtifactPortfolio.model_rebuild() +TriggerFieldsScopeArtifactSequence.model_rebuild() +TriggerFieldsEventFilterEventTriggeringCondition.model_rebuild() +TriggerFieldsActionGenericWebhookTriggeredAction.model_rebuild() +GenericWebhookActionFieldsIntegrationGenericWebhookIntegration.model_rebuild() +GenericWebhookIntegrationConnectionFieldsEdgesNodeGenericWebhookIntegration.model_rebuild() +IntegrationConnectionFieldsEdgesNodeGenericWebhookIntegration.model_rebuild() +TriggerFieldsActionNoOpTriggeredAction.model_rebuild() +TriggerFieldsActionNotificationTriggeredAction.model_rebuild() +PageInfoFields.model_rebuild() +PageInfoFields.model_rebuild() +PageInfoFields.model_rebuild() +PageInfoFields.model_rebuild() +TriggerFieldsScopeProject.model_rebuild() +TriggerFieldsActionQueueJobTriggeredAction.model_rebuild() +IntegrationConnectionFieldsEdgesNodeSlackIntegration.model_rebuild() +NotificationActionFieldsIntegrationSlackIntegration.model_rebuild() +SlackIntegrationConnectionFieldsEdgesNodeSlackIntegration.model_rebuild() +TriggerFields.model_rebuild() +TriggerFields.model_rebuild() +TriggerFields.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/automations/_generated/generic_webhook_integrations_by_entity.py b/lib/python3.12/site-packages/wandb/automations/_generated/generic_webhook_integrations_by_entity.py new file mode 100644 index 0000000000000000000000000000000000000000..4374455b8da8c051a43df1e189ce986c64d9d64c --- /dev/null +++ b/lib/python3.12/site-packages/wandb/automations/_generated/generic_webhook_integrations_by_entity.py @@ -0,0 +1,22 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/automations/ + +from __future__ import annotations + +from typing import Optional + +from wandb._pydantic import GQLBase + +from .fragments import GenericWebhookIntegrationConnectionFields + + +class GenericWebhookIntegrationsByEntity(GQLBase): + entity: Optional[GenericWebhookIntegrationsByEntityEntity] + + +class GenericWebhookIntegrationsByEntityEntity(GQLBase): + integrations: Optional[GenericWebhookIntegrationConnectionFields] + + +GenericWebhookIntegrationsByEntity.model_rebuild() +GenericWebhookIntegrationsByEntityEntity.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/automations/_generated/get_automations.py b/lib/python3.12/site-packages/wandb/automations/_generated/get_automations.py new file mode 100644 index 0000000000000000000000000000000000000000..d9be5c3fe2f55010e134e140342a9032fd07ae7d --- /dev/null +++ b/lib/python3.12/site-packages/wandb/automations/_generated/get_automations.py @@ -0,0 +1,24 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/automations/ + +from __future__ import annotations + +from typing import Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase + +from .fragments import ProjectConnectionFields + + +class GetAutomations(GQLBase): + search_scope: Optional[GetAutomationsSearchScope] = Field(alias="searchScope") + + +class GetAutomationsSearchScope(GQLBase): + projects: Optional[ProjectConnectionFields] + + +GetAutomations.model_rebuild() +GetAutomationsSearchScope.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/automations/_generated/get_automations_by_entity.py b/lib/python3.12/site-packages/wandb/automations/_generated/get_automations_by_entity.py new file mode 100644 index 0000000000000000000000000000000000000000..2be7ecdce1186fbf902fe2342fcd0feca35ce603 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/automations/_generated/get_automations_by_entity.py @@ -0,0 +1,26 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/automations/ + +from __future__ import annotations + +from typing import Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase + +from .fragments import ProjectConnectionFields + + +class GetAutomationsByEntity(GQLBase): + search_scope: Optional[GetAutomationsByEntitySearchScope] = Field( + alias="searchScope" + ) + + +class GetAutomationsByEntitySearchScope(GQLBase): + projects: Optional[ProjectConnectionFields] + + +GetAutomationsByEntity.model_rebuild() +GetAutomationsByEntitySearchScope.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/automations/_generated/input_types.py b/lib/python3.12/site-packages/wandb/automations/_generated/input_types.py new file mode 100644 index 0000000000000000000000000000000000000000..968d461fd06d9639f8960b8114ce0c7d1219b708 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/automations/_generated/input_types.py @@ -0,0 +1,104 @@ +# Generated by ariadne-codegen +# Source: core/api/graphql/schemas/schema-latest.graphql + +from __future__ import annotations + +from typing import Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase, GQLId + +from .enums import ( + AlertSeverity, + EventTriggeringConditionType, + TriggeredActionType, + TriggerScopeType, +) + + +class CreateGenericWebhookIntegrationInput(GQLBase): + entity_name: str = Field(alias="entityName") + url_endpoint: str = Field(alias="urlEndpoint") + name: str + secret_ref: Optional[str] = Field(alias="secretRef", default=None) + access_token_ref: Optional[str] = Field(alias="accessTokenRef", default=None) + client_mutation_id: Optional[str] = Field(alias="clientMutationId", default=None) + + +class QueueJobActionInput(GQLBase): + queue_id: GQLId = Field(alias="queueID") + template: str + + +class NotificationActionInput(GQLBase): + integration_id: GQLId = Field(alias="integrationID") + title: Optional[str] = None + message: Optional[str] = None + severity: Optional[AlertSeverity] = None + + +class GenericWebhookActionInput(GQLBase): + integration_id: GQLId = Field(alias="integrationID") + request_payload: Optional[str] = Field(alias="requestPayload", default=None) + + +class NoOpTriggeredActionInput(GQLBase): + no_op: Optional[bool] = Field(alias="noOp", default=None) + + +class TriggeredActionConfig(GQLBase): + queue_job_action_input: Optional[QueueJobActionInput] = Field( + alias="queueJobActionInput", default=None + ) + notification_action_input: Optional[NotificationActionInput] = Field( + alias="notificationActionInput", default=None + ) + generic_webhook_action_input: Optional[GenericWebhookActionInput] = Field( + alias="genericWebhookActionInput", default=None + ) + no_op_action_input: Optional[NoOpTriggeredActionInput] = Field( + alias="noOpActionInput", default=None + ) + + +class CreateFilterTriggerInput(GQLBase): + name: str + description: Optional[str] = None + triggering_event_type: EventTriggeringConditionType = Field( + alias="triggeringEventType" + ) + scope_type: TriggerScopeType = Field(alias="scopeType") + scope_id: GQLId = Field(alias="scopeID") + event_filter: str = Field(alias="eventFilter") + triggered_action_type: TriggeredActionType = Field(alias="triggeredActionType") + triggered_action_config: TriggeredActionConfig = Field( + alias="triggeredActionConfig" + ) + enabled: bool + client_mutation_id: Optional[str] = Field(alias="clientMutationId", default=None) + + +class UpdateFilterTriggerInput(GQLBase): + id: GQLId + name: Optional[str] = None + description: Optional[str] = None + triggering_event_type: Optional[EventTriggeringConditionType] = Field( + alias="triggeringEventType", default=None + ) + scope_type: Optional[TriggerScopeType] = Field(alias="scopeType", default=None) + scope_id: Optional[GQLId] = Field(alias="scopeID", default=None) + event_filter: Optional[str] = Field(alias="eventFilter", default=None) + triggered_action_type: Optional[TriggeredActionType] = Field( + alias="triggeredActionType", default=None + ) + triggered_action_config: Optional[TriggeredActionConfig] = Field( + alias="triggeredActionConfig", default=None + ) + enabled: Optional[bool] = None + client_mutation_id: Optional[str] = Field(alias="clientMutationId", default=None) + + +TriggeredActionConfig.model_rebuild() +CreateFilterTriggerInput.model_rebuild() +UpdateFilterTriggerInput.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/automations/_generated/integrations_by_entity.py b/lib/python3.12/site-packages/wandb/automations/_generated/integrations_by_entity.py new file mode 100644 index 0000000000000000000000000000000000000000..993c0ad33149e3fd184fc7eddf461d929df0d3ce --- /dev/null +++ b/lib/python3.12/site-packages/wandb/automations/_generated/integrations_by_entity.py @@ -0,0 +1,22 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/automations/ + +from __future__ import annotations + +from typing import Optional + +from wandb._pydantic import GQLBase + +from .fragments import IntegrationConnectionFields + + +class IntegrationsByEntity(GQLBase): + entity: Optional[IntegrationsByEntityEntity] + + +class IntegrationsByEntityEntity(GQLBase): + integrations: Optional[IntegrationConnectionFields] + + +IntegrationsByEntity.model_rebuild() +IntegrationsByEntityEntity.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/automations/_generated/operations.py b/lib/python3.12/site-packages/wandb/automations/_generated/operations.py new file mode 100644 index 0000000000000000000000000000000000000000..5d67acd8d6541180d8ab7a99f8f876364592a338 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/automations/_generated/operations.py @@ -0,0 +1,647 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/automations/ + +__all__ = [ + "CREATE_AUTOMATION_GQL", + "CREATE_GENERIC_WEBHOOK_INTEGRATION_GQL", + "DELETE_AUTOMATION_GQL", + "GENERIC_WEBHOOK_INTEGRATIONS_BY_ENTITY_GQL", + "GET_AUTOMATIONS_BY_ENTITY_GQL", + "GET_AUTOMATIONS_GQL", + "INTEGRATIONS_BY_ENTITY_GQL", + "SLACK_INTEGRATIONS_BY_ENTITY_GQL", + "UPDATE_AUTOMATION_GQL", +] + +GET_AUTOMATIONS_GQL = """ +query GetAutomations($cursor: String, $perPage: Int) { + searchScope: viewer { + projects(after: $cursor, first: $perPage) { + ...ProjectConnectionFields + } + } +} + +fragment ArtifactPortfolioScopeFields on ArtifactPortfolio { + __typename + id + name +} + +fragment ArtifactSequenceScopeFields on ArtifactSequence { + __typename + id + name +} + +fragment FilterEventFields on FilterEventTriggeringCondition { + __typename + eventType + filter +} + +fragment GenericWebhookActionFields on GenericWebhookTriggeredAction { + __typename + integration { + __typename + ...GenericWebhookIntegrationFields + } + requestPayload +} + +fragment GenericWebhookIntegrationFields on GenericWebhookIntegration { + __typename + id + name + urlEndpoint +} + +fragment NoOpActionFields on NoOpTriggeredAction { + __typename + noOp +} + +fragment NotificationActionFields on NotificationTriggeredAction { + __typename + integration { + __typename + ...SlackIntegrationFields + } + title + message + severity +} + +fragment PageInfoFields on PageInfo { + endCursor + hasNextPage +} + +fragment ProjectConnectionFields on ProjectConnection { + __typename + pageInfo { + ...PageInfoFields + } + edges { + cursor + node { + triggers { + ...TriggerFields + } + } + } +} + +fragment ProjectScopeFields on Project { + __typename + id + name +} + +fragment QueueJobActionFields on QueueJobTriggeredAction { + __typename + queue { + id + name + } + template +} + +fragment SlackIntegrationFields on SlackIntegration { + __typename + id + teamName + channelName +} + +fragment TriggerFields on Trigger { + __typename + id + createdAt + updatedAt + name + description + enabled + scope { + __typename + ...ProjectScopeFields + ...ArtifactPortfolioScopeFields + ...ArtifactSequenceScopeFields + } + event: triggeringCondition { + __typename + ...FilterEventFields + } + action: triggeredAction { + __typename + ...QueueJobActionFields + ...NotificationActionFields + ...GenericWebhookActionFields + ...NoOpActionFields + } +} +""" + +GET_AUTOMATIONS_BY_ENTITY_GQL = """ +query GetAutomationsByEntity($entityName: String!, $cursor: String, $perPage: Int) { + searchScope: entity(name: $entityName) { + projects(after: $cursor, first: $perPage) { + ...ProjectConnectionFields + } + } +} + +fragment ArtifactPortfolioScopeFields on ArtifactPortfolio { + __typename + id + name +} + +fragment ArtifactSequenceScopeFields on ArtifactSequence { + __typename + id + name +} + +fragment FilterEventFields on FilterEventTriggeringCondition { + __typename + eventType + filter +} + +fragment GenericWebhookActionFields on GenericWebhookTriggeredAction { + __typename + integration { + __typename + ...GenericWebhookIntegrationFields + } + requestPayload +} + +fragment GenericWebhookIntegrationFields on GenericWebhookIntegration { + __typename + id + name + urlEndpoint +} + +fragment NoOpActionFields on NoOpTriggeredAction { + __typename + noOp +} + +fragment NotificationActionFields on NotificationTriggeredAction { + __typename + integration { + __typename + ...SlackIntegrationFields + } + title + message + severity +} + +fragment PageInfoFields on PageInfo { + endCursor + hasNextPage +} + +fragment ProjectConnectionFields on ProjectConnection { + __typename + pageInfo { + ...PageInfoFields + } + edges { + cursor + node { + triggers { + ...TriggerFields + } + } + } +} + +fragment ProjectScopeFields on Project { + __typename + id + name +} + +fragment QueueJobActionFields on QueueJobTriggeredAction { + __typename + queue { + id + name + } + template +} + +fragment SlackIntegrationFields on SlackIntegration { + __typename + id + teamName + channelName +} + +fragment TriggerFields on Trigger { + __typename + id + createdAt + updatedAt + name + description + enabled + scope { + __typename + ...ProjectScopeFields + ...ArtifactPortfolioScopeFields + ...ArtifactSequenceScopeFields + } + event: triggeringCondition { + __typename + ...FilterEventFields + } + action: triggeredAction { + __typename + ...QueueJobActionFields + ...NotificationActionFields + ...GenericWebhookActionFields + ...NoOpActionFields + } +} +""" + +CREATE_AUTOMATION_GQL = """ +mutation CreateAutomation($params: CreateFilterTriggerInput!) { + result: createFilterTrigger(input: $params) { + ...CreateAutomationResult + } +} + +fragment ArtifactPortfolioScopeFields on ArtifactPortfolio { + __typename + id + name +} + +fragment ArtifactSequenceScopeFields on ArtifactSequence { + __typename + id + name +} + +fragment CreateAutomationResult on CreateFilterTriggerPayload { + __typename + trigger { + ...TriggerFields + } +} + +fragment FilterEventFields on FilterEventTriggeringCondition { + __typename + eventType + filter +} + +fragment GenericWebhookActionFields on GenericWebhookTriggeredAction { + __typename + integration { + __typename + ...GenericWebhookIntegrationFields + } + requestPayload +} + +fragment GenericWebhookIntegrationFields on GenericWebhookIntegration { + __typename + id + name + urlEndpoint +} + +fragment NoOpActionFields on NoOpTriggeredAction { + __typename + noOp +} + +fragment NotificationActionFields on NotificationTriggeredAction { + __typename + integration { + __typename + ...SlackIntegrationFields + } + title + message + severity +} + +fragment ProjectScopeFields on Project { + __typename + id + name +} + +fragment QueueJobActionFields on QueueJobTriggeredAction { + __typename + queue { + id + name + } + template +} + +fragment SlackIntegrationFields on SlackIntegration { + __typename + id + teamName + channelName +} + +fragment TriggerFields on Trigger { + __typename + id + createdAt + updatedAt + name + description + enabled + scope { + __typename + ...ProjectScopeFields + ...ArtifactPortfolioScopeFields + ...ArtifactSequenceScopeFields + } + event: triggeringCondition { + __typename + ...FilterEventFields + } + action: triggeredAction { + __typename + ...QueueJobActionFields + ...NotificationActionFields + ...GenericWebhookActionFields + ...NoOpActionFields + } +} +""" + +UPDATE_AUTOMATION_GQL = """ +mutation UpdateAutomation($params: UpdateFilterTriggerInput!) { + result: updateFilterTrigger(input: $params) { + ...UpdateAutomationResult + } +} + +fragment ArtifactPortfolioScopeFields on ArtifactPortfolio { + __typename + id + name +} + +fragment ArtifactSequenceScopeFields on ArtifactSequence { + __typename + id + name +} + +fragment FilterEventFields on FilterEventTriggeringCondition { + __typename + eventType + filter +} + +fragment GenericWebhookActionFields on GenericWebhookTriggeredAction { + __typename + integration { + __typename + ...GenericWebhookIntegrationFields + } + requestPayload +} + +fragment GenericWebhookIntegrationFields on GenericWebhookIntegration { + __typename + id + name + urlEndpoint +} + +fragment NoOpActionFields on NoOpTriggeredAction { + __typename + noOp +} + +fragment NotificationActionFields on NotificationTriggeredAction { + __typename + integration { + __typename + ...SlackIntegrationFields + } + title + message + severity +} + +fragment ProjectScopeFields on Project { + __typename + id + name +} + +fragment QueueJobActionFields on QueueJobTriggeredAction { + __typename + queue { + id + name + } + template +} + +fragment SlackIntegrationFields on SlackIntegration { + __typename + id + teamName + channelName +} + +fragment TriggerFields on Trigger { + __typename + id + createdAt + updatedAt + name + description + enabled + scope { + __typename + ...ProjectScopeFields + ...ArtifactPortfolioScopeFields + ...ArtifactSequenceScopeFields + } + event: triggeringCondition { + __typename + ...FilterEventFields + } + action: triggeredAction { + __typename + ...QueueJobActionFields + ...NotificationActionFields + ...GenericWebhookActionFields + ...NoOpActionFields + } +} + +fragment UpdateAutomationResult on UpdateFilterTriggerPayload { + __typename + trigger { + ...TriggerFields + } +} +""" + +DELETE_AUTOMATION_GQL = """ +mutation DeleteAutomation($id: ID!) { + result: deleteTrigger(input: {triggerID: $id}) { + ...DeleteAutomationResult + } +} + +fragment DeleteAutomationResult on DeleteTriggerPayload { + __typename + success +} +""" + +INTEGRATIONS_BY_ENTITY_GQL = """ +query IntegrationsByEntity($entityName: String!, $cursor: String, $perPage: Int) { + entity(name: $entityName) { + integrations(after: $cursor, first: $perPage) { + ...IntegrationConnectionFields + } + } +} + +fragment GenericWebhookIntegrationFields on GenericWebhookIntegration { + __typename + id + name + urlEndpoint +} + +fragment IntegrationConnectionFields on IntegrationConnection { + __typename + pageInfo { + ...PageInfoFields + } + edges { + cursor + node { + __typename + ...SlackIntegrationFields + ...GenericWebhookIntegrationFields + } + } +} + +fragment PageInfoFields on PageInfo { + endCursor + hasNextPage +} + +fragment SlackIntegrationFields on SlackIntegration { + __typename + id + teamName + channelName +} +""" + +SLACK_INTEGRATIONS_BY_ENTITY_GQL = """ +query SlackIntegrationsByEntity($entityName: String!, $cursor: String, $perPage: Int) { + entity(name: $entityName) { + integrations(after: $cursor, first: $perPage) { + ...SlackIntegrationConnectionFields + } + } +} + +fragment PageInfoFields on PageInfo { + endCursor + hasNextPage +} + +fragment SlackIntegrationConnectionFields on IntegrationConnection { + __typename + pageInfo { + ...PageInfoFields + } + edges { + cursor + node { + __typename + ...SlackIntegrationFields + } + } +} + +fragment SlackIntegrationFields on SlackIntegration { + __typename + id + teamName + channelName +} +""" + +GENERIC_WEBHOOK_INTEGRATIONS_BY_ENTITY_GQL = """ +query GenericWebhookIntegrationsByEntity($entityName: String!, $cursor: String, $perPage: Int) { + entity(name: $entityName) { + integrations(after: $cursor, first: $perPage) { + ...GenericWebhookIntegrationConnectionFields + } + } +} + +fragment GenericWebhookIntegrationConnectionFields on IntegrationConnection { + __typename + pageInfo { + ...PageInfoFields + } + edges { + cursor + node { + __typename + ...GenericWebhookIntegrationFields + } + } +} + +fragment GenericWebhookIntegrationFields on GenericWebhookIntegration { + __typename + id + name + urlEndpoint +} + +fragment PageInfoFields on PageInfo { + endCursor + hasNextPage +} +""" + +CREATE_GENERIC_WEBHOOK_INTEGRATION_GQL = """ +mutation CreateGenericWebhookIntegration($params: CreateGenericWebhookIntegrationInput!) { + createGenericWebhookIntegration(input: $params) { + integration { + __typename + ...GenericWebhookIntegrationFields + } + } +} + +fragment GenericWebhookIntegrationFields on GenericWebhookIntegration { + __typename + id + name + urlEndpoint +} +""" diff --git a/lib/python3.12/site-packages/wandb/automations/_generated/slack_integrations_by_entity.py b/lib/python3.12/site-packages/wandb/automations/_generated/slack_integrations_by_entity.py new file mode 100644 index 0000000000000000000000000000000000000000..6576c2d585135b6171572682f4d69fef498e1a07 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/automations/_generated/slack_integrations_by_entity.py @@ -0,0 +1,22 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/automations/ + +from __future__ import annotations + +from typing import Optional + +from wandb._pydantic import GQLBase + +from .fragments import SlackIntegrationConnectionFields + + +class SlackIntegrationsByEntity(GQLBase): + entity: Optional[SlackIntegrationsByEntityEntity] + + +class SlackIntegrationsByEntityEntity(GQLBase): + integrations: Optional[SlackIntegrationConnectionFields] + + +SlackIntegrationsByEntity.model_rebuild() +SlackIntegrationsByEntityEntity.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/automations/_generated/update_automation.py b/lib/python3.12/site-packages/wandb/automations/_generated/update_automation.py new file mode 100644 index 0000000000000000000000000000000000000000..12eb1be3549da444ca5532a0c02bede8e4def571 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/automations/_generated/update_automation.py @@ -0,0 +1,17 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/automations/ + +from __future__ import annotations + +from typing import Optional + +from wandb._pydantic import GQLBase + +from .fragments import UpdateAutomationResult + + +class UpdateAutomation(GQLBase): + result: Optional[UpdateAutomationResult] + + +UpdateAutomation.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/automations/_utils.py b/lib/python3.12/site-packages/wandb/automations/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c2b874d152496b7cd16b59fbabc65542cd95f7f6 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/automations/_utils.py @@ -0,0 +1,235 @@ +from __future__ import annotations + +from typing import Any, Collection, Final, Optional, Protocol, TypedDict + +from pydantic import Field +from typing_extensions import Annotated, Self, Unpack + +from wandb._pydantic import GQLBase, GQLId, computed_field, model_validator, to_json + +from ._filters import MongoLikeFilter +from ._generated import ( + CreateFilterTriggerInput, + QueueJobActionInput, + TriggeredActionConfig, + UpdateFilterTriggerInput, +) +from ._validators import to_input_action +from .actions import ( + ActionType, + DoNothing, + InputAction, + SavedAction, + SendNotification, + SendWebhook, +) +from .automations import Automation, NewAutomation +from .events import EventType, InputEvent, RunMetricFilter, _WrappedSavedEventFilter +from .scopes import AutomationScope, ScopeType + +EXCLUDED_INPUT_EVENTS: Final[Collection[EventType]] = frozenset( + { + EventType.UPDATE_ARTIFACT_ALIAS, + } +) +"""Event types that should not be assigned when creating/updating automations.""" + +EXCLUDED_INPUT_ACTIONS: Final[Collection[ActionType]] = frozenset( + { + ActionType.QUEUE_JOB, + } +) +"""Action types that should not be assigned when creating/updating automations.""" + +ALWAYS_SUPPORTED_EVENTS: Final[Collection[EventType]] = frozenset( + { + EventType.CREATE_ARTIFACT, + EventType.LINK_ARTIFACT, + EventType.ADD_ARTIFACT_ALIAS, + } +) +"""Event types that we can safely assume all contemporary server versions support.""" + +ALWAYS_SUPPORTED_ACTIONS: Final[Collection[ActionType]] = frozenset( + { + ActionType.NOTIFICATION, + ActionType.GENERIC_WEBHOOK, + } +) +"""Action types that we can safely assume all contemporary server versions support.""" + + +class HasId(Protocol): + id: str + + +def extract_id(obj: HasId | str) -> str: + return obj.id if hasattr(obj, "id") else obj + + +# --------------------------------------------------------------------------- +ACTION_CONFIG_KEYS: dict[ActionType, str] = { + ActionType.NOTIFICATION: "notification_action_input", + ActionType.GENERIC_WEBHOOK: "generic_webhook_action_input", + ActionType.NO_OP: "no_op_action_input", + ActionType.QUEUE_JOB: "queue_job_action_input", +} + + +class InputActionConfig(TriggeredActionConfig): + """A `TriggeredActionConfig` that prepares the action config for saving an automation.""" + + # NOTE: `QueueJobActionInput` for defining a Launch job is deprecated, + # so while it's allowed here to update EXISTING mutations, we don't + # currently expose it through the public API to create NEW automations. + queue_job_action_input: Optional[QueueJobActionInput] = None + + notification_action_input: Optional[SendNotification] = None + generic_webhook_action_input: Optional[SendWebhook] = None + no_op_action_input: Optional[DoNothing] = None + + +def prepare_action_config_input(obj: SavedAction | InputAction) -> dict[str, Any]: + """Prepare the `TriggeredActionConfig` input, nesting the action input inside the appropriate key. + + This is necessary to conform to the schemas for: + - CreateFilterTriggerInput + - UpdateFilterTriggerInput + """ + # Delegate to inner validators to convert SavedAction -> InputAction types, if needed. + obj = to_input_action(obj) + return InputActionConfig(**{ACTION_CONFIG_KEYS[obj.action_type]: obj}).model_dump() + + +def prepare_event_filter_input( + obj: _WrappedSavedEventFilter | MongoLikeFilter | RunMetricFilter, +) -> str: + """Prepare the `EventFilter` input, unnesting the filter if needed and serializing to JSON. + + This is necessary to conform to the schemas for: + - CreateFilterTriggerInput + - UpdateFilterTriggerInput + """ + # Input event filters are nested one level deeper than saved event filters. + # Note that this is NOT the case for run/run metric filters. + # + # Yes, this is confusing. It's also necessary to conform to under-the-hood + # schemas and logic in the backend. + filter_to_serialize = ( + obj.filter if isinstance(obj, _WrappedSavedEventFilter) else obj + ) + return to_json(filter_to_serialize) + + +class WriteAutomationsKwargs(TypedDict, total=False): + """Keyword arguments that can be passed to create or update an automation.""" + + name: str + description: str + enabled: bool + scope: AutomationScope + event: InputEvent + action: InputAction + + +class ValidatedCreateInput(GQLBase, extra="forbid", frozen=True): + """Validated automation parameters, prepared for creating a new automation. + + Note: Users should never need to instantiate this class directly. + """ + + name: str + description: Optional[str] = None + enabled: bool = True + + # ------------------------------------------------------------------------------ + # Set on instantiation, but used to derive other fields and deliberately + # EXCLUDED from the final GraphQL request vars + event: Annotated[InputEvent, Field(exclude=True)] + action: Annotated[InputAction, Field(exclude=True)] + + # ------------------------------------------------------------------------------ + # Derived fields to match the input schemas + @computed_field + def scope_type(self) -> ScopeType: + return self.event.scope.scope_type + + @computed_field + def scope_id(self) -> GQLId: + return self.event.scope.id + + @computed_field + def triggering_event_type(self) -> EventType: + return self.event.event_type + + @computed_field + def event_filter(self) -> str: + return prepare_event_filter_input(self.event.filter) + + @computed_field + def triggered_action_type(self) -> ActionType: + return self.action.action_type + + @computed_field + def triggered_action_config(self) -> dict[str, Any]: + return prepare_action_config_input(self.action) + + # ------------------------------------------------------------------------------ + # Custom validation + @model_validator(mode="after") + def _forbid_legacy_event_types(self) -> Self: + if (type_ := self.event.event_type) in EXCLUDED_INPUT_EVENTS: + raise ValueError(f"{type_!r} events cannot be assigned to automations.") + return self + + @model_validator(mode="after") + def _forbid_legacy_action_types(self) -> Self: + if (type_ := self.action.action_type) in EXCLUDED_INPUT_ACTIONS: + raise ValueError(f"{type_!r} actions cannot be assigned to automations.") + return self + + +def prepare_to_create( + obj: NewAutomation | None = None, + /, + **kwargs: Unpack[WriteAutomationsKwargs], +) -> CreateFilterTriggerInput: + """Prepares the payload to create an automation in a GraphQL request.""" + # Validate all input variables, and prepare as expected by the GraphQL request. + # - if an object is provided, override its fields with any keyword args + # - otherwise, instantiate from the keyword args + + # NOTE: `exclude_none=True` drops fields that are still `None`. + # + # This assumes that `None` is good enough for now as a sentinel + # "unset" value. If this proves insufficient, revisit in the future, + # as it should be reasonably easy to implement a custom sentinel + # type later on. + obj_dict = {**obj.model_dump(exclude_none=True), **kwargs} if obj else kwargs + validated = ValidatedCreateInput(**obj_dict) + return CreateFilterTriggerInput.model_validate(validated) + + +def prepare_to_update( + obj: Automation | None = None, + /, + **kwargs: Unpack[WriteAutomationsKwargs], +) -> UpdateFilterTriggerInput: + """Prepares the payload to update an automation in a GraphQL request.""" + # Validate all values: + # - if an object is provided, override its fields with any keyword args + # - otherwise, instantiate from the keyword args + v_obj = Automation(**{**dict(obj or {}), **kwargs}) + + return UpdateFilterTriggerInput( + id=v_obj.id, + name=v_obj.name, + description=v_obj.description, + enabled=v_obj.enabled, + scope_type=v_obj.scope.scope_type, + scope_id=v_obj.scope.id, + triggering_event_type=v_obj.event.event_type, + event_filter=prepare_event_filter_input(v_obj.event.filter), + triggered_action_type=v_obj.action.action_type, + triggered_action_config=prepare_action_config_input(v_obj.action), + ) diff --git a/lib/python3.12/site-packages/wandb/automations/_validators.py b/lib/python3.12/site-packages/wandb/automations/_validators.py new file mode 100644 index 0000000000000000000000000000000000000000..7c2f63285d1c3b70397e95c0b21c6f1cbfd91a97 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/automations/_validators.py @@ -0,0 +1,165 @@ +from __future__ import annotations + +from enum import Enum +from functools import singledispatch +from itertools import chain +from typing import Any, TypeVar + +from pydantic_core import PydanticUseDefault + +from ._filters import And, FilterExpr, In, Nor, Not, NotIn, Op, Or + +T = TypeVar("T") + + +class LenientStrEnum(str, Enum): + """A string enum allowing for case-insensitive lookups by value. + + May include other internal customizations if needed. + + Note: This is a bespoke, internal implementation and NOT intended as a + backport of `enum.StrEnum` from Python 3.11+. + """ + + def __repr__(self) -> str: + return self.name + + @classmethod + def _missing_(cls, value: object) -> Any: + # Accept case-insensitive enum values + if isinstance(value, str): + v = value.lower() + return next((e for e in cls if e.value.lower() == v), None) + return None + + +def default_if_none(v: Any) -> Any: + """A before-validator validator that coerces `None` to the default field value instead.""" + # https://docs.pydantic.dev/2.11/api/pydantic_core/#pydantic_core.PydanticUseDefault + if v is None: + raise PydanticUseDefault + return v + + +def upper_if_str(v: Any) -> Any: + return v.strip().upper() if isinstance(v, str) else v + + +# ---------------------------------------------------------------------------- +def to_scope(v: Any) -> Any: + """Convert eligible objects, including pre-existing `wandb` types, to an automation scope.""" + from wandb.apis.public import ArtifactCollection, Project + + from .scopes import ProjectScope, _ArtifactPortfolioScope, _ArtifactSequenceScope + + if isinstance(v, Project): + return ProjectScope(id=v.id, name=v.name) + if isinstance(v, ArtifactCollection): + cls = _ArtifactSequenceScope if v.is_sequence() else _ArtifactPortfolioScope + return cls(id=v.id, name=v.name) + return v + + +def to_saved_action(v: Any) -> Any: + """If necessary (and possible), convert the object to a saved action.""" + from .actions import ( + DoNothing, + SavedNoOpAction, + SavedNotificationAction, + SavedWebhookAction, + SendNotification, + SendWebhook, + ) + + if isinstance(v, SendNotification): + return SavedNotificationAction( + integration={"id": v.integration_id}, + **v.model_dump(exclude={"integration_id"}), + ) + if isinstance(v, SendWebhook): + return SavedWebhookAction( + integration={"id": v.integration_id}, + **v.model_dump(exclude={"integration_id"}), + ) + if isinstance(v, DoNothing): + return SavedNoOpAction.model_validate(v) + + return v + + +def to_input_action(v: Any) -> Any: + """If necessary (and possible), convert the object to an input action.""" + from .actions import ( + DoNothing, + SavedNoOpAction, + SavedNotificationAction, + SavedWebhookAction, + SendNotification, + SendWebhook, + ) + + if isinstance(v, SavedNotificationAction): + return SendNotification( + integration_id=v.integration.id, + **v.model_dump(exclude={"integration"}), + ) + if isinstance(v, SavedWebhookAction): + return SendWebhook( + integration_id=v.integration.id, + **v.model_dump(exclude={"integration"}), + ) + if isinstance(v, SavedNoOpAction): + return DoNothing.model_validate(v) + + return v + + +# ---------------------------------------------------------------------------- +@singledispatch +def simplify_op(op: Op | FilterExpr) -> Op | FilterExpr: + """Simplify a MongoDB filter by removing and unnesting redundant operators.""" + return op + + +@simplify_op.register +def _(op: And) -> Op: + # {"$and": []} -> {"$and": []} + if not (args := op.and_): + return op + + # {"$and": [op]} -> op + if len(args) == 1: + return simplify_op(args[0]) + + # {"$and": [op, {"$and": [op2, ...]}]} -> {"$and": [op, op2, ...]} + flattened = chain.from_iterable(x.and_ if isinstance(x, And) else [x] for x in args) + return And(and_=map(simplify_op, flattened)) + + +@simplify_op.register +def _(op: Or) -> Op: + # {"$or": []} -> {"$or": []} + if not (args := op.or_): + return op + + # {"$or": [op]} -> op + if len(args) == 1: + return simplify_op(args[0]) + + # {"$or": [op, {"$or": [op2, ...]}]} -> {"$or": [op, op2, ...]} + flattened = chain.from_iterable(x.or_ if isinstance(x, Or) else [x] for x in args) + return Or(or_=map(simplify_op, flattened)) + + +@simplify_op.register +def _(op: Not) -> Op: + inner = op.not_ + + # {"$not": {"$not": op}} -> op + # {"$not": {"$or": [op, ...]}} -> {"$nor": [op, ...]} + # {"$not": {"$nor": [op, ...]}} -> {"$or": [op, ...]} + # {"$not": {"$in": [op, ...]}} -> {"$nin": [op, ...]} + # {"$not": {"$nin": [op, ...]}} -> {"$in": [op, ...]} + if isinstance(inner, (Not, Or, Nor, In, NotIn)): + return simplify_op(~inner) + return Not(not_=simplify_op(inner)) diff --git a/lib/python3.12/site-packages/wandb/automations/actions.py b/lib/python3.12/site-packages/wandb/automations/actions.py new file mode 100644 index 0000000000000000000000000000000000000000..1b61b7037f2bb28778f887cd3e8df535db6beab0 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/automations/actions.py @@ -0,0 +1,218 @@ +"""Actions that are triggered by W&B Automations.""" + +from __future__ import annotations + +from typing import Any, Literal, Optional, Union + +from pydantic import BeforeValidator, Field +from typing_extensions import Annotated, Self, get_args + +from wandb._pydantic import GQLBase, GQLId, SerializedToJson, Typename + +from ._generated import ( + AlertSeverity, + GenericWebhookActionFields, + GenericWebhookActionInput, + NoOpActionFields, + NoOpTriggeredActionInput, + NotificationActionFields, + NotificationActionInput, + QueueJobActionFields, +) +from ._validators import ( + LenientStrEnum, + default_if_none, + to_input_action, + to_saved_action, + upper_if_str, +) +from .integrations import SlackIntegration, WebhookIntegration + + +# NOTE: Name shortened for readability and defined publicly for easier access +class ActionType(LenientStrEnum): + """The type of action triggered by an automation.""" + + QUEUE_JOB = "QUEUE_JOB" # NOTE: Deprecated for creation + NOTIFICATION = "NOTIFICATION" + GENERIC_WEBHOOK = "GENERIC_WEBHOOK" + NO_OP = "NO_OP" + + +# ------------------------------------------------------------------------------ +# Saved types: for parsing response data from saved automations + + +# NOTE: `QueueJobActionInput` for defining a Launch job is deprecated, +# so while we allow parsing it from previously saved Automations, we deliberately +# don't currently expose it in the API for creating automations. +class SavedLaunchJobAction(QueueJobActionFields): + action_type: Literal[ActionType.QUEUE_JOB] = ActionType.QUEUE_JOB + + +# FIXME: Find a better place to put these OR a better way to handle the +# conversion from `InputAction` -> `SavedAction`. +# +# Necessary placeholder class defs for converting: +# - `SendNotification -> SavedNotificationAction` +# - `SendWebhook -> SavedWebhookAction` +# +# The "input" types (`Send{Notification,Webhook}`) will only have an `integration_id`, +# and we don't want/need to fetch the other `{Slack,Webhook}Integration` fields if +# we can avoid it. +class _SavedActionSlackIntegration(GQLBase, extra="allow"): + typename__: Typename[Literal["SlackIntegration"]] = "SlackIntegration" + id: GQLId + + +class _SavedActionWebhookIntegration(GQLBase, extra="allow"): + typename__: Typename[Literal["GenericWebhookIntegration"]] = ( + "GenericWebhookIntegration" + ) + id: GQLId + + +class SavedNotificationAction(NotificationActionFields): + action_type: Literal[ActionType.NOTIFICATION] = ActionType.NOTIFICATION + integration: _SavedActionSlackIntegration + + +class SavedWebhookAction(GenericWebhookActionFields): + action_type: Literal[ActionType.GENERIC_WEBHOOK] = ActionType.GENERIC_WEBHOOK + integration: _SavedActionWebhookIntegration + + # We override the type of the `requestPayload` field since the original GraphQL + # schema (and generated class) effectively defines it as a string, when we know + # and need to anticipate the expected structure of the JSON-serialized data. + request_payload: Annotated[ + Optional[SerializedToJson[dict[str, Any]]], + Field(alias="requestPayload"), + ] = None # type: ignore[assignment] + + +class SavedNoOpAction(NoOpActionFields, frozen=True): + action_type: Literal[ActionType.NO_OP] = ActionType.NO_OP + + no_op: Annotated[bool, BeforeValidator(default_if_none)] = True + """Placeholder field, only needed to conform to schema requirements. + + There should never be a need to set this field explicitly, as its value is ignored. + """ + + +# for type annotations +SavedAction = Annotated[ + Union[ + SavedLaunchJobAction, + SavedNotificationAction, + SavedWebhookAction, + SavedNoOpAction, + ], + BeforeValidator(to_saved_action), + Field(discriminator="typename__"), +] +# for runtime type checks +SavedActionTypes: tuple[type, ...] = get_args(SavedAction.__origin__) # type: ignore[attr-defined] + + +# ------------------------------------------------------------------------------ +# Input types: for creating or updating automations +class _BaseActionInput(GQLBase): + action_type: Annotated[ActionType, Field(frozen=True)] + """The kind of action to be triggered.""" + + +class SendNotification(_BaseActionInput, NotificationActionInput): + """Defines an automation action that sends a (Slack) notification.""" + + action_type: Literal[ActionType.NOTIFICATION] = ActionType.NOTIFICATION + + integration_id: GQLId + """The ID of the Slack integration that will be used to send the notification.""" + + # Note: Validation aliases are meant to provide continuity with prior `wandb.alert()` API. + title: str = "" + """The title of the sent notification.""" + + message: Annotated[str, Field(validation_alias="text")] = "" + """The message body of the sent notification.""" + + severity: Annotated[ + AlertSeverity, + BeforeValidator(upper_if_str), # Be helpful by ensuring uppercase strings + Field(validation_alias="level"), + ] = AlertSeverity.INFO + """The severity (`INFO`, `WARN`, `ERROR`) of the sent notification.""" + + @classmethod + def from_integration( + cls, + integration: SlackIntegration, + *, + title: str = "", + text: str = "", + level: AlertSeverity = AlertSeverity.INFO, + ) -> Self: + """Define a notification action that sends to the given (Slack) integration.""" + return cls( + integration_id=integration.id, + title=title, + message=text, + severity=level, + ) + + +class SendWebhook(_BaseActionInput, GenericWebhookActionInput): + """Defines an automation action that sends a webhook request.""" + + action_type: Literal[ActionType.GENERIC_WEBHOOK] = ActionType.GENERIC_WEBHOOK + + integration_id: GQLId + """The ID of the webhook integration that will be used to send the request.""" + + # overrides the generated field type to parse/serialize JSON strings + request_payload: Optional[SerializedToJson[dict[str, Any]]] = Field( # type: ignore[assignment] + default=None, alias="requestPayload" + ) + """The payload, possibly with template variables, to send in the webhook request.""" + + @classmethod + def from_integration( + cls, + integration: WebhookIntegration, + *, + payload: Optional[SerializedToJson[dict[str, Any]]] = None, + ) -> Self: + """Define a webhook action that sends to the given (webhook) integration.""" + return cls(integration_id=integration.id, request_payload=payload) + + +class DoNothing(_BaseActionInput, NoOpTriggeredActionInput, frozen=True): + """Defines an automation action that intentionally does nothing.""" + + action_type: Literal[ActionType.NO_OP] = ActionType.NO_OP + + no_op: Annotated[bool, BeforeValidator(default_if_none)] = True + """Placeholder field which exists only to satisfy backend schema requirements. + + There should never be a need to set this field explicitly, as its value is ignored. + """ + + +# for type annotations +InputAction = Annotated[ + Union[ + SendNotification, + SendWebhook, + DoNothing, + ], + BeforeValidator(to_input_action), + Field(discriminator="action_type"), +] +# for runtime type checks +InputActionTypes: tuple[type, ...] = get_args(InputAction.__origin__) # type: ignore[attr-defined] + +__all__ = [ + "ActionType", + *(cls.__name__ for cls in InputActionTypes), +] diff --git a/lib/python3.12/site-packages/wandb/automations/automations.py b/lib/python3.12/site-packages/wandb/automations/automations.py new file mode 100644 index 0000000000000000000000000000000000000000..9ebd34046e90b651c79c2b8ab857e88108216561 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/automations/automations.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +from datetime import datetime +from typing import Optional + +from pydantic import Field +from typing_extensions import Annotated + +from wandb._pydantic import GQLBase, GQLId + +from ._generated import TriggerFields +from .actions import InputAction, SavedAction +from .events import InputEvent, SavedEvent +from .scopes import AutomationScope + + +# ------------------------------------------------------------------------------ +# Saved types: for parsing response data from saved automations +class Automation(TriggerFields): + """A local instance of a saved W&B automation.""" + + id: GQLId + + created_at: Annotated[datetime, Field(repr=False, frozen=True, alias="createdAt")] + """The date and time when this automation was created.""" + + updated_at: Annotated[ + Optional[datetime], Field(repr=False, frozen=True, alias="updatedAt") + ] = None + """The date and time when this automation was last updated, if applicable.""" + + name: str + """The name of this automation.""" + + description: Optional[str] + """An optional description of this automation.""" + + enabled: bool + """Whether this automation is enabled. Only enabled automations will trigger.""" + + event: SavedEvent + """The event that will trigger this automation.""" + + scope: AutomationScope + """The scope in which the triggering event must occur.""" + + action: SavedAction + """The action that will execute when this automation is triggered.""" + + +class NewAutomation(GQLBase, extra="forbid", validate_default=False): + """A new automation to be created.""" + + name: Optional[str] = None + """The name of this automation.""" + + description: Optional[str] = None + """An optional description of this automation.""" + + enabled: Optional[bool] = None + """Whether this automation is enabled. Only enabled automations will trigger.""" + + event: Optional[InputEvent] = None + """The event that will trigger this automation.""" + + # Ensure that the event and its scope are always consistent, if the event is set. + @property + def scope(self) -> Optional[AutomationScope]: + """The scope in which the triggering event must occur.""" + return self.event.scope if self.event else None + + @scope.setter + def scope(self, value: AutomationScope) -> None: + if self.event is None: + raise ValueError("Cannot set `scope` for an automation with no `event`") + self.event.scope = value + + action: Optional[InputAction] = None + """The action that will execute when this automation is triggered.""" + + +__all__ = [ + "Automation", + "NewAutomation", +] diff --git a/lib/python3.12/site-packages/wandb/automations/events.py b/lib/python3.12/site-packages/wandb/automations/events.py new file mode 100644 index 0000000000000000000000000000000000000000..8666e7faafa9b7c1de3c05d2888eb8ed9011c2a2 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/automations/events.py @@ -0,0 +1,285 @@ +"""Events that trigger W&B Automations.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Literal, Optional, Union + +from pydantic import Field +from typing_extensions import Annotated, Self, get_args + +from wandb._pydantic import ( + GQLBase, + SerializedToJson, + ensure_json, + field_validator, + model_validator, + pydantic_isinstance, +) + +from ._filters import And, MongoLikeFilter, Or +from ._filters.expressions import FilterableField +from ._filters.run_metrics import MetricChangeFilter, MetricThresholdFilter, MetricVal +from ._generated import FilterEventFields +from ._validators import LenientStrEnum, simplify_op +from .actions import InputAction, InputActionTypes, SavedActionTypes +from .scopes import ArtifactCollectionScope, AutomationScope, ProjectScope + +if TYPE_CHECKING: + from .automations import NewAutomation + + +# NOTE: Re-defined publicly with a more readable name for easier access +class EventType(LenientStrEnum): + """The type of event that triggers an automation.""" + + # --------------------------------------------------------------------------- + # Events triggered by GraphQL mutations + UPDATE_ARTIFACT_ALIAS = "UPDATE_ARTIFACT_ALIAS" # NOTE: Avoid in new automations + + CREATE_ARTIFACT = "CREATE_ARTIFACT" + ADD_ARTIFACT_ALIAS = "ADD_ARTIFACT_ALIAS" + LINK_ARTIFACT = "LINK_MODEL" + # Note: "LINK_MODEL" is the (legacy) value expected by the backend, but we + # name it "LINK_ARTIFACT" here in the public API for clarity and consistency. + + # --------------------------------------------------------------------------- + # Events triggered by Run conditions + RUN_METRIC_THRESHOLD = "RUN_METRIC" + RUN_METRIC_CHANGE = "RUN_METRIC_CHANGE" + + +# ------------------------------------------------------------------------------ +# Saved types: for parsing response data from saved automations + + +# Note: In GQL responses containing saved automation data, the filter is wrapped in an extra `filter` key. +class _WrappedSavedEventFilter(GQLBase): # from: TriggeringFilterEvent + filter: SerializedToJson[MongoLikeFilter] = And() + + +class _WrappedMetricFilter(GQLBase): # from: RunMetricFilter + threshold_filter: Optional[MetricThresholdFilter] = None + change_filter: Optional[MetricChangeFilter] = None + + @model_validator(mode="before") + @classmethod + def _wrap_metric_filter(cls, v: Any) -> Any: + if pydantic_isinstance(v, MetricThresholdFilter): + return cls(threshold_filter=v) + if pydantic_isinstance(v, MetricChangeFilter): + return cls(change_filter=v) + return v + + @model_validator(mode="after") + def _ensure_exactly_one_set(self) -> Self: + set_fields = [name for name, val in self if (val is not None)] + + if not set_fields: + all_names = ", ".join(map(repr, type(self).model_fields)) + raise ValueError(f"Expected one of: {all_names}") + + if len(set_fields) > 1: + set_names = ", ".join(map(repr, set_fields)) + raise ValueError(f"Expected exactly one metric filter, got: {set_names}") + + return self + + @property + def event_type(self) -> EventType: + if self.threshold_filter is not None: + return EventType.RUN_METRIC_THRESHOLD + if self.change_filter is not None: + return EventType.RUN_METRIC_CHANGE + raise RuntimeError("Expected one of: `threshold_filter` or `change_filter`") + + +class RunMetricFilter(GQLBase): # from: TriggeringRunMetricEvent + run: Annotated[SerializedToJson[MongoLikeFilter], Field(alias="run_filter")] = And() + metric: Annotated[_WrappedMetricFilter, Field(alias="run_metric_filter")] + + # ------------------------------------------------------------------------------ + legacy_metric_filter: Annotated[ + Optional[SerializedToJson[MetricThresholdFilter]], + Field(alias="metric_filter", deprecated=True), + ] = None + """Deprecated legacy field that was previously used to define run metric threshold events. + + For new automations, use the `metric` field (`run_metric_filter` JSON alias) instead. + """ + + @model_validator(mode="before") + @classmethod + def _wrap_metric_filter(cls, v: Any) -> Any: + if pydantic_isinstance(v, (MetricThresholdFilter, MetricChangeFilter)): + # If only an (unnested) metric filter is given, nest it under the + # `metric` field, delegating to inner validator(s) for further + # wrapping/nesting, if needed. + # This is necessary to conform to the expected backend schema. + return cls(metric=v) + return v + + @field_validator("run", mode="after") + def _wrap_run_filter(cls, v: MongoLikeFilter) -> Any: + v_new = simplify_op(v) + return v_new if pydantic_isinstance(v_new, And) else And(and_=[v_new]) + + +class SavedEvent(FilterEventFields): # from: FilterEventTriggeringCondition + """A triggering event from a saved automation.""" + + event_type: Annotated[EventType, Field(frozen=True)] # type: ignore[assignment] + + # We override the type of the `filter` field in order to enforce the expected + # structure for the JSON data when validating and serializing. + filter: SerializedToJson[Union[_WrappedSavedEventFilter, RunMetricFilter]] + """The condition(s) under which this event triggers an automation.""" + + +# ------------------------------------------------------------------------------ +# Input types: for creating or updating automations + + +# Note: The GQL input for "eventFilter" does NOT wrap the filter in an extra `filter` key, unlike the +# eventFilter returned in responses for saved automations. +class _BaseEventInput(GQLBase): + event_type: EventType + + scope: AutomationScope + """The scope of the event.""" + + filter: SerializedToJson[Any] + + def then(self, action: InputAction) -> NewAutomation: + """Define a new Automation in which this event triggers the given action.""" + from .automations import NewAutomation + + if isinstance(action, (InputActionTypes, SavedActionTypes)): + return NewAutomation(event=self, action=action) + + raise TypeError(f"Expected a valid action, got: {type(action).__qualname__!r}") + + def __rshift__(self, other: InputAction) -> NewAutomation: + """Implements `event >> action` to define an Automation with this event and action.""" + return self.then(other) + + +# ------------------------------------------------------------------------------ +# Events that trigger on specific mutations in the backend +class _BaseMutationEventInput(_BaseEventInput): + filter: SerializedToJson[MongoLikeFilter] = And() + """Additional condition(s), if any, that must be met for this event to trigger an automation.""" + + @field_validator("filter", mode="after") + def _wrap_filter(cls, v: Any) -> Any: + """Ensure the given filter is wrapped like: `{"$or": [{"$and": []}]}`. + + This is awkward but necessary, because the frontend expects this format. + """ + v_new = simplify_op(v) + v_new = v_new if pydantic_isinstance(v_new, And) else And(and_=[v_new]) + return Or(or_=[v_new]) + + +class OnLinkArtifact(_BaseMutationEventInput): + """A new artifact is linked to a collection.""" + + event_type: Literal[EventType.LINK_ARTIFACT] = EventType.LINK_ARTIFACT + + +class OnAddArtifactAlias(_BaseMutationEventInput): + """A new alias is assigned to an artifact.""" + + event_type: Literal[EventType.ADD_ARTIFACT_ALIAS] = EventType.ADD_ARTIFACT_ALIAS + + +class OnCreateArtifact(_BaseMutationEventInput): + """A new artifact is created.""" + + event_type: Literal[EventType.CREATE_ARTIFACT] = EventType.CREATE_ARTIFACT + + scope: ArtifactCollectionScope + """The scope of the event: only artifact collections are valid scopes for this event.""" + + +# ------------------------------------------------------------------------------ +# Events that trigger on run conditions +class _BaseRunEventInput(_BaseEventInput): + scope: ProjectScope + """The scope of the event: only projects are valid scopes for this event.""" + + +class OnRunMetric(_BaseRunEventInput): + """A run metric satisfies a user-defined condition.""" + + event_type: Literal[EventType.RUN_METRIC_THRESHOLD, EventType.RUN_METRIC_CHANGE] + + filter: SerializedToJson[RunMetricFilter] + """Run and/or metric condition(s) that must be satisfied for this event to trigger an automation.""" + + @model_validator(mode="before") + @classmethod + def _infer_event_type(cls, data: Any) -> Any: + """Infer the event type at validation time from the inner filter. + + This allows this class to accommodate both "threshold" and "change" metric + filter types, which are can only be determined after parsing and validating + the inner JSON data. + """ + if isinstance(data, dict) and (raw_filter := data.get("filter")): + # At this point, `raw_filter` may or may not be JSON-serialized + parsed_filter = RunMetricFilter.model_validate_json(ensure_json(raw_filter)) + return {**data, "event_type": parsed_filter.metric.event_type} + + return data + + +# for type annotations +InputEvent = Annotated[ + Union[ + OnLinkArtifact, + OnAddArtifactAlias, + OnCreateArtifact, + OnRunMetric, + ], + Field(discriminator="event_type"), +] +# for runtime type checks +InputEventTypes: tuple[type, ...] = get_args(InputEvent.__origin__) # type: ignore[attr-defined] + + +# ---------------------------------------------------------------------------- + + +class RunEvent: + name = FilterableField(server_name="display_name") + # `Run.name` is actually filtered on `Run.display_name` in the backend. + # We can't reasonably expect users to know this a priori, so + # automatically fix it here. + + @staticmethod + def metric(name: str) -> MetricVal: + """Define a metric filter condition.""" + return MetricVal(name=name) + + +class ArtifactEvent: + alias = FilterableField() + + +MetricThresholdFilter.model_rebuild() +RunMetricFilter.model_rebuild() +_WrappedSavedEventFilter.model_rebuild() + +OnLinkArtifact.model_rebuild() +OnAddArtifactAlias.model_rebuild() +OnCreateArtifact.model_rebuild() +OnRunMetric.model_rebuild() + +__all__ = [ + "EventType", + *(cls.__name__ for cls in InputEventTypes), + "RunEvent", + "ArtifactEvent", + "MetricThresholdFilter", + "MetricChangeFilter", +] diff --git a/lib/python3.12/site-packages/wandb/automations/integrations.py b/lib/python3.12/site-packages/wandb/automations/integrations.py new file mode 100644 index 0000000000000000000000000000000000000000..b296d6d82830ca02483bbd367478c517f1d749a4 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/automations/integrations.py @@ -0,0 +1,45 @@ +from typing import Union + +from pydantic import Field +from typing_extensions import Annotated + +from wandb._pydantic import GQLBase +from wandb.automations._generated import ( + GenericWebhookIntegrationFields, + SlackIntegrationFields, +) + + +class SlackIntegration(SlackIntegrationFields): + team_name: str + """The name of the Slack workspace (not the W&B team) that this integration is associated with.""" + + channel_name: str + """The name of the Slack channel that this integration will post messages to.""" + + +class WebhookIntegration(GenericWebhookIntegrationFields): + name: str + """The name of this webhook integration.""" + + url_endpoint: str + """The URL that this webhook will POST events to.""" + + +Integration = Annotated[ + Union[SlackIntegration, WebhookIntegration], + Field(discriminator="typename__"), +] + + +# For parsing integration instances from paginated responses +class _IntegrationEdge(GQLBase): + cursor: str + node: Integration + + +__all__ = [ + "Integration", + "SlackIntegration", + "WebhookIntegration", +] diff --git a/lib/python3.12/site-packages/wandb/automations/scopes.py b/lib/python3.12/site-packages/wandb/automations/scopes.py new file mode 100644 index 0000000000000000000000000000000000000000..630a67bcc3bd63f22b2a653be31889a9ea631527 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/automations/scopes.py @@ -0,0 +1,78 @@ +"""Scopes in which a W&B Automation can be triggered.""" + +from __future__ import annotations + +from typing import Literal, Union + +from pydantic import BeforeValidator, Field +from typing_extensions import Annotated, TypeAlias, get_args + +from wandb._pydantic import GQLBase +from wandb.automations._generated import ( + ArtifactPortfolioScopeFields, + ArtifactSequenceScopeFields, + ProjectScopeFields, +) + +from ._validators import LenientStrEnum, to_scope + + +# NOTE: Re-defined publicly with a more readable name for easier access +class ScopeType(LenientStrEnum): + """The kind of scope that triggers an automation.""" + + PROJECT = "PROJECT" + ARTIFACT_COLLECTION = "ARTIFACT_COLLECTION" + + +class _BaseScope(GQLBase): + scope_type: Annotated[ScopeType, Field(frozen=True)] + + +class _ArtifactSequenceScope(_BaseScope, ArtifactSequenceScopeFields): + """An automation scope defined by a specific `ArtifactSequence`.""" + + scope_type: Literal[ScopeType.ARTIFACT_COLLECTION] = ScopeType.ARTIFACT_COLLECTION + + +class _ArtifactPortfolioScope(_BaseScope, ArtifactPortfolioScopeFields): + """An automation scope defined by a specific `ArtifactPortfolio` (e.g. a registry collection).""" + + scope_type: Literal[ScopeType.ARTIFACT_COLLECTION] = ScopeType.ARTIFACT_COLLECTION + + +# for type annotations +ArtifactCollectionScope = Annotated[ + Union[_ArtifactSequenceScope, _ArtifactPortfolioScope], + BeforeValidator(to_scope), + Field(discriminator="typename__"), +] +"""An automation scope defined by a specific `ArtifactCollection`.""" + +# for runtime type checks +ArtifactCollectionScopeTypes: tuple[type, ...] = get_args( + ArtifactCollectionScope.__origin__ # type: ignore[attr-defined] +) + + +class ProjectScope(_BaseScope, ProjectScopeFields): + """An automation scope defined by a specific `Project`.""" + + scope_type: Literal[ScopeType.PROJECT] = ScopeType.PROJECT + + +# for type annotations +AutomationScope: TypeAlias = Annotated[ + Union[_ArtifactSequenceScope, _ArtifactPortfolioScope, ProjectScope], + BeforeValidator(to_scope), + Field(discriminator="typename__"), +] +# for runtime type checks +AutomationScopeTypes: tuple[type, ...] = get_args(AutomationScope.__origin__) # type: ignore[attr-defined] + + +__all__ = [ + "ScopeType", + "ArtifactCollectionScope", + "ProjectScope", +] diff --git a/lib/python3.12/site-packages/wandb/beta/__pycache__/workflows.cpython-312.pyc b/lib/python3.12/site-packages/wandb/beta/__pycache__/workflows.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..58427b5dc7745ec96c5f94458b852441913ef4a5 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/beta/__pycache__/workflows.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/beta/workflows.py b/lib/python3.12/site-packages/wandb/beta/workflows.py new file mode 100644 index 0000000000000000000000000000000000000000..84bf4a8f0e8460d0e7b32a396e5d941431f2f90a --- /dev/null +++ b/lib/python3.12/site-packages/wandb/beta/workflows.py @@ -0,0 +1,288 @@ +import json +import os +from typing import Any, Dict, List, Optional, Union + +import wandb +import wandb.data_types as data_types +from wandb.data_types import _SavedModel +from wandb.sdk import wandb_setup +from wandb.sdk.artifacts.artifact import Artifact +from wandb.sdk.artifacts.artifact_manifest_entry import ArtifactManifestEntry + + +def _add_any( + artifact: Artifact, + path_or_obj: Union[ + str, ArtifactManifestEntry, data_types.WBValue + ], # todo: add dataframe + name: Optional[str], +) -> Any: + """Add an object to an artifact. + + High-level wrapper to add object(s) to an artifact - calls any of the .add* methods + under Artifact depending on the type of object that's passed in. This will probably + be moved to the Artifact class in the future. + + Args: + artifact: `Artifact` - artifact created with `wandb.Artifact(...)` + path_or_obj: `Union[str, ArtifactManifestEntry, data_types.WBValue]` - either a + str or valid object which indicates what to add to an artifact. + + name: `str` - the name of the object which is added to an artifact. + + Returns: + Type[Any] - Union[None, ArtifactManifestEntry, etc] + + """ + if isinstance(path_or_obj, ArtifactManifestEntry): + return artifact.add_reference(path_or_obj, name) + elif isinstance(path_or_obj, data_types.WBValue): + return artifact.add(path_or_obj, name) + elif isinstance(path_or_obj, str): + if os.path.isdir(path_or_obj): + return artifact.add_dir(path_or_obj) + elif os.path.isfile(path_or_obj): + return artifact.add_file(path_or_obj) + else: + with artifact.new_file(name) as f: + f.write(json.dumps(path_or_obj, sort_keys=True)) + else: + raise TypeError( + "Expected `path_or_obj` to be instance of `ArtifactManifestEntry`," + f" `WBValue`, or `str, found {type(path_or_obj)}" + ) + + +def _log_artifact_version( + name: str, + type: str, + entries: Dict[str, Union[str, ArtifactManifestEntry, data_types.WBValue]], + aliases: Optional[Union[str, List[str]]] = None, + description: Optional[str] = None, + metadata: Optional[dict] = None, + project: Optional[str] = None, + scope_project: Optional[bool] = None, + job_type: str = "auto", +) -> Artifact: + """Create an artifact, populate it, and log it with a run. + + If a run is not present, we create one. + + Args: + name: `str` - name of the artifact. If not scoped to a project, name will be + suffixed by "-{run_id}". + type: `str` - type of the artifact, used in the UI to group artifacts of the + same type. + entries: `Dict` - dictionary containing the named objects we want added to this + artifact. + description: `str` - text description of artifact. + metadata: `Dict` - users can pass in artifact-specific metadata here, will be + visible in the UI. + project: `str` - project under which to place this artifact. + scope_project: `bool` - if True, we will not suffix `name` with "-{run_id}". + job_type: `str` - Only applied if run is not present and we create one. + Used to identify runs of a certain job type, i.e "evaluation". + + Returns: + Artifact + + """ + run = wandb_setup.singleton().most_recent_active_run + if not run: + run = wandb.init( + project=project, + job_type=job_type, + settings=wandb.Settings(silent=True), + ) + + if not scope_project: + name = f"{name}-{run.id}" + + if metadata is None: + metadata = {} + + art = wandb.Artifact(name, type, description, metadata, False, None) + + for path in entries: + _add_any(art, entries[path], path) + + # "latest" should always be present as an alias + aliases = wandb.util._resolve_aliases(aliases) + run.log_artifact(art, aliases=aliases) + + return art + + +def log_model( + model_obj: Any, + name: str = "model", + aliases: Optional[Union[str, List[str]]] = None, + description: Optional[str] = None, + metadata: Optional[dict] = None, + project: Optional[str] = None, + scope_project: Optional[bool] = None, + **kwargs: Dict[str, Any], +) -> "_SavedModel": + """Log a model object to enable model-centric workflows in the UI. + + Supported frameworks include PyTorch, Keras, Tensorflow, Scikit-learn, etc. Under + the hood, we create a model artifact, bind it to the run that produced this model, + associate it with the latest metrics logged with `run.log(...)` and more. + + Args: + model_obj: any model object created with the following ML frameworks: PyTorch, + Keras, Tensorflow, Scikit-learn. name: `str` - name of the model artifact + that will be created to house this model_obj. + aliases: `str, List[str]` - optional alias(es) that will be applied on this + model and allow for unique identification. The alias "latest" will always be + applied to the latest version of a model. + description: `str` - text description/notes about the model - will be visible in + the Model Card UI. + metadata: `Dict` - model-specific metadata goes here - will be visible the UI. + project: `str` - project under which to place this artifact. + scope_project: `bool` - If true, name of this model artifact will not be + suffixed by `-{run_id}`. + + Returns: + _SavedModel instance + + Example: + ```python + import torch.nn as nn + import torch.nn.functional as F + + + class Net(nn.Module): + def __init__(self): + super(Net, self).__init__() + self.fc1 = nn.Linear(10, 10) + + def forward(self, x): + x = self.fc1(x) + x = F.relu(x) + return x + + + model = Net() + sm = log_model(model, "my-simple-model", aliases=["best"]) + ``` + + """ + model = data_types._SavedModel.init(model_obj, **kwargs) + _ = _log_artifact_version( + name=name, + type="model", + entries={ + "index": model, + }, + aliases=aliases, + description=description, + metadata=metadata, + project=project, + scope_project=scope_project, + job_type="log_model", + ) + # TODO: handle offline mode appropriately. + return model + + +def use_model(aliased_path: str, unsafe: bool = False) -> "_SavedModel": + """Fetch a saved model from an alias. + + Under the hood, we use the alias to fetch the model artifact containing the + serialized model files and rebuild the model object from these files. We also + declare the fetched model artifact as an input to the run (with `run.use_artifact`). + + Args: + aliased_path: `str` - the following forms are valid: "name:version", + "name:alias". May be prefixed with "entity/project". + unsafe: `bool` - must be True to indicate the user understands the risks + associated with loading external models. + + Returns: + _SavedModel instance + + Example: + ```python + # Assuming the model with the name "my-simple-model" is trusted: + sm = use_model("my-simple-model:latest", unsafe=True) + model = sm.model_obj() + ``` + """ + if not unsafe: + raise ValueError("The 'unsafe' parameter must be set to True to load a model.") + + if ":" not in aliased_path: + raise ValueError( + "aliased_path must be of the form 'name:alias' or 'name:version'." + ) + + # Returns a _SavedModel instance + if run := wandb_setup.singleton().most_recent_active_run: + artifact = run.use_artifact(aliased_path) + sm = artifact.get("index") + + if sm is None or not isinstance(sm, _SavedModel): + raise ValueError( + "Deserialization into model object failed: _SavedModel instance could not be initialized properly." + ) + + return sm + else: + raise ValueError( + "use_model can only be called inside a run. Please call wandb.init() before use_model(...)" + ) + + +def link_model( + model: "_SavedModel", + target_path: str, + aliases: Optional[Union[str, List[str]]] = None, +) -> None: + """Link the given model to a portfolio. + + A portfolio is a promoted collection which contains (in this case) model artifacts. + Linking to a portfolio allows for useful model-centric workflows in the UI. + + Args: + model: `_SavedModel` - an instance of _SavedModel, most likely from the output + of `log_model` or `use_model`. + target_path: `str` - the target portfolio. The following forms are valid for the + string: {portfolio}, {project/portfolio},{entity}/{project}/{portfolio}. + aliases: `str, List[str]` - optional alias(es) that will only be applied on this + linked model inside the portfolio. The alias "latest" will always be applied + to the latest version of a model. + + Returns: + None + + Example: + sm = use_model("my-simple-model:latest") + link_model(sm, "my-portfolio") + + """ + aliases = wandb.util._resolve_aliases(aliases) + + if run := wandb_setup.singleton().most_recent_active_run: + # _artifact_source, if it exists, points to a Public Artifact. + # Its existence means that _SavedModel was deserialized from a logged artifact, most likely from `use_model`. + if model._artifact_source: + artifact = model._artifact_source.artifact + # If the _SavedModel has been added to a Local Artifact (most likely through `.add(WBValue)`), then + # model._artifact_target will point to that Local Artifact. + elif model._artifact_target and model._artifact_target.artifact._final: + artifact = model._artifact_target.artifact + else: + raise ValueError( + "Linking requires that the given _SavedModel belongs to an artifact" + ) + + run.link_artifact(artifact, target_path, aliases) + + else: + if model._artifact_source is not None: + model._artifact_source.artifact.link(target_path, aliases) + else: + raise ValueError( + "Linking requires that the given _SavedModel belongs to a logged artifact." + ) diff --git a/lib/python3.12/site-packages/wandb/docker/__init__.py b/lib/python3.12/site-packages/wandb/docker/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..761748bb8bece16f4b95a25995f6d05ef710748d --- /dev/null +++ b/lib/python3.12/site-packages/wandb/docker/__init__.py @@ -0,0 +1,290 @@ +import json +import logging +import os +import shutil +import subprocess +from typing import Any, Dict, List, Optional, Tuple, Union + +from wandb.docker import names +from wandb.errors import Error + + +class DockerError(Error): + """Raised when attempting to execute a docker command.""" + + def __init__( + self, + command_launched: List[str], + return_code: int, + stdout: Optional[bytes] = None, + stderr: Optional[bytes] = None, + ) -> None: + command_launched_str = " ".join(command_launched) + error_msg = ( + f"The docker command executed was `{command_launched_str}`.\n" + f"It returned with code {return_code}\n" + ) + if stdout is not None: + error_msg += f"The content of stdout is '{stdout.decode()}'\n" + else: + error_msg += ( + "The content of stdout can be found above the " + "stacktrace (it wasn't captured).\n" + ) + if stderr is not None: + error_msg += f"The content of stderr is '{stderr.decode()}'\n" + else: + error_msg += ( + "The content of stderr can be found above the " + "stacktrace (it wasn't captured)." + ) + super().__init__(error_msg) + + +entrypoint = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "wandb-entrypoint.sh" +) +log = logging.getLogger(__name__) + + +def shell(cmd: List[str]) -> Optional[str]: + """Simple wrapper for calling docker,. + + returning None on error and the output on success + """ + try: + return ( + subprocess.check_output(["docker"] + cmd, stderr=subprocess.STDOUT) + .decode("utf8") + .strip() + ) + except subprocess.CalledProcessError as e: + print(e) # noqa: T201 + return None + + +_buildx_installed = None + + +def is_buildx_installed() -> bool: + """Return `True` if docker buildx is installed and working.""" + global _buildx_installed + if _buildx_installed is not None: + return _buildx_installed # type: ignore + if not shutil.which("docker"): + _buildx_installed = False + else: + help_output = shell(["buildx", "--help"]) + _buildx_installed = help_output is not None and "buildx" in help_output + return _buildx_installed + + +def is_docker_installed() -> bool: + """Return `True` if docker is installed and working, else `False`.""" + try: + # Run the docker --version command + result = subprocess.run( + ["docker", "--version"], + capture_output=True, + ) + if result.returncode == 0: + return True + else: + return False + except FileNotFoundError: + # If docker command is not found + return False + + +def build( + tags: List[str], file: str, context_path: str, platform: Optional[str] = None +) -> str: + use_buildx = is_buildx_installed() + command = ["buildx", "build"] if use_buildx else ["build"] + command += ["--load"] if should_add_load_argument(platform) and use_buildx else [] + if platform: + command += ["--platform", platform] + build_tags = [] + for tag in tags: + build_tags += ["-t", tag] + args = ["docker"] + command + build_tags + ["-f", file, context_path] + stdout = run_command_live_output( + args, + ) + return stdout + + +def should_add_load_argument(platform: Optional[str]) -> bool: + # the load option does not work when multiple platforms are specified: + # https://github.com/docker/buildx/issues/59 + if platform is None or (platform and "," not in platform): + return True + return False + + +def run_command_live_output(args: List[Any]) -> str: + with subprocess.Popen( + args, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + universal_newlines=True, + bufsize=1, + ) as process: + stdout = "" + while True: + chunk = os.read(process.stdout.fileno(), 4096) # type: ignore + if not chunk: + break + index = chunk.find(b"\r") + if index != -1: + print(chunk.decode(), end="") # noqa: T201 + else: + stdout += chunk.decode() + print(chunk.decode(), end="\r") # noqa: T201 + + print(stdout) # noqa: T201 + + return_code = process.wait() + if return_code != 0: + raise DockerError(args, return_code, stdout.encode()) + + return stdout + + +def run( + args: List[Any], + capture_stdout: bool = True, + capture_stderr: bool = True, + input: Optional[bytes] = None, + return_stderr: bool = False, + env: Optional[Dict[str, str]] = None, +) -> Union[str, Tuple[str, str]]: + args = [str(x) for x in args] + subprocess_env = dict(os.environ) + subprocess_env.update(env or {}) + if args[1] == "buildx": + subprocess_env["DOCKER_CLI_EXPERIMENTAL"] = "enabled" + stdout_dest: Optional[int] = subprocess.PIPE if capture_stdout else None + stderr_dest: Optional[int] = subprocess.PIPE if capture_stderr else None + + completed_process = subprocess.run( + args, input=input, stdout=stdout_dest, stderr=stderr_dest, env=subprocess_env + ) + if completed_process.returncode != 0: + raise DockerError( + args, + completed_process.returncode, + completed_process.stdout, + completed_process.stderr, + ) + + if return_stderr: + return ( + _post_process_stream(completed_process.stdout), + _post_process_stream(completed_process.stderr), + ) + else: + return _post_process_stream(completed_process.stdout) + + +def _post_process_stream(stream: Optional[bytes]) -> str: + if stream is None: + return "" + decoded_stream = stream.decode() + if len(decoded_stream) != 0 and decoded_stream[-1] == "\n": + decoded_stream = decoded_stream[:-1] + return decoded_stream + + +def default_image(gpu: bool = False) -> str: + tag = "all" + if not gpu: + tag += "-cpu" + return f"wandb/deepo:{tag}" + + +def parse_repository_tag(repo_name: str) -> Tuple[str, Optional[str]]: + parts = repo_name.rsplit("@", 1) + if len(parts) == 2: + return parts[0], parts[1] + parts = repo_name.rsplit(":", 1) + if len(parts) == 2 and "/" not in parts[1]: + return parts[0], parts[1] + return repo_name, None + + +def parse(image_name: str) -> Tuple[str, str, str]: + repository, tag = parse_repository_tag(image_name) + registry, repo_name = names.resolve_repository_name(repository) + if registry == "docker.io": + registry = "index.docker.io" + return registry, repo_name, (tag or "latest") + + +def image_id_from_registry(image_name: str) -> Optional[str]: + """Query the image manifest to get its full ID including the digest. + + Args: + image_name: The image name, such as "wandb/local". + + Returns: + The image name followed by its digest, like "wandb/local@sha256:...". + """ + # https://docs.docker.com/reference/cli/docker/buildx/imagetools/inspect + inspect_cmd = ["buildx", "imagetools", "inspect", image_name] + format_args = ["--format", r"{{.Name}}@{{.Manifest.Digest}}"] + return shell([*inspect_cmd, *format_args]) + + +def image_id(image_name: str) -> Optional[str]: + """Retrieve the image id from the local docker daemon or remote registry.""" + if "@sha256:" in image_name: + return image_name + else: + digests = shell(["inspect", image_name, "--format", "{{json .RepoDigests}}"]) + + if digests is None: + return image_id_from_registry(image_name) + + try: + return json.loads(digests)[0] + except (ValueError, IndexError): + return image_id_from_registry(image_name) + + +def get_image_uid(image_name: str) -> int: + """Retrieve the image default uid through brute force.""" + image_uid = shell(["run", image_name, "id", "-u"]) + return int(image_uid) if image_uid else -1 + + +def push(image: str, tag: str) -> Optional[str]: + """Push an image to a remote registry.""" + return shell(["push", f"{image}:{tag}"]) + + +def login(username: str, password: str, registry: str) -> Optional[str]: + """Login to a registry.""" + return shell(["login", "--username", username, "--password", password, registry]) + + +def tag(image_name: str, tag: str) -> Optional[str]: + """Tag an image.""" + return shell(["tag", image_name, tag]) + + +__all__ = [ + "shell", + "build", + "run", + "image_id", + "image_id_from_registry", + "is_docker_installed", + "parse", + "parse_repository_tag", + "default_image", + "get_image_uid", + "push", + "login", + "tag", +] diff --git a/lib/python3.12/site-packages/wandb/docker/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/docker/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..796d5fc126f627ce2531022b6eebe8b7ed63dde7 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/docker/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/docker/__pycache__/names.cpython-312.pyc b/lib/python3.12/site-packages/wandb/docker/__pycache__/names.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e5f9b11d105b8aa7af85f4f028af78f95de38616 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/docker/__pycache__/names.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/docker/names.py b/lib/python3.12/site-packages/wandb/docker/names.py new file mode 100644 index 0000000000000000000000000000000000000000..384d39a9e21329bcb71abfbcacdcfa9b3c78bbd2 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/docker/names.py @@ -0,0 +1,40 @@ +from __future__ import annotations + + +class InvalidRepositoryError(Exception): + """The given string is not a valid repository name.""" + + +def resolve_repository_name(repo_name: str) -> tuple[str, str]: + if "://" in repo_name: + raise InvalidRepositoryError( + f"Repository name cannot contain a scheme ({repo_name})" + ) + + index_name, remote_name = split_repo_name(repo_name) + if index_name[0] == "-" or index_name[-1] == "-": + raise InvalidRepositoryError( + f"Invalid index name ({index_name}). Cannot begin or end with a hyphen." + ) + return resolve_index_name(index_name), remote_name + + +def resolve_index_name(index_name: str) -> str: + index_name = convert_to_hostname(index_name) + if index_name == "index.docker.io": + index_name = "docker.io" + return index_name + + +def split_repo_name(repo_name: str) -> tuple[str, str]: + parts = repo_name.split("/", 1) + if len(parts) == 1 or ( + "." not in parts[0] and ":" not in parts[0] and parts[0] != "localhost" + ): + # This is a docker index repo (ex: username/foobar or ubuntu) + return "docker.io", repo_name + return parts[0], parts[1] + + +def convert_to_hostname(url: str) -> str: + return url.replace("http://", "").replace("https://", "").split("/", 1)[0] diff --git a/lib/python3.12/site-packages/wandb/docker/wandb-entrypoint.sh b/lib/python3.12/site-packages/wandb/docker/wandb-entrypoint.sh new file mode 100644 index 0000000000000000000000000000000000000000..f2d4af61de2eb53973f8ee5eb6e619521d66a63a --- /dev/null +++ b/lib/python3.12/site-packages/wandb/docker/wandb-entrypoint.sh @@ -0,0 +1,33 @@ +#!/bin/sh +set -e + +wandb="\x1b[34m\x1b[1mwandb\x1b[0m" +/bin/echo -e "${wandb}: Checking image for required packages." + +if ! [ -x "$(command -v python)" ]; then + /bin/echo -e "${wandb}: python not installed, can't use wandb with this image." + exit 1 +fi + +if ! [ -x "$(command -v wandb)" ]; then + /bin/echo -e "${wandb}: wandb not installed, installing." + pip install wandb --upgrade +else + ver=$(wandb --version) + /bin/echo -e "${wandb}: Found $ver" +fi + +if [ "$WANDB_ENSURE_JUPYTER" = "1" ]; then + if ! [ -x "$(command -v jupyter-lab)" ]; then + /bin/echo -e "${wandb}: jupyter not installed, installing." + pip install jupyterlab + /bin/echo -e "${wandb}: starting jupyter, you can access it at: http://127.0.0.1:8888" + fi +fi + +if ! [ -z "$WANDB_COMMAND" ]; then + /bin/echo $WANDB_COMMAND >> ~/.bash_history + /bin/echo -e "${wandb}: Command added to history, press up arrow to access it." + /bin/echo -e "${wandb}: $WANDB_COMMAND" +fi +exec "$@" \ No newline at end of file diff --git a/lib/python3.12/site-packages/wandb/filesync/__init__.py b/lib/python3.12/site-packages/wandb/filesync/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/wandb/filesync/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/filesync/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dc9bb33bbb609bda3d66597f78b2158af6d058a6 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/filesync/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/filesync/__pycache__/dir_watcher.cpython-312.pyc b/lib/python3.12/site-packages/wandb/filesync/__pycache__/dir_watcher.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fa344884cc48db7bb37935683f6b27d9b53e02c6 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/filesync/__pycache__/dir_watcher.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/filesync/__pycache__/stats.cpython-312.pyc b/lib/python3.12/site-packages/wandb/filesync/__pycache__/stats.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..49b73ae93738ba7d428c761117d855747d318b01 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/filesync/__pycache__/stats.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/filesync/__pycache__/step_checksum.cpython-312.pyc b/lib/python3.12/site-packages/wandb/filesync/__pycache__/step_checksum.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..69fb79ee8f0f4a88a6453bd226e0110449bc7243 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/filesync/__pycache__/step_checksum.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/filesync/__pycache__/step_prepare.cpython-312.pyc b/lib/python3.12/site-packages/wandb/filesync/__pycache__/step_prepare.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ca6fd8b407d4a12cddcdfe733fc2b312d028ac53 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/filesync/__pycache__/step_prepare.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/filesync/__pycache__/step_upload.cpython-312.pyc b/lib/python3.12/site-packages/wandb/filesync/__pycache__/step_upload.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..459b681f9f536de2dc67157cf63ad82786c5cda1 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/filesync/__pycache__/step_upload.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/filesync/__pycache__/upload_job.cpython-312.pyc b/lib/python3.12/site-packages/wandb/filesync/__pycache__/upload_job.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9aa5d957b8f42a78d27f386255ff73851ca8a159 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/filesync/__pycache__/upload_job.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/filesync/dir_watcher.py b/lib/python3.12/site-packages/wandb/filesync/dir_watcher.py new file mode 100644 index 0000000000000000000000000000000000000000..4a423fc8a88f3794489752e84f969f434e300a10 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/filesync/dir_watcher.py @@ -0,0 +1,404 @@ +import abc +import fnmatch +import glob +import logging +import os +import queue +import time +from typing import TYPE_CHECKING, Any, Mapping, MutableMapping, MutableSet, Optional + +from wandb import util +from wandb.sdk.interface.interface import GlobStr +from wandb.sdk.lib.paths import LogicalPath + +if TYPE_CHECKING: + import wandb.vendor.watchdog_0_9_0.observers.api as wd_api + import wandb.vendor.watchdog_0_9_0.observers.polling as wd_polling + import wandb.vendor.watchdog_0_9_0.watchdog.events as wd_events + from wandb.sdk.interface.interface import PolicyName + from wandb.sdk.internal.file_pusher import FilePusher + from wandb.sdk.internal.settings_static import SettingsStatic +else: + wd_polling = util.vendor_import("wandb_watchdog.observers.polling") + wd_events = util.vendor_import("wandb_watchdog.events") + +PathStr = str # TODO(spencerpearson): would be nice to use Path here + + +logger = logging.getLogger(__name__) + + +class FileEventHandler(abc.ABC): + def __init__( + self, + file_path: PathStr, + save_name: LogicalPath, + file_pusher: "FilePusher", + *args: Any, + **kwargs: Any, + ) -> None: + self.file_path = file_path + # Convert windows paths to unix paths + self.save_name = LogicalPath(save_name) + self._file_pusher = file_pusher + self._last_sync: Optional[float] = None + + @property + @abc.abstractmethod + def policy(self) -> "PolicyName": + raise NotImplementedError + + @abc.abstractmethod + def on_modified(self, force: bool = False) -> None: + raise NotImplementedError + + @abc.abstractmethod + def finish(self) -> None: + raise NotImplementedError + + def on_renamed(self, new_path: PathStr, new_name: LogicalPath) -> None: + self.file_path = new_path + self.save_name = new_name + self.on_modified() + + +class PolicyNow(FileEventHandler): + """This policy only uploads files now.""" + + def on_modified(self, force: bool = False) -> None: + # only upload if we've never uploaded or when .save is called + if self._last_sync is None or force: + self._file_pusher.file_changed(self.save_name, self.file_path) + self._last_sync = os.path.getmtime(self.file_path) + + def finish(self) -> None: + pass + + @property + def policy(self) -> "PolicyName": + return "now" + + +class PolicyEnd(FileEventHandler): + """This policy only updates at the end of the run.""" + + def on_modified(self, force: bool = False) -> None: + pass + + # TODO: make sure we call this + def finish(self) -> None: + # We use copy=False to avoid possibly expensive copies, and because + # user files shouldn't still be changing at the end of the run. + self._last_sync = os.path.getmtime(self.file_path) + self._file_pusher.file_changed(self.save_name, self.file_path, copy=False) + + @property + def policy(self) -> "PolicyName": + return "end" + + +class PolicyLive(FileEventHandler): + """Event handler that uploads respecting throttling. + + Uploads files every RATE_LIMIT_SECONDS, which changes as the size increases to deal + with throttling. + """ + + RATE_LIMIT_SECONDS = 15 + unit_dict = dict(util.POW_10_BYTES) + # Wait to upload until size has increased 20% from last upload + RATE_LIMIT_SIZE_INCREASE = 1.2 + + def __init__( + self, + file_path: PathStr, + save_name: LogicalPath, + file_pusher: "FilePusher", + settings: Optional["SettingsStatic"] = None, + *args: Any, + **kwargs: Any, + ) -> None: + super().__init__(file_path, save_name, file_pusher, *args, **kwargs) + self._last_uploaded_time: Optional[float] = None + self._last_uploaded_size: int = 0 + if settings is not None: + if settings.x_live_policy_rate_limit is not None: + self.RATE_LIMIT_SECONDS = settings.x_live_policy_rate_limit + self._min_wait_time: Optional[float] = settings.x_live_policy_wait_time + else: + self._min_wait_time = None + + @property + def current_size(self) -> int: + return os.path.getsize(self.file_path) + + @classmethod + def min_wait_for_size(cls, size: int) -> float: + if size < 10 * cls.unit_dict["MB"]: + return 60 + elif size < 100 * cls.unit_dict["MB"]: + return 5 * 60 + elif size < cls.unit_dict["GB"]: + return 10 * 60 + else: + return 20 * 60 + + def should_update(self) -> bool: + if self._last_uploaded_time is not None: + # Check rate limit by time elapsed + time_elapsed = time.time() - self._last_uploaded_time + # if more than 15 seconds has passed potentially upload it + if time_elapsed < self.RATE_LIMIT_SECONDS: + return False + + # Check rate limit by size increase + if float(self._last_uploaded_size) > 0: + size_increase = self.current_size / float(self._last_uploaded_size) + if size_increase < self.RATE_LIMIT_SIZE_INCREASE: + return False + return time_elapsed > ( + self._min_wait_time or self.min_wait_for_size(self.current_size) + ) + + # if the file has never been uploaded, we'll upload it + return True + + def on_modified(self, force: bool = False) -> None: + if self.current_size == 0: + return + if self._last_sync == os.path.getmtime(self.file_path): + return + if force or self.should_update(): + self.save_file() + + def save_file(self) -> None: + self._last_sync = os.path.getmtime(self.file_path) + self._last_uploaded_time = time.time() + self._last_uploaded_size = self.current_size + self._file_pusher.file_changed(self.save_name, self.file_path) + + def finish(self) -> None: + self.on_modified(force=True) + + @property + def policy(self) -> "PolicyName": + return "live" + + +class DirWatcher: + def __init__( + self, + settings: "SettingsStatic", + file_pusher: "FilePusher", + file_dir: Optional[PathStr] = None, + ) -> None: + self._file_count = 0 + self._dir = file_dir or settings.files_dir + self._settings = settings + self._savename_file_policies: MutableMapping[LogicalPath, PolicyName] = {} + self._user_file_policies: Mapping[PolicyName, MutableSet[GlobStr]] = { + "end": set(), + "live": set(), + "now": set(), + } + self._file_pusher = file_pusher + self._file_event_handlers: MutableMapping[LogicalPath, FileEventHandler] = {} + self._file_observer = wd_polling.PollingObserver() + self._file_observer.schedule( + self._per_file_event_handler(), self._dir, recursive=True + ) + self._file_observer.start() + logger.info("watching files in: %s", settings.files_dir) + + @property + def emitter(self) -> Optional["wd_api.EventEmitter"]: + try: + return next(iter(self._file_observer.emitters)) + except StopIteration: + return None + + def update_policy(self, path: GlobStr, policy: "PolicyName") -> None: + # When we're dealing with one of our own media files, there's no need + # to store the policy in memory. _get_file_event_handler will always + # return PolicyNow. Using the path makes syncing historic runs much + # faster if the name happens to include glob escapable characters. In + # the future we may add a flag to "files" records that indicates it's + # policy is not dynamic and doesn't need to be stored / checked. + save_name = LogicalPath( + os.path.relpath(os.path.join(self._dir, path), self._dir) + ) + if save_name.startswith("media/"): + pass + elif path == glob.escape(path): + self._savename_file_policies[save_name] = policy + else: + self._user_file_policies[policy].add(path) + + for src_path in glob.glob(os.path.join(self._dir, path)): + save_name = LogicalPath(os.path.relpath(src_path, self._dir)) + feh = self._get_file_event_handler(src_path, save_name) + # handle the case where the policy changed + if feh.policy != policy: + try: + del self._file_event_handlers[save_name] + except KeyError: + # TODO: probably should do locking, but this handles moved files for now + pass + feh = self._get_file_event_handler(src_path, save_name) + feh.on_modified(force=True) + + def _per_file_event_handler(self) -> "wd_events.FileSystemEventHandler": + """Create a Watchdog file event handler that does different things for every file.""" + file_event_handler = wd_events.PatternMatchingEventHandler() + file_event_handler.on_created = self._on_file_created + file_event_handler.on_modified = self._on_file_modified + file_event_handler.on_moved = self._on_file_moved + file_event_handler._patterns = [os.path.join(self._dir, os.path.normpath("*"))] + # Ignore hidden files/folders + # TODO: what other files should we skip? + file_event_handler._ignore_patterns = [ + "*.tmp", + "*.wandb", + "wandb-summary.json", + os.path.join(self._dir, ".*"), + os.path.join(self._dir, "*/.*"), + ] + for glb in self._settings.ignore_globs: + file_event_handler._ignore_patterns.append(os.path.join(self._dir, glb)) + + return file_event_handler + + def _on_file_created(self, event: "wd_events.FileCreatedEvent") -> None: + logger.info("file/dir created: %s", event.src_path) + if os.path.isdir(event.src_path): + return None + self._file_count += 1 + # We do the directory scan less often as it grows + if self._file_count % 100 == 0: + emitter = self.emitter + if emitter: + emitter._timeout = int(self._file_count / 100) + 1 + save_name = LogicalPath(os.path.relpath(event.src_path, self._dir)) + self._get_file_event_handler(event.src_path, save_name).on_modified() + + # TODO(spencerpearson): this pattern repeats so many times we should have a method/function for it + # def _save_name(self, path: PathStr) -> LogicalPath: + # return LogicalPath(os.path.relpath(path, self._dir)) + + def _on_file_modified(self, event: "wd_events.FileModifiedEvent") -> None: + logger.info(f"file/dir modified: {event.src_path}") + if os.path.isdir(event.src_path): + return None + save_name = LogicalPath(os.path.relpath(event.src_path, self._dir)) + self._get_file_event_handler(event.src_path, save_name).on_modified() + + def _on_file_moved(self, event: "wd_events.FileMovedEvent") -> None: + # TODO: test me... + logger.info(f"file/dir moved: {event.src_path} -> {event.dest_path}") + if os.path.isdir(event.dest_path): + return None + old_save_name = LogicalPath(os.path.relpath(event.src_path, self._dir)) + new_save_name = LogicalPath(os.path.relpath(event.dest_path, self._dir)) + + # We have to move the existing file handler to the new name + handler = self._get_file_event_handler(event.src_path, old_save_name) + self._file_event_handlers[new_save_name] = handler + del self._file_event_handlers[old_save_name] + + handler.on_renamed(event.dest_path, new_save_name) + + def _get_file_event_handler( + self, file_path: PathStr, save_name: LogicalPath + ) -> FileEventHandler: + """Get or create an event handler for a particular file. + + file_path: the file's actual path + save_name: its path relative to the run directory (aka the watch directory) + """ + # Always return PolicyNow for any of our media files. + if save_name.startswith("media/"): + return PolicyNow(file_path, save_name, self._file_pusher, self._settings) + if save_name not in self._file_event_handlers: + # TODO: we can use PolicyIgnore if there are files we never want to sync + if "tfevents" in save_name or "graph.pbtxt" in save_name: + self._file_event_handlers[save_name] = PolicyLive( + file_path, save_name, self._file_pusher, self._settings + ) + elif save_name in self._savename_file_policies: + policy_name = self._savename_file_policies[save_name] + make_handler = ( + PolicyLive + if policy_name == "live" + else PolicyNow + if policy_name == "now" + else PolicyEnd + ) + self._file_event_handlers[save_name] = make_handler( + file_path, save_name, self._file_pusher, self._settings + ) + else: + make_handler = PolicyEnd + for policy, globs in self._user_file_policies.items(): + if policy == "end": + continue + # Convert set to list to avoid RuntimeError's + # TODO: we may need to add locks + for g in list(globs): + paths = glob.glob(os.path.join(self._dir, g)) + if any(save_name in p for p in paths): + if policy == "live": + make_handler = PolicyLive + elif policy == "now": + make_handler = PolicyNow + self._file_event_handlers[save_name] = make_handler( + file_path, save_name, self._file_pusher, self._settings + ) + return self._file_event_handlers[save_name] + + def finish(self) -> None: + logger.info("shutting down directory watcher") + try: + # avoid hanging if we crashed before the observer was started + if self._file_observer.is_alive(): + # rather unfortunately we need to manually do a final scan of the dir + # with `queue_events`, then iterate through all events before stopping + # the observer to catch all files written. First we need to prevent the + # existing thread from consuming our final events, then we process them + self._file_observer._timeout = 0 + self._file_observer._stopped_event.set() + self._file_observer.join() + self.emitter.queue_events(0) # type: ignore[union-attr] + while True: + try: + self._file_observer.dispatch_events( + self._file_observer.event_queue, 0 + ) + except queue.Empty: + break + # Calling stop unschedules any inflight events so we handled them above + self._file_observer.stop() + # TODO: py2 TypeError: PyCObject_AsVoidPtr called with null pointer + except TypeError: + pass + # TODO: py3 SystemError: returned an error + except SystemError: + pass + + # Ensure we've at least noticed every file in the run directory. Sometimes + # we miss things because asynchronously watching filesystems isn't reliable. + logger.info("scan: %s", self._dir) + + for dirpath, _, filenames in os.walk(self._dir): + for fname in filenames: + file_path = os.path.join(dirpath, fname) + save_name = LogicalPath(os.path.relpath(file_path, self._dir)) + ignored = False + for glb in self._settings.ignore_globs: + if len(fnmatch.filter([save_name], glb)) > 0: + ignored = True + logger.info("ignored: %s matching glob %s", save_name, glb) + break + if ignored: + continue + logger.info("scan save: %s %s", file_path, save_name) + self._get_file_event_handler(file_path, save_name).finish() diff --git a/lib/python3.12/site-packages/wandb/filesync/stats.py b/lib/python3.12/site-packages/wandb/filesync/stats.py new file mode 100644 index 0000000000000000000000000000000000000000..95c3523a37a2a82ce17bd40a6b4a667db04ff0ff --- /dev/null +++ b/lib/python3.12/site-packages/wandb/filesync/stats.py @@ -0,0 +1,100 @@ +import threading +from typing import MutableMapping, NamedTuple + +import wandb + + +class FileStats(NamedTuple): + deduped: bool + total: int + uploaded: int + failed: bool + artifact_file: bool + + +class Summary(NamedTuple): + uploaded_bytes: int + total_bytes: int + deduped_bytes: int + + +class FileCountsByCategory(NamedTuple): + artifact: int + wandb: int + media: int + other: int + + +class Stats: + def __init__(self) -> None: + self._stats: MutableMapping[str, FileStats] = {} + self._lock = threading.Lock() + + def init_file( + self, save_name: str, size: int, is_artifact_file: bool = False + ) -> None: + with self._lock: + self._stats[save_name] = FileStats( + deduped=False, + total=size, + uploaded=0, + failed=False, + artifact_file=is_artifact_file, + ) + + def set_file_deduped(self, save_name: str) -> None: + with self._lock: + orig = self._stats[save_name] + self._stats[save_name] = orig._replace( + deduped=True, + uploaded=orig.total, + ) + + def update_uploaded_file(self, save_name: str, total_uploaded: int) -> None: + with self._lock: + self._stats[save_name] = self._stats[save_name]._replace( + uploaded=total_uploaded, + ) + + def update_failed_file(self, save_name: str) -> None: + with self._lock: + self._stats[save_name] = self._stats[save_name]._replace( + uploaded=0, + failed=True, + ) + + def summary(self) -> Summary: + # Need to use list to ensure we get a copy, since other threads may + # modify this while we iterate + with self._lock: + stats = list(self._stats.values()) + return Summary( + uploaded_bytes=sum(f.uploaded for f in stats), + total_bytes=sum(f.total for f in stats), + deduped_bytes=sum(f.total for f in stats if f.deduped), + ) + + def file_counts_by_category(self) -> FileCountsByCategory: + artifact_files = 0 + wandb_files = 0 + media_files = 0 + other_files = 0 + # Need to use list to ensure we get a copy, since other threads may + # modify this while we iterate + with self._lock: + file_stats = list(self._stats.items()) + for save_name, stats in file_stats: + if stats.artifact_file: + artifact_files += 1 + elif wandb.wandb_lib.filenames.is_wandb_file(save_name): # type: ignore[attr-defined] # TODO(spencerpearson): this is probably synonymous with wandb.sdk.lib.filenames...? + wandb_files += 1 + elif save_name.startswith("media"): + media_files += 1 + else: + other_files += 1 + return FileCountsByCategory( + artifact=artifact_files, + wandb=wandb_files, + media=media_files, + other=other_files, + ) diff --git a/lib/python3.12/site-packages/wandb/filesync/step_checksum.py b/lib/python3.12/site-packages/wandb/filesync/step_checksum.py new file mode 100644 index 0000000000000000000000000000000000000000..c0acd96e80e8ad7b4fb8027da718a3b28780f86e --- /dev/null +++ b/lib/python3.12/site-packages/wandb/filesync/step_checksum.py @@ -0,0 +1,142 @@ +"""Batching file prepare requests to our API.""" + +import concurrent.futures +import functools +import os +import queue +import shutil +import threading +from typing import TYPE_CHECKING, NamedTuple, Optional, Union, cast + +from wandb.filesync import step_upload +from wandb.sdk.lib import filesystem, runid +from wandb.sdk.lib.paths import LogicalPath + +if TYPE_CHECKING: + import tempfile + + from wandb.filesync import stats + from wandb.sdk.artifacts.artifact_manifest import ArtifactManifest + from wandb.sdk.artifacts.artifact_saver import SaveFn + from wandb.sdk.internal import internal_api + + +class RequestUpload(NamedTuple): + path: str + save_name: LogicalPath + copy: bool + + +class RequestStoreManifestFiles(NamedTuple): + manifest: "ArtifactManifest" + artifact_id: str + save_fn: "SaveFn" + + +class RequestCommitArtifact(NamedTuple): + artifact_id: str + finalize: bool + before_commit: step_upload.PreCommitFn + result_future: "concurrent.futures.Future[None]" + + +class RequestFinish(NamedTuple): + callback: Optional[step_upload.OnRequestFinishFn] + + +Event = Union[ + RequestUpload, RequestStoreManifestFiles, RequestCommitArtifact, RequestFinish +] + + +class StepChecksum: + def __init__( + self, + api: "internal_api.Api", + tempdir: "tempfile.TemporaryDirectory", + request_queue: "queue.Queue[Event]", + output_queue: "queue.Queue[step_upload.Event]", + stats: "stats.Stats", + ) -> None: + self._api = api + self._tempdir = tempdir + self._request_queue = request_queue + self._output_queue = output_queue + self._stats = stats + + self._thread = threading.Thread(target=self._thread_body) + self._thread.daemon = True + + def _thread_body(self) -> None: + while True: + req = self._request_queue.get() + if isinstance(req, RequestUpload): + path = req.path + if req.copy: + path = os.path.join( + self._tempdir.name, + f"{runid.generate_id()}-{req.save_name}", + ) + filesystem.mkdir_exists_ok(os.path.dirname(path)) + try: + # certain linux distros throw an exception when copying + # large files: https://bugs.python.org/issue43743 + shutil.copy2(req.path, path) + except OSError: + shutil._USE_CP_SENDFILE = False # type: ignore[attr-defined] + shutil.copy2(req.path, path) + self._stats.init_file(req.save_name, os.path.getsize(path)) + self._output_queue.put( + step_upload.RequestUpload( + path, + req.save_name, + None, + None, + req.copy, + None, + None, + ) + ) + elif isinstance(req, RequestStoreManifestFiles): + for entry in req.manifest.entries.values(): + if entry.local_path: + self._stats.init_file( + entry.local_path, + cast(int, entry.size), + is_artifact_file=True, + ) + self._output_queue.put( + step_upload.RequestUpload( + entry.local_path, + entry.path, + req.artifact_id, + entry.digest, + False, + functools.partial(req.save_fn, entry), + entry.digest, + ) + ) + elif isinstance(req, RequestCommitArtifact): + self._output_queue.put( + step_upload.RequestCommitArtifact( + req.artifact_id, + req.finalize, + req.before_commit, + req.result_future, + ) + ) + elif isinstance(req, RequestFinish): + break + else: + raise TypeError + + self._output_queue.put(step_upload.RequestFinish(req.callback)) + + def start(self) -> None: + self._thread.start() + + def is_alive(self) -> bool: + return self._thread.is_alive() + + def finish(self) -> None: + self._request_queue.put(RequestFinish(None)) diff --git a/lib/python3.12/site-packages/wandb/filesync/step_prepare.py b/lib/python3.12/site-packages/wandb/filesync/step_prepare.py new file mode 100644 index 0000000000000000000000000000000000000000..95a4a21b4593a62dd21d080e847c212eef7e2ca7 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/filesync/step_prepare.py @@ -0,0 +1,179 @@ +"""Batching file prepare requests to our API.""" + +import queue +import threading +import time +from typing import ( + TYPE_CHECKING, + Callable, + Dict, + List, + Mapping, + NamedTuple, + Optional, + Sequence, + Tuple, + Union, +) + +if TYPE_CHECKING: + from wandb.sdk.internal.internal_api import ( + Api, + CreateArtifactFileSpecInput, + CreateArtifactFilesResponseFile, + ) + + +# Request for a file to be prepared. +class RequestPrepare(NamedTuple): + file_spec: "CreateArtifactFileSpecInput" + response_channel: "queue.Queue[ResponsePrepare]" + + +class RequestFinish(NamedTuple): + pass + + +class ResponsePrepare(NamedTuple): + birth_artifact_id: str + upload_url: Optional[str] + upload_headers: Sequence[str] + upload_id: Optional[str] + storage_path: Optional[str] + multipart_upload_urls: Optional[Dict[int, str]] + + +Request = Union[RequestPrepare, RequestFinish] + + +def _clamp(x: float, low: float, high: float) -> float: + return max(low, min(x, high)) + + +def gather_batch( + request_queue: "queue.Queue[Request]", + batch_time: float, + inter_event_time: float, + max_batch_size: int, + clock: Callable[[], float] = time.monotonic, +) -> Tuple[bool, Sequence[RequestPrepare]]: + batch_start_time = clock() + remaining_time = batch_time + + first_request = request_queue.get() + if isinstance(first_request, RequestFinish): + return True, [] + + batch: List[RequestPrepare] = [first_request] + + while remaining_time > 0 and len(batch) < max_batch_size: + try: + request = request_queue.get( + timeout=_clamp( + x=inter_event_time, + low=1e-12, # 0 = "block forever", so just use something tiny + high=remaining_time, + ), + ) + if isinstance(request, RequestFinish): + return True, batch + + batch.append(request) + remaining_time = batch_time - (clock() - batch_start_time) + + except queue.Empty: + break + + return False, batch + + +def prepare_response(response: "CreateArtifactFilesResponseFile") -> ResponsePrepare: + multipart_resp = response.get("uploadMultipartUrls") + part_list = multipart_resp["uploadUrlParts"] if multipart_resp else [] + multipart_parts = {u["partNumber"]: u["uploadUrl"] for u in part_list} or None + + return ResponsePrepare( + birth_artifact_id=response["artifact"]["id"], + upload_url=response["uploadUrl"], + upload_headers=response["uploadHeaders"], + upload_id=multipart_resp and multipart_resp.get("uploadID"), + storage_path=response.get("storagePath"), + multipart_upload_urls=multipart_parts, + ) + + +class StepPrepare: + """A thread that batches requests to our file prepare API. + + Any number of threads may call prepare() in parallel. The PrepareBatcher thread + will batch requests up and send them all to the backend at once. + """ + + def __init__( + self, + api: "Api", + batch_time: float, + inter_event_time: float, + max_batch_size: int, + request_queue: Optional["queue.Queue[Request]"] = None, + ) -> None: + self._api = api + self._inter_event_time = inter_event_time + self._batch_time = batch_time + self._max_batch_size = max_batch_size + self._request_queue: queue.Queue[Request] = request_queue or queue.Queue() + self._thread = threading.Thread(target=self._thread_body) + self._thread.daemon = True + + def _thread_body(self) -> None: + while True: + finish, batch = gather_batch( + request_queue=self._request_queue, + batch_time=self._batch_time, + inter_event_time=self._inter_event_time, + max_batch_size=self._max_batch_size, + ) + if batch: + batch_response = self._prepare_batch(batch) + # send responses + for prepare_request in batch: + name = prepare_request.file_spec["name"] + response_file = batch_response[name] + response = prepare_response(response_file) + prepare_request.response_channel.put(response) + if finish: + break + + def _prepare_batch( + self, batch: Sequence[RequestPrepare] + ) -> Mapping[str, "CreateArtifactFilesResponseFile"]: + """Execute the prepareFiles API call. + + Args: + batch: List of RequestPrepare objects + Returns: + dict of (save_name: ResponseFile) pairs where ResponseFile is a dict with + an uploadUrl key. The value of the uploadUrl key is None if the file + already exists, or a url string if the file should be uploaded. + """ + return self._api.create_artifact_files([req.file_spec for req in batch]) + + def prepare( + self, file_spec: "CreateArtifactFileSpecInput" + ) -> "queue.Queue[ResponsePrepare]": + response_queue: queue.Queue[ResponsePrepare] = queue.Queue() + self._request_queue.put(RequestPrepare(file_spec, response_queue)) + return response_queue + + def start(self) -> None: + self._thread.start() + + def finish(self) -> None: + self._request_queue.put(RequestFinish()) + + def is_alive(self) -> bool: + return self._thread.is_alive() + + def shutdown(self) -> None: + self.finish() + self._thread.join() diff --git a/lib/python3.12/site-packages/wandb/filesync/step_upload.py b/lib/python3.12/site-packages/wandb/filesync/step_upload.py new file mode 100644 index 0000000000000000000000000000000000000000..0840293d35aaf53aeaf8d211783f9be957beafe2 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/filesync/step_upload.py @@ -0,0 +1,287 @@ +"""Batching file prepare requests to our API.""" + +import concurrent.futures +import logging +import queue +import sys +import threading +from typing import ( + TYPE_CHECKING, + Callable, + MutableMapping, + MutableSequence, + MutableSet, + NamedTuple, + Optional, + Union, +) + +from wandb.errors.term import termerror +from wandb.filesync import upload_job +from wandb.sdk.lib.paths import LogicalPath + +if TYPE_CHECKING: + from typing import TypedDict + + from wandb.filesync import stats + from wandb.sdk.internal import file_stream, internal_api, progress + from wandb.sdk.internal.settings_static import SettingsStatic + + class ArtifactStatus(TypedDict): + finalize: bool + pending_count: int + commit_requested: bool + pre_commit_callbacks: MutableSet["PreCommitFn"] + result_futures: MutableSet["concurrent.futures.Future[None]"] + + +PreCommitFn = Callable[[], None] +OnRequestFinishFn = Callable[[], None] +SaveFn = Callable[["progress.ProgressFn"], bool] + +logger = logging.getLogger(__name__) + + +class RequestUpload(NamedTuple): + path: str + save_name: LogicalPath + artifact_id: Optional[str] + md5: Optional[str] + copied: bool + save_fn: Optional[SaveFn] + digest: Optional[str] + + +class RequestCommitArtifact(NamedTuple): + artifact_id: str + finalize: bool + before_commit: PreCommitFn + result_future: "concurrent.futures.Future[None]" + + +class RequestFinish(NamedTuple): + callback: Optional[OnRequestFinishFn] + + +class EventJobDone(NamedTuple): + job: RequestUpload + exc: Optional[BaseException] + + +Event = Union[RequestUpload, RequestCommitArtifact, RequestFinish, EventJobDone] + + +class StepUpload: + def __init__( + self, + api: "internal_api.Api", + stats: "stats.Stats", + event_queue: "queue.Queue[Event]", + max_threads: int, + file_stream: "file_stream.FileStreamApi", + settings: Optional["SettingsStatic"] = None, + ) -> None: + self._api = api + self._stats = stats + self._event_queue = event_queue + self._file_stream = file_stream + + self._thread = threading.Thread(target=self._thread_body) + self._thread.daemon = True + + self._pool = concurrent.futures.ThreadPoolExecutor( + thread_name_prefix="wandb-upload", + max_workers=max_threads, + ) + + # Indexed by files' `save_name`'s, which are their ID's in the Run. + self._running_jobs: MutableMapping[LogicalPath, RequestUpload] = {} + self._pending_jobs: MutableSequence[RequestUpload] = [] + + self._artifacts: MutableMapping[str, ArtifactStatus] = {} + + self.silent = bool(settings.silent) if settings else False + + def _thread_body(self) -> None: + event: Optional[Event] + # Wait for event in the queue, and process one by one until a + # finish event is received + finish_callback = None + while True: + event = self._event_queue.get() + if isinstance(event, RequestFinish): + finish_callback = event.callback + break + self._handle_event(event) + + # We've received a finish event. At this point, further Upload requests + # are invalid. + + # After a finish event is received, iterate through the event queue + # one by one and process all remaining events. + while True: + try: + event = self._event_queue.get(True, 0.2) + except queue.Empty: + event = None + if event: + self._handle_event(event) + elif not self._running_jobs: + # Queue was empty and no jobs left. + self._pool.shutdown(wait=False) + if finish_callback: + finish_callback() + break + + def _handle_event(self, event: Event) -> None: + if isinstance(event, EventJobDone): + job = event.job + + if event.exc is not None: + logger.exception( + "Failed to upload file: %s", job.path, exc_info=event.exc + ) + + if job.artifact_id: + if event.exc is None: + self._artifacts[job.artifact_id]["pending_count"] -= 1 + self._maybe_commit_artifact(job.artifact_id) + else: + if not self.silent: + termerror( + "Uploading artifact file failed. Artifact won't be committed." + ) + self._fail_artifact_futures(job.artifact_id, event.exc) + self._running_jobs.pop(job.save_name) + # If we have any pending jobs, start one now + if self._pending_jobs: + event = self._pending_jobs.pop(0) + self._start_upload_job(event) + elif isinstance(event, RequestCommitArtifact): + if event.artifact_id not in self._artifacts: + self._init_artifact(event.artifact_id) + self._artifacts[event.artifact_id]["commit_requested"] = True + self._artifacts[event.artifact_id]["finalize"] = event.finalize + self._artifacts[event.artifact_id]["pre_commit_callbacks"].add( + event.before_commit + ) + self._artifacts[event.artifact_id]["result_futures"].add( + event.result_future + ) + self._maybe_commit_artifact(event.artifact_id) + elif isinstance(event, RequestUpload): + if event.artifact_id is not None: + if event.artifact_id not in self._artifacts: + self._init_artifact(event.artifact_id) + self._artifacts[event.artifact_id]["pending_count"] += 1 + self._start_upload_job(event) + else: + raise TypeError(f"Event has unexpected type: {event!s}") + + def _start_upload_job(self, event: RequestUpload) -> None: + # Operations on a single backend file must be serialized. if + # we're already uploading this file, put the event on the + # end of the queue + if event.save_name in self._running_jobs: + self._pending_jobs.append(event) + return + + self._spawn_upload(event) + + def _spawn_upload(self, event: RequestUpload) -> None: + """Spawn an upload job, and handles the bookkeeping of `self._running_jobs`. + + Context: it's important that, whenever we add an entry to `self._running_jobs`, + we ensure that a corresponding `EventJobDone` message will eventually get handled; + otherwise, the `_running_jobs` entry will never get removed, and the StepUpload + will never shut down. + + The sole purpose of this function is to make sure that the code that adds an entry + to `self._running_jobs` is textually right next to the code that eventually enqueues + the `EventJobDone` message. This should help keep them in sync. + """ + # Adding the entry to `self._running_jobs` MUST happen in the main thread, + # NOT in the job that gets submitted to the thread-pool, to guard against + # this sequence of events: + # - StepUpload receives a RequestUpload + # ...and therefore spawns a thread to do the upload + # - StepUpload receives a RequestFinish + # ...and checks `self._running_jobs` to see if there are any tasks to wait for... + # ...and there are none, because the addition to `self._running_jobs` happens in + # the background thread, which the scheduler hasn't yet run... + # ...so the StepUpload shuts down. Even though we haven't uploaded the file! + # + # This would be very bad! + # So, this line has to happen _outside_ the `pool.submit()`. + self._running_jobs[event.save_name] = event + + def run_and_notify() -> None: + try: + self._do_upload(event) + finally: + self._event_queue.put(EventJobDone(event, exc=sys.exc_info()[1])) + + self._pool.submit(run_and_notify) + + def _do_upload(self, event: RequestUpload) -> None: + job = upload_job.UploadJob( + self._stats, + self._api, + self._file_stream, + self.silent, + event.save_name, + event.path, + event.artifact_id, + event.md5, + event.copied, + event.save_fn, + event.digest, + ) + job.run() + + def _init_artifact(self, artifact_id: str) -> None: + self._artifacts[artifact_id] = { + "finalize": False, + "pending_count": 0, + "commit_requested": False, + "pre_commit_callbacks": set(), + "result_futures": set(), + } + + def _maybe_commit_artifact(self, artifact_id: str) -> None: + artifact_status = self._artifacts[artifact_id] + if ( + artifact_status["pending_count"] == 0 + and artifact_status["commit_requested"] + ): + try: + for pre_callback in artifact_status["pre_commit_callbacks"]: + pre_callback() + if artifact_status["finalize"]: + self._api.commit_artifact(artifact_id) + except Exception as exc: + termerror( + f"Committing artifact failed. Artifact {artifact_id} won't be finalized." + ) + termerror(str(exc)) + self._fail_artifact_futures(artifact_id, exc) + else: + self._resolve_artifact_futures(artifact_id) + + def _fail_artifact_futures(self, artifact_id: str, exc: BaseException) -> None: + futures = self._artifacts[artifact_id]["result_futures"] + for result_future in futures: + result_future.set_exception(exc) + futures.clear() + + def _resolve_artifact_futures(self, artifact_id: str) -> None: + futures = self._artifacts[artifact_id]["result_futures"] + for result_future in futures: + result_future.set_result(None) + futures.clear() + + def start(self) -> None: + self._thread.start() + + def is_alive(self) -> bool: + return self._thread.is_alive() diff --git a/lib/python3.12/site-packages/wandb/filesync/upload_job.py b/lib/python3.12/site-packages/wandb/filesync/upload_job.py new file mode 100644 index 0000000000000000000000000000000000000000..3db449571b3f3e06310e76cb267d0dc042557915 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/filesync/upload_job.py @@ -0,0 +1,142 @@ +import logging +import os +from typing import TYPE_CHECKING, Optional + +import wandb +from wandb.sdk.lib.paths import LogicalPath + +if TYPE_CHECKING: + from wandb.filesync import dir_watcher, stats, step_upload + from wandb.sdk.internal import file_stream, internal_api + + +logger = logging.getLogger(__name__) + + +class UploadJob: + def __init__( + self, + stats: "stats.Stats", + api: "internal_api.Api", + file_stream: "file_stream.FileStreamApi", + silent: bool, + save_name: LogicalPath, + path: "dir_watcher.PathStr", + artifact_id: Optional[str], + md5: Optional[str], + copied: bool, + save_fn: Optional["step_upload.SaveFn"], + digest: Optional[str], + ) -> None: + """A file uploader. + + Args: + push_function: function(save_name, actual_path) which actually uploads + the file. + save_name: string logical location of the file relative to the run + directory. + path: actual string path of the file to upload on the filesystem. + """ + self._stats = stats + self._api = api + self._file_stream = file_stream + self.silent = silent + self.save_name = save_name + self.save_path = path + self.artifact_id = artifact_id + self.md5 = md5 + self.copied = copied + self.save_fn = save_fn + self.digest = digest + super().__init__() + + def run(self) -> None: + success = False + try: + self.push() + success = True + finally: + if self.copied and os.path.isfile(self.save_path): + os.remove(self.save_path) + if success: + self._file_stream.push_success(self.artifact_id, self.save_name) # type: ignore + + def push(self) -> None: + if self.save_fn: + # Retry logic must happen in save_fn currently + try: + deduped = self.save_fn( + lambda _, t: self._stats.update_uploaded_file(self.save_path, t) + ) + except Exception as e: + self._stats.update_failed_file(self.save_path) + logger.exception("Failed to upload file: %s", self.save_path) + wandb._sentry.exception(e) + message = str(e) + # TODO: this is usually XML, but could be JSON + if hasattr(e, "response"): + message = e.response.content + wandb.termerror( + f'Error uploading "{self.save_path}": {type(e).__name__}, {message}' + ) + raise + + if deduped: + logger.info("Skipped uploading %s", self.save_path) + self._stats.set_file_deduped(self.save_path) + else: + logger.info("Uploaded file %s", self.save_path) + return + + if self.md5: + # This is the new artifact manifest upload flow, in which we create the + # database entry for the manifest file before creating it. This is used for + # artifact L0 files. Which now is only artifact_manifest.json + _, response = self._api.create_artifact_manifest( + self.save_name, self.md5, self.artifact_id + ) + upload_url = response["uploadUrl"] + upload_headers = response["uploadHeaders"] + else: + # The classic file upload flow. We get a signed url and upload the file + # then the backend handles the cloud storage metadata callback to create the + # file entry. This flow has aged like a fine wine. + project = self._api.get_project() + _, upload_headers, result = self._api.upload_urls(project, [self.save_name]) + file_info = result[self.save_name] + upload_url = file_info["uploadUrl"] + + if upload_url is None: + logger.info("Skipped uploading %s", self.save_path) + self._stats.set_file_deduped(self.save_name) + else: + extra_headers = self._api._extra_http_headers + for upload_header in upload_headers: + key, val = upload_header.split(":", 1) + extra_headers[key] = val + # Copied from push TODO(artifacts): clean up + # If the upload URL is relative, fill it in with the base URL, + # since its a proxied file store like the on-prem VM. + if upload_url.startswith("/"): + upload_url = f"{self._api.api_url}{upload_url}" + try: + with open(self.save_path, "rb") as f: + self._api.upload_file_retry( + upload_url, + f, + lambda _, t: self.progress(t), + extra_headers=extra_headers, + ) + logger.info("Uploaded file %s", self.save_path) + except Exception as e: + self._stats.update_failed_file(self.save_name) + logger.exception("Failed to upload file: %s", self.save_path) + wandb._sentry.exception(e) + if not self.silent: + wandb.termerror( + f'Error uploading "{self.save_name}": {type(e).__name__}, {e}' + ) + raise + + def progress(self, total_bytes: int) -> None: + self._stats.update_uploaded_file(self.save_name, total_bytes) diff --git a/lib/python3.12/site-packages/wandb/integration/__init__.py b/lib/python3.12/site-packages/wandb/integration/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/wandb/integration/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4a34941e0beb5fa77ab2625cfceac1ff6523e8b2 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/catboost/__init__.py b/lib/python3.12/site-packages/wandb/integration/catboost/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ccd732b7e26cead03e12f0d6ab00f6d3165c2ba4 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/catboost/__init__.py @@ -0,0 +1,5 @@ +"""W&B callback for CatBoost.""" + +from .catboost import WandbCallback, log_summary + +__all__ = ["log_summary", "WandbCallback"] diff --git a/lib/python3.12/site-packages/wandb/integration/catboost/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/catboost/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..29c961d34c733238cad0bf9018d8dc1fea2bcf4b Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/catboost/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/catboost/__pycache__/catboost.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/catboost/__pycache__/catboost.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c8e5a369e422400b8db60439ffde8c20baef2892 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/catboost/__pycache__/catboost.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/catboost/catboost.py b/lib/python3.12/site-packages/wandb/integration/catboost/catboost.py new file mode 100644 index 0000000000000000000000000000000000000000..09dc31b84ffaec7826d70536eb59ef56452a9062 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/catboost/catboost.py @@ -0,0 +1,182 @@ +"""catboost init.""" + +from pathlib import Path +from types import SimpleNamespace +from typing import List, Union + +from catboost import CatBoostClassifier, CatBoostRegressor # type: ignore + +import wandb +from wandb.sdk.lib import telemetry as wb_telemetry + + +class WandbCallback: + """`WandbCallback` automatically integrates CatBoost with wandb. + + Args: + - metric_period: (int) if you are passing `metric_period` to your CatBoost model please pass the same value here (default=1). + + Passing `WandbCallback` to CatBoost will: + - log training and validation metrics at every `metric_period` + - log iteration at every `metric_period` + + Example: + ``` + train_pool = Pool( + train[features], label=train["label"], cat_features=cat_features + ) + test_pool = Pool(test[features], label=test["label"], cat_features=cat_features) + + model = CatBoostRegressor( + iterations=100, + loss_function="Cox", + eval_metric="Cox", + ) + + model.fit( + train_pool, + eval_set=test_pool, + callbacks=[WandbCallback()], + ) + ``` + """ + + def __init__(self, metric_period: int = 1): + if wandb.run is None: + raise wandb.Error("You must call `wandb.init()` before `WandbCallback()`") + + with wb_telemetry.context() as tel: + tel.feature.catboost_wandb_callback = True + + self.metric_period: int = metric_period + + def after_iteration(self, info: SimpleNamespace) -> bool: + if info.iteration % self.metric_period == 0: + for data, metric in info.metrics.items(): + for metric_name, log in metric.items(): + # todo: replace with wandb.run._log once available + wandb.log({f"{data}-{metric_name}": log[-1]}, commit=False) + # todo: replace with wandb.run._log once available + wandb.log({f"iteration@metric-period-{self.metric_period}": info.iteration}) + + return True + + +def _checkpoint_artifact( + model: Union[CatBoostClassifier, CatBoostRegressor], aliases: List[str] +) -> None: + """Upload model checkpoint as W&B artifact.""" + if wandb.run is None: + raise wandb.Error( + "You must call `wandb.init()` before `_checkpoint_artifact()`" + ) + + model_name = f"model_{wandb.run.id}" + # save the model in the default `cbm` format + model_path = Path(wandb.run.dir) / "model" + + model.save_model(model_path) + + model_artifact = wandb.Artifact(name=model_name, type="model") + model_artifact.add_file(str(model_path)) + wandb.log_artifact(model_artifact, aliases=aliases) + + +def _log_feature_importance( + model: Union[CatBoostClassifier, CatBoostRegressor], +) -> None: + """Log feature importance with default settings.""" + if wandb.run is None: + raise wandb.Error( + "You must call `wandb.init()` before `_checkpoint_artifact()`" + ) + + feat_df = model.get_feature_importance(prettified=True) + + fi_data = [ + [feat, feat_imp] + for feat, feat_imp in zip(feat_df["Feature Id"], feat_df["Importances"]) + ] + table = wandb.Table(data=fi_data, columns=["Feature", "Importance"]) + # todo: replace with wandb.run._log once available + wandb.log( + { + "Feature Importance": wandb.plot.bar( + table, "Feature", "Importance", title="Feature Importance" + ) + }, + commit=False, + ) + + +def log_summary( + model: Union[CatBoostClassifier, CatBoostRegressor], + log_all_params: bool = True, + save_model_checkpoint: bool = False, + log_feature_importance: bool = True, +) -> None: + """`log_summary` logs useful metrics about catboost model after training is done. + + Args: + model: it can be CatBoostClassifier or CatBoostRegressor. + log_all_params: (boolean) if True (default) log the model hyperparameters as W&B config. + save_model_checkpoint: (boolean) if True saves the model upload as W&B artifacts. + log_feature_importance: (boolean) if True (default) logs feature importance as W&B bar chart using the default setting of `get_feature_importance`. + + Using this along with `wandb_callback` will: + + - save the hyperparameters as W&B config, + - log `best_iteration` and `best_score` as `wandb.summary`, + - save and upload your trained model to Weights & Biases Artifacts (when `save_model_checkpoint = True`) + - log feature importance plot. + + Example: + ```python + train_pool = Pool( + train[features], label=train["label"], cat_features=cat_features + ) + test_pool = Pool(test[features], label=test["label"], cat_features=cat_features) + + model = CatBoostRegressor( + iterations=100, + loss_function="Cox", + eval_metric="Cox", + ) + + model.fit( + train_pool, + eval_set=test_pool, + callbacks=[WandbCallback()], + ) + + log_summary(model) + ``` + """ + if wandb.run is None: + raise wandb.Error("You must call `wandb.init()` before `log_summary()`") + + if not (isinstance(model, (CatBoostClassifier, CatBoostRegressor))): + raise wandb.Error( + "Model should be an instance of CatBoostClassifier or CatBoostRegressor" + ) + + with wb_telemetry.context() as tel: + tel.feature.catboost_log_summary = True + + # log configs + params = model.get_all_params() + if log_all_params: + wandb.config.update(params) + + # log best score and iteration + wandb.run.summary["best_iteration"] = model.get_best_iteration() + wandb.run.summary["best_score"] = model.get_best_score() + + # log model + if save_model_checkpoint: + aliases = ["best"] if params["use_best_model"] else ["last"] + _checkpoint_artifact(model, aliases=aliases) + + # Feature importance + if log_feature_importance: + _log_feature_importance(model) diff --git a/lib/python3.12/site-packages/wandb/integration/cohere/__init__.py b/lib/python3.12/site-packages/wandb/integration/cohere/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2d367dc6988bda4251745dc6b3610a3e92b4c85e --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/cohere/__init__.py @@ -0,0 +1,3 @@ +__all__ = ("autolog",) + +from .cohere import autolog diff --git a/lib/python3.12/site-packages/wandb/integration/cohere/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/cohere/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d0ffdbc6e0d357e6d5c467d2abaa657353e115d2 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/cohere/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/cohere/__pycache__/cohere.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/cohere/__pycache__/cohere.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..80398c70bcde9bd3eba5233a26e1cd943297e7f7 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/cohere/__pycache__/cohere.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/cohere/__pycache__/resolver.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/cohere/__pycache__/resolver.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a8aa44b77d9f4222ac42f97c7a427069b6223989 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/cohere/__pycache__/resolver.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/cohere/cohere.py b/lib/python3.12/site-packages/wandb/integration/cohere/cohere.py new file mode 100644 index 0000000000000000000000000000000000000000..91f9a43e23150a6882dbb87512e6fbe657a7b8d4 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/cohere/cohere.py @@ -0,0 +1,21 @@ +import logging + +from wandb.sdk.integration_utils.auto_logging import AutologAPI + +from .resolver import CohereRequestResponseResolver + +logger = logging.getLogger(__name__) + + +autolog = AutologAPI( + name="Cohere", + symbols=( + "Client.generate", + "Client.chat", + "Client.classify", + "Client.summarize", + "Client.rerank", + ), + resolver=CohereRequestResponseResolver(), + telemetry_feature="cohere_autolog", +) diff --git a/lib/python3.12/site-packages/wandb/integration/cohere/resolver.py b/lib/python3.12/site-packages/wandb/integration/cohere/resolver.py new file mode 100644 index 0000000000000000000000000000000000000000..6cfcdf020b87c8a5a3c15f164226ac50f8670d4d --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/cohere/resolver.py @@ -0,0 +1,347 @@ +import logging +from datetime import datetime +from typing import Any, Dict, List, Optional, Sequence, Tuple + +import wandb +from wandb.sdk.integration_utils.auto_logging import Response +from wandb.sdk.lib.runid import generate_id + +logger = logging.getLogger(__name__) + + +def subset_dict( + original_dict: Dict[str, Any], keys_subset: Sequence[str] +) -> Dict[str, Any]: + """Create a subset of a dictionary using a subset of keys. + + :param original_dict: The original dictionary. + :param keys_subset: The subset of keys to extract. + :return: A dictionary containing only the specified keys. + """ + return {key: original_dict[key] for key in keys_subset if key in original_dict} + + +def reorder_and_convert_dict_list_to_table( + data: List[Dict[str, Any]], order: List[str] +) -> Tuple[List[str], List[List[Any]]]: + """Convert a list of dictionaries to a pair of column names and corresponding values, with the option to order specific dictionaries. + + :param data: A list of dictionaries. + :param order: A list of keys specifying the desired order for specific dictionaries. The remaining dictionaries will be ordered based on their original order. + :return: A pair of column names and corresponding values. + """ + final_columns = [] + keys_present = set() + + # First, add all ordered keys to the final columns + for key in order: + if key not in keys_present: + final_columns.append(key) + keys_present.add(key) + + # Then, add any keys present in the dictionaries but not in the order + for d in data: + for key in d: + if key not in keys_present: + final_columns.append(key) + keys_present.add(key) + + # Then, construct the table of values + values = [] + for d in data: + row = [] + for key in final_columns: + row.append(d.get(key, None)) + values.append(row) + + return final_columns, values + + +def flatten_dict( + dictionary: Dict[str, Any], parent_key: str = "", sep: str = "-" +) -> Dict[str, Any]: + """Flatten a nested dictionary, joining keys using a specified separator. + + :param dictionary: The dictionary to flatten. + :param parent_key: The base key to prepend to each key. + :param sep: The separator to use when joining keys. + :return: A flattened dictionary. + """ + flattened_dict = {} + for key, value in dictionary.items(): + new_key = f"{parent_key}{sep}{key}" if parent_key else key + if isinstance(value, dict): + flattened_dict.update(flatten_dict(value, new_key, sep=sep)) + else: + flattened_dict[new_key] = value + return flattened_dict + + +def collect_common_keys(list_of_dicts: List[Dict[str, Any]]) -> Dict[str, List[Any]]: + """Collect the common keys of a list of dictionaries. For each common key, put its values into a list in the order they appear in the original dictionaries. + + :param list_of_dicts: The list of dictionaries to inspect. + :return: A dictionary with each common key and its corresponding list of values. + """ + common_keys = set.intersection(*map(set, list_of_dicts)) + common_dict = {key: [] for key in common_keys} + for d in list_of_dicts: + for key in common_keys: + common_dict[key].append(d[key]) + return common_dict + + +class CohereRequestResponseResolver: + """Class to resolve the request/response from the Cohere API and convert it to a dictionary that can be logged.""" + + def __call__( + self, + args: Sequence[Any], + kwargs: Dict[str, Any], + response: Response, + start_time: float, + time_elapsed: float, + ) -> Optional[Dict[str, Any]]: + """Process the response from the Cohere API and convert it to a dictionary that can be logged. + + :param args: The arguments of the original function. + :param kwargs: The keyword arguments of the original function. + :param response: The response from the Cohere API. + :param start_time: The start time of the request. + :param time_elapsed: The time elapsed for the request. + :return: A dictionary containing the parsed response and timing information. + """ + try: + # Each of the different endpoints map to one specific response type + # We want to 'type check' the response without directly importing the packages type + # It may make more sense to pass the invoked symbol from the AutologAPI instead + response_type = str(type(response)).split("'")[1].split(".")[-1] + + # Initialize parsed_response to None to handle the case where the response type is unsupported + parsed_response = None + if response_type == "Generations": + parsed_response = self._resolve_generate_response(response) + # TODO: Remove hard-coded default model name + table_column_order = [ + "start_time", + "query_id", + "model", + "prompt", + "text", + "token_likelihoods", + "likelihood", + "time_elapsed_(seconds)", + "end_time", + ] + default_model = "command" + elif response_type == "Chat": + parsed_response = self._resolve_chat_response(response) + table_column_order = [ + "start_time", + "query_id", + "model", + "conversation_id", + "response_id", + "query", + "text", + "prompt", + "preamble", + "chat_history", + "chatlog", + "time_elapsed_(seconds)", + "end_time", + ] + default_model = "command" + elif response_type == "Classifications": + parsed_response = self._resolve_classify_response(response) + kwargs = self._resolve_classify_kwargs(kwargs) + table_column_order = [ + "start_time", + "query_id", + "model", + "id", + "input", + "prediction", + "confidence", + "time_elapsed_(seconds)", + "end_time", + ] + default_model = "embed-english-v2.0" + elif response_type == "SummarizeResponse": + parsed_response = self._resolve_summarize_response(response) + table_column_order = [ + "start_time", + "query_id", + "model", + "response_id", + "text", + "additional_command", + "summary", + "time_elapsed_(seconds)", + "end_time", + "length", + "format", + ] + default_model = "summarize-xlarge" + elif response_type == "Reranking": + parsed_response = self._resolve_rerank_response(response) + table_column_order = [ + "start_time", + "query_id", + "model", + "id", + "query", + "top_n", + # This is a nested dict key that got flattened + "document-text", + "relevance_score", + "index", + "time_elapsed_(seconds)", + "end_time", + ] + default_model = "rerank-english-v2.0" + else: + logger.info(f"Unsupported Cohere response object: {response}") + + return self._resolve( + args, + kwargs, + parsed_response, + start_time, + time_elapsed, + response_type, + table_column_order, + default_model, + ) + except Exception as e: + logger.warning(f"Failed to resolve request/response: {e}") + return None + + # These helper functions process the response from different endpoints of the Cohere API. + # Since the response objects for different endpoints have different structures, + # we need different logic to process them. + + def _resolve_generate_response(self, response: Response) -> List[Dict[str, Any]]: + return_list = [] + for _response in response: + # Built in Cohere.*.Generations function to color token_likelihoods and return a dict of response data + _response_dict = _response._visualize_helper() + try: + _response_dict["token_likelihoods"] = wandb.Html( + _response_dict["token_likelihoods"] + ) + except (KeyError, ValueError): + pass + return_list.append(_response_dict) + + return return_list + + def _resolve_chat_response(self, response: Response) -> List[Dict[str, Any]]: + return [ + subset_dict( + response.__dict__, + [ + "response_id", + "generation_id", + "query", + "text", + "conversation_id", + "prompt", + "chatlog", + "preamble", + ], + ) + ] + + def _resolve_classify_response(self, response: Response) -> List[Dict[str, Any]]: + # The labels key is a dict returning the scores for the classification probability for each label provided + # We flatten this nested dict for ease of consumption in the wandb UI + return [flatten_dict(_response.__dict__) for _response in response] + + def _resolve_classify_kwargs(self, kwargs: Dict[str, Any]) -> Dict[str, Any]: + # Example texts look strange when rendered in Wandb UI as it is a list of text and label + # We extract each value into its own column + example_texts = [] + example_labels = [] + for example in kwargs["examples"]: + example_texts.append(example.text) + example_labels.append(example.label) + kwargs.pop("examples") + kwargs["example_texts"] = example_texts + kwargs["example_labels"] = example_labels + return kwargs + + def _resolve_summarize_response(self, response: Response) -> List[Dict[str, Any]]: + return [{"response_id": response.id, "summary": response.summary}] + + def _resolve_rerank_response(self, response: Response) -> List[Dict[str, Any]]: + # The documents key contains a dict containing the content of the document which is at least "text" + # We flatten this nested dict for ease of consumption in the wandb UI + flattened_response_dicts = [ + flatten_dict(_response.__dict__) for _response in response + ] + # ReRank returns each document provided a top_n value so we aggregate into one view so users can paginate a row + # As opposed to each row being one of the top_n responses + return_dict = collect_common_keys(flattened_response_dicts) + return_dict["id"] = response.id + return [return_dict] + + def _resolve( + self, + args: Sequence[Any], + kwargs: Dict[str, Any], + parsed_response: List[Dict[str, Any]], + start_time: float, + time_elapsed: float, + response_type: str, + table_column_order: List[str], + default_model: str, + ) -> Dict[str, Any]: + """Convert a list of dictionaries to a pair of column names and corresponding values, with the option to order specific dictionaries. + + :param args: The arguments passed to the API client. + :param kwargs: The keyword arguments passed to the API client. + :param parsed_response: The parsed response from the API. + :param start_time: The start time of the API request. + :param time_elapsed: The time elapsed during the API request. + :param response_type: The type of the API response. + :param table_column_order: The desired order of columns in the resulting table. + :param default_model: The default model to use if not specified in the response. + :return: A dictionary containing the formatted response. + """ + # Args[0] is the client object where we can grab specific metadata about the underlying API status + query_id = generate_id(length=16) + parsed_args = subset_dict( + args[0].__dict__, + ["api_version", "batch_size", "max_retries", "num_workers", "timeout"], + ) + + start_time_dt = datetime.fromtimestamp(start_time) + end_time_dt = datetime.fromtimestamp(start_time + time_elapsed) + + timings = { + "start_time": start_time_dt, + "end_time": end_time_dt, + "time_elapsed_(seconds)": time_elapsed, + } + + packed_data = [] + for _parsed_response in parsed_response: + _packed_dict = { + "query_id": query_id, + **kwargs, + **_parsed_response, + **timings, + **parsed_args, + } + if "model" not in _packed_dict: + _packed_dict["model"] = default_model + packed_data.append(_packed_dict) + + columns, data = reorder_and_convert_dict_list_to_table( + packed_data, table_column_order + ) + + request_response_table = wandb.Table(data=data, columns=columns) + + return {f"{response_type}": request_response_table} diff --git a/lib/python3.12/site-packages/wandb/integration/diffusers/__init__.py b/lib/python3.12/site-packages/wandb/integration/diffusers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5bcf7980133c41393fdc22db3bbff89be29be94e --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/diffusers/__init__.py @@ -0,0 +1,3 @@ +from .autologger import autolog + +__all__ = ["autolog"] diff --git a/lib/python3.12/site-packages/wandb/integration/diffusers/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/diffusers/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e78d4fdd10201a436fe17de5b94a04007e766bc1 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/diffusers/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/diffusers/__pycache__/autologger.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/diffusers/__pycache__/autologger.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8e2b1f2590ce0d535e8a976361199f4f79551318 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/diffusers/__pycache__/autologger.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/diffusers/__pycache__/pipeline_resolver.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/diffusers/__pycache__/pipeline_resolver.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e9da573376ca206fe804d97f6409873d91e30bd3 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/diffusers/__pycache__/pipeline_resolver.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/diffusers/autologger.py b/lib/python3.12/site-packages/wandb/integration/diffusers/autologger.py new file mode 100644 index 0000000000000000000000000000000000000000..ad21a77edc6a1125f8d2f5621c046c6bbe1f0b9b --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/diffusers/autologger.py @@ -0,0 +1,76 @@ +import logging + +from wandb.sdk.integration_utils.auto_logging import AutologAPI + +from .pipeline_resolver import DiffusersPipelineResolver + +logger = logging.getLogger(__name__) + +autolog = AutologAPI( + name="diffusers", + symbols=( + "DiffusionPipeline.__call__", + "AutoPipelineForText2Image.__call__", + "AutoPipelineForImage2Image.__call__", + "AutoPipelineForInpainting.__call__", + "StableDiffusionPipeline.__call__", + "KandinskyCombinedPipeline.__call__", + "KandinskyV22CombinedPipeline.__call__", + "LatentConsistencyModelPipeline.__call__", + "LDMTextToImagePipeline.__call__", + "StableDiffusionPanoramaPipeline.__call__", + "StableDiffusionParadigmsPipeline.__call__", + "PixArtAlphaPipeline.__call__", + "StableDiffusionSAGPipeline.__call__", + "SemanticStableDiffusionPipeline.__call__", + "WuerstchenCombinedPipeline.__call__", + "AltDiffusionPipeline.__call__", + "StableDiffusionAttendAndExcitePipeline.__call__", + "StableDiffusionXLPipeline.__call__", + "StableDiffusionXLImg2ImgPipeline.__call__", + "IFPipeline.__call__", + "BlipDiffusionPipeline.__call__", + "BlipDiffusionControlNetPipeline.__call__", + "StableDiffusionControlNetPipeline.__call__", + "StableDiffusionControlNetImg2ImgPipeline.__call__", + "StableDiffusionControlNetInpaintPipeline.__call__", + "CycleDiffusionPipeline.__call__", + "StableDiffusionInstructPix2PixPipeline.__call__", + "PaintByExamplePipeline.__call__", + "RePaintPipeline.__call__", + "KandinskyImg2ImgCombinedPipeline.__call__", + "KandinskyInpaintCombinedPipeline.__call__", + "KandinskyV22Img2ImgCombinedPipeline.__call__", + "KandinskyV22InpaintCombinedPipeline.__call__", + "Kandinsky3Pipeline.__call__", + "Kandinsky3Img2ImgPipeline.__call__", + "AnimateDiffPipeline.__call__", + "AudioLDMPipeline.__call__", + "AudioLDM2Pipeline.__call__", + "MusicLDMPipeline.__call__", + "StableDiffusionPix2PixZeroPipeline.__call__", + "PNDMPipeline.__call__", + "ShapEPipeline.__call__", + "StableDiffusionImg2ImgPipeline.__call__", + "StableDiffusionInpaintPipeline.__call__", + "StableDiffusionDepth2ImgPipeline.__call__", + "StableDiffusionImageVariationPipeline.__call__", + "StableDiffusionPipelineSafe.__call__", + "StableDiffusionUpscalePipeline.__call__", + "StableDiffusionAdapterPipeline.__call__", + "StableDiffusionGLIGENPipeline.__call__", + "StableDiffusionModelEditingPipeline.__call__", + "VersatileDiffusionTextToImagePipeline.__call__", + "VersatileDiffusionImageVariationPipeline.__call__", + "VersatileDiffusionDualGuidedPipeline.__call__", + "LDMPipeline.__call__", + "TextToVideoSDPipeline.__call__", + "TextToVideoZeroPipeline.__call__", + "StableVideoDiffusionPipeline.__call__", + "AmusedPipeline.__call__", + "StableDiffusionXLControlNetPipeline.__call__", + "StableDiffusionXLControlNetImg2ImgPipeline.__call__", + ), + resolver=DiffusersPipelineResolver(), + telemetry_feature="diffusers_autolog", +) diff --git a/lib/python3.12/site-packages/wandb/integration/diffusers/pipeline_resolver.py b/lib/python3.12/site-packages/wandb/integration/diffusers/pipeline_resolver.py new file mode 100644 index 0000000000000000000000000000000000000000..a5c4d73511c20f53e75f4b1cbf5a6cfc444c8875 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/diffusers/pipeline_resolver.py @@ -0,0 +1,50 @@ +from typing import Any, Dict, Sequence + +from wandb.sdk.integration_utils.auto_logging import Response + +from .resolvers import ( + SUPPORTED_MULTIMODAL_PIPELINES, + DiffusersMultiModalPipelineResolver, +) + + +class DiffusersPipelineResolver: + """Resolver for `DiffusionPipeline` request and responses from [HuggingFace Diffusers](https://huggingface.co/docs/diffusers/index), providing necessary data transformations, formatting, and logging. + + This is based off `wandb.sdk.integration_utils.auto_logging.RequestResponseResolver`. + """ + + def __init__(self) -> None: + self.wandb_table = None + self.pipeline_call_count = 1 + + def __call__( + self, + args: Sequence[Any], + kwargs: Dict[str, Any], + response: Response, + start_time: float, + time_elapsed: float, + ) -> Any: + """Main call method for the `DiffusersPipelineResolver` class. + + Args: + args: (Sequence[Any]) List of arguments. + kwargs: (Dict[str, Any]) Dictionary of keyword arguments. + response: (wandb.sdk.integration_utils.auto_logging.Response) The response from + the request. + start_time: (float) Time when request started. + time_elapsed: (float) Time elapsed for the request. + + Returns: + Packed data as a dictionary for logging to wandb, None if an exception occurred. + """ + pipeline_name = args[0].__class__.__name__ + resolver = None + if pipeline_name in SUPPORTED_MULTIMODAL_PIPELINES: + resolver = DiffusersMultiModalPipelineResolver( + pipeline_name, self.pipeline_call_count + ) + self.pipeline_call_count += 1 + loggable_dict = resolver(args, kwargs, response, start_time, time_elapsed) + return loggable_dict diff --git a/lib/python3.12/site-packages/wandb/integration/diffusers/resolvers/__init__.py b/lib/python3.12/site-packages/wandb/integration/diffusers/resolvers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d6211bd4bc2b9901593e477c1b5713bc074f8804 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/diffusers/resolvers/__init__.py @@ -0,0 +1,9 @@ +from .multimodal import ( + SUPPORTED_MULTIMODAL_PIPELINES, + DiffusersMultiModalPipelineResolver, +) + +__all__ = [ + "SUPPORTED_MULTIMODAL_PIPELINES", + "DiffusersMultiModalPipelineResolver", +] diff --git a/lib/python3.12/site-packages/wandb/integration/diffusers/resolvers/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/diffusers/resolvers/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..178b1b95ac49c156d3efc0cf31b6e42f12f5c740 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/diffusers/resolvers/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/diffusers/resolvers/__pycache__/multimodal.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/diffusers/resolvers/__pycache__/multimodal.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..98efaea7473d2fd078c4f2ce3398a8e267b00eda Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/diffusers/resolvers/__pycache__/multimodal.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/diffusers/resolvers/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/diffusers/resolvers/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5ab1be5945bedadeca13f01897a890792d9743e3 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/diffusers/resolvers/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/diffusers/resolvers/multimodal.py b/lib/python3.12/site-packages/wandb/integration/diffusers/resolvers/multimodal.py new file mode 100644 index 0000000000000000000000000000000000000000..34ef33e8639193b8f87a65434e6261ff112fbd7a --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/diffusers/resolvers/multimodal.py @@ -0,0 +1,881 @@ +import logging +from typing import Any, Dict, List, Sequence + +import wandb +from wandb.sdk.integration_utils.auto_logging import Response + +from .utils import ( + chunkify, + decode_sdxl_t2i_latents, + get_updated_kwargs, + postprocess_np_arrays_for_video, + postprocess_pils_to_np, +) + +logger = logging.getLogger(__name__) + + +SUPPORTED_MULTIMODAL_PIPELINES = { + "BlipDiffusionPipeline": { + "table-schema": [ + "Reference-Image", + "Prompt", + "Negative-Prompt", + "Source-Subject-Category", + "Target-Subject-Category", + "Generated-Image", + ], + "kwarg-logging": [ + "reference_image", + "prompt", + "neg_prompt", + "source_subject_category", + "target_subject_category", + ], + "kwarg-actions": [wandb.Image, None, None, None, None], + }, + "BlipDiffusionControlNetPipeline": { + "table-schema": [ + "Reference-Image", + "Control-Image", + "Prompt", + "Negative-Prompt", + "Source-Subject-Category", + "Target-Subject-Category", + "Generated-Image", + ], + "kwarg-logging": [ + "reference_image", + "condtioning_image", + "prompt", + "neg_prompt", + "source_subject_category", + "target_subject_category", + ], + "kwarg-actions": [wandb.Image, wandb.Image, None, None, None, None], + }, + "StableDiffusionControlNetPipeline": { + "table-schema": [ + "Control-Image", + "Prompt", + "Negative-Prompt", + "Generated-Image", + ], + "kwarg-logging": ["image", "prompt", "negative_prompt"], + "kwarg-actions": [wandb.Image, None, None], + }, + "StableDiffusionControlNetImg2ImgPipeline": { + "table-schema": [ + "Source-Image", + "Control-Image", + "Prompt", + "Negative-Prompt", + "Generated-Image", + ], + "kwarg-logging": ["image", "control_image", "prompt", "negative_prompt"], + "kwarg-actions": [wandb.Image, wandb.Image, None, None], + }, + "StableDiffusionControlNetInpaintPipeline": { + "table-schema": [ + "Source-Image", + "Mask-Image", + "Control-Image", + "Prompt", + "Negative-Prompt", + "Generated-Image", + ], + "kwarg-logging": [ + "image", + "mask_image", + "control_image", + "prompt", + "negative_prompt", + ], + "kwarg-actions": [wandb.Image, wandb.Image, wandb.Image, None, None], + }, + "CycleDiffusionPipeline": { + "table-schema": [ + "Source-Image", + "Prompt", + "Source-Prompt", + "Generated-Image", + ], + "kwarg-logging": [ + "image", + "prompt", + "source_prompt", + ], + "kwarg-actions": [wandb.Image, None, None], + }, + "StableDiffusionInstructPix2PixPipeline": { + "table-schema": [ + "Source-Image", + "Prompt", + "Negative-Prompt", + "Generated-Image", + ], + "kwarg-logging": [ + "image", + "prompt", + "negative_prompt", + ], + "kwarg-actions": [wandb.Image, None, None], + }, + "PaintByExamplePipeline": { + "table-schema": [ + "Source-Image", + "Example-Image", + "Mask-Prompt", + "Generated-Image", + ], + "kwarg-logging": [ + "image", + "example_image", + "mask_image", + ], + "kwarg-actions": [wandb.Image, wandb.Image, wandb.Image], + }, + "RePaintPipeline": { + "table-schema": [ + "Source-Image", + "Mask-Prompt", + "Generated-Image", + ], + "kwarg-logging": [ + "image", + "mask_image", + ], + "kwarg-actions": [wandb.Image, wandb.Image], + }, + "StableDiffusionPipeline": { + "table-schema": ["Prompt", "Negative-Prompt", "Generated-Image"], + "kwarg-logging": ["prompt", "negative_prompt"], + "kwarg-actions": [None, None], + }, + "KandinskyCombinedPipeline": { + "table-schema": ["Prompt", "Negative-Prompt", "Generated-Image"], + "kwarg-logging": ["prompt", "negative_prompt"], + "kwarg-actions": [None, None], + }, + "KandinskyV22CombinedPipeline": { + "table-schema": ["Prompt", "Negative-Prompt", "Generated-Image"], + "kwarg-logging": ["prompt", "negative_prompt"], + "kwarg-actions": [None, None], + }, + "LatentConsistencyModelPipeline": { + "table-schema": ["Prompt", "Generated-Image"], + "kwarg-logging": ["prompt"], + "kwarg-actions": [None], + }, + "LDMTextToImagePipeline": { + "table-schema": ["Prompt", "Generated-Image"], + "kwarg-logging": ["prompt"], + "kwarg-actions": [None], + }, + "StableDiffusionPanoramaPipeline": { + "table-schema": ["Prompt", "Negative-Prompt", "Generated-Image"], + "kwarg-logging": ["prompt", "negative_prompt"], + "kwarg-actions": [None, None], + }, + "PixArtAlphaPipeline": { + "table-schema": ["Prompt", "Negative-Prompt", "Generated-Image"], + "kwarg-logging": ["prompt", "negative_prompt"], + "kwarg-actions": [None, None], + }, + "StableDiffusionSAGPipeline": { + "table-schema": ["Prompt", "Negative-Prompt", "Generated-Image"], + "kwarg-logging": ["prompt", "negative_prompt"], + "kwarg-actions": [None, None], + }, + "SemanticStableDiffusionPipeline": { + "table-schema": ["Prompt", "Negative-Prompt", "Generated-Image"], + "kwarg-logging": ["prompt", "negative_prompt"], + "kwarg-actions": [None, None], + }, + "WuerstchenCombinedPipeline": { + "table-schema": ["Prompt", "Negative-Prompt", "Generated-Image"], + "kwarg-logging": ["prompt", "negative_prompt"], + "kwarg-actions": [None, None], + }, + "IFPipeline": { + "table-schema": ["Prompt", "Negative-Prompt", "Generated-Image"], + "kwarg-logging": ["prompt", "negative_prompt"], + "kwarg-actions": [None, None], + }, + "AltDiffusionPipeline": { + "table-schema": ["Prompt", "Negative-Prompt", "Generated-Image"], + "kwarg-logging": ["prompt", "negative_prompt"], + "kwarg-actions": [None, None], + }, + "StableDiffusionAttendAndExcitePipeline": { + "table-schema": ["Prompt", "Negative-Prompt", "Generated-Image"], + "kwarg-logging": ["prompt", "negative_prompt"], + "kwarg-actions": [None, None], + }, + "KandinskyImg2ImgCombinedPipeline": { + "table-schema": [ + "Source-Image", + "Prompt", + "Negative-Prompt", + "Generated-Image", + ], + "kwarg-logging": ["image", "prompt", "negative_prompt"], + "kwarg-actions": [wandb.Image, None, None], + }, + "KandinskyInpaintCombinedPipeline": { + "table-schema": [ + "Source-Image", + "Prompt", + "Negative-Prompt", + "Generated-Image", + ], + "kwarg-logging": ["image", "prompt", "negative_prompt"], + "kwarg-actions": [wandb.Image, None, None], + }, + "KandinskyV22Img2ImgCombinedPipeline": { + "table-schema": [ + "Source-Image", + "Prompt", + "Negative-Prompt", + "Generated-Image", + ], + "kwarg-logging": ["image", "prompt", "negative_prompt"], + "kwarg-actions": [wandb.Image, None, None], + }, + "KandinskyV22InpaintCombinedPipeline": { + "table-schema": [ + "Source-Image", + "Prompt", + "Negative-Prompt", + "Generated-Image", + ], + "kwarg-logging": ["image", "prompt", "negative_prompt"], + "kwarg-actions": [wandb.Image, None, None], + }, + "AnimateDiffPipeline": { + "table-schema": [ + "Prompt", + "Negative-Prompt", + "Number-of-Frames", + "Generated-Video", + ], + "kwarg-logging": ["prompt", "negative_prompt", "num_frames"], + "kwarg-actions": [None, None, None], + "output-type": "video", + }, + "StableVideoDiffusionPipeline": { + "table-schema": [ + "Input-Image", + "Frames-Per-Second", + "Generated-Video", + ], + "kwarg-logging": ["image", "fps"], + "kwarg-actions": [wandb.Image, None], + "output-type": "video", + }, + "AudioLDMPipeline": { + "table-schema": [ + "Prompt", + "Negative-Prompt", + "Audio-Length-in-Seconds", + "Generated-Audio", + ], + "kwarg-logging": ["prompt", "negative_prompt", "audio_length_in_s"], + "kwarg-actions": [None, None, None], + "output-type": "audio", + }, + "AudioLDM2Pipeline": { + "table-schema": [ + "Prompt", + "Negative-Prompt", + "Audio-Length-in-Seconds", + "Generated-Audio", + ], + "kwarg-logging": ["prompt", "negative_prompt", "audio_length_in_s"], + "kwarg-actions": [None, None, None], + "output-type": "audio", + }, + "MusicLDMPipeline": { + "table-schema": [ + "Prompt", + "Negative-Prompt", + "Audio-Length-in-Seconds", + "Generated-Audio", + ], + "kwarg-logging": ["prompt", "negative_prompt", "audio_length_in_s"], + "kwarg-actions": [None, None, None], + "output-type": "audio", + }, + "StableDiffusionPix2PixZeroPipeline": { + "table-schema": [ + "Prompt", + "Negative-Prompt", + "Generated-Image", + ], + "kwarg-logging": ["prompt", "negative_prompt"], + "kwarg-actions": [None, None], + }, + "PNDMPipeline": { + "table-schema": [ + "Batch-Size", + "Number-of-Inference-Steps", + "Generated-Image", + ], + "kwarg-logging": ["batch_size", "num_inference_steps"], + "kwarg-actions": [None, None], + }, + "ShapEPipeline": { + "table-schema": [ + "Prompt", + "Generated-Video", + ], + "kwarg-logging": ["prompt"], + "kwarg-actions": [None], + "output-type": "video", + }, + "StableDiffusionImg2ImgPipeline": { + "table-schema": [ + "Source-Image", + "Prompt", + "Negative-Prompt", + "Generated-Image", + ], + "kwarg-logging": ["image", "prompt", "negative_prompt"], + "kwarg-actions": [wandb.Image, None, None], + }, + "StableDiffusionInpaintPipeline": { + "table-schema": [ + "Source-Image", + "Mask-Image", + "Prompt", + "Negative-Prompt", + "Generated-Image", + ], + "kwarg-logging": ["image", "mask_image", "prompt", "negative_prompt"], + "kwarg-actions": [wandb.Image, wandb.Image, None, None], + }, + "StableDiffusionDepth2ImgPipeline": { + "table-schema": [ + "Source-Image", + "Prompt", + "Negative-Prompt", + "Generated-Image", + ], + "kwarg-logging": ["image", "prompt", "negative_prompt"], + "kwarg-actions": [wandb.Image, None, None], + }, + "StableDiffusionImageVariationPipeline": { + "table-schema": [ + "Source-Image", + "Generated-Image", + ], + "kwarg-logging": [ + "image", + ], + "kwarg-actions": [wandb.Image], + }, + "StableDiffusionPipelineSafe": { + "table-schema": [ + "Prompt", + "Negative-Prompt", + "Generated-Image", + ], + "kwarg-logging": ["prompt", "negative_prompt"], + "kwarg-actions": [None, None], + }, + "StableDiffusionUpscalePipeline": { + "table-schema": [ + "Source-Image", + "Prompt", + "Negative-Prompt", + "Upscaled-Image", + ], + "kwarg-logging": ["image", "prompt", "negative_prompt"], + "kwarg-actions": [wandb.Image, None, None], + }, + "StableDiffusionAdapterPipeline": { + "table-schema": [ + "Source-Image", + "Prompt", + "Negative-Prompt", + "Generated-Image", + ], + "kwarg-logging": ["image", "prompt", "negative_prompt"], + "kwarg-actions": [wandb.Image, None, None], + }, + "StableDiffusionGLIGENPipeline": { + "table-schema": [ + "Prompt", + "GLIGEN-Phrases", + "GLIGEN-Boxes", + "GLIGEN-Inpaint-Image", + "Negative-Prompt", + "Generated-Image", + ], + "kwarg-logging": [ + "prompt", + "gligen_phrases", + "gligen_boxes", + "gligen_inpaint_image", + "negative_prompt", + ], + "kwarg-actions": [None, None, None, wandb.Image, None], + }, + "VersatileDiffusionTextToImagePipeline": { + "table-schema": [ + "Prompt", + "Negative-Prompt", + "Generated-Image", + ], + "kwarg-logging": ["prompt", "negative_prompt"], + "kwarg-actions": [None, None], + }, + "VersatileDiffusionImageVariationPipeline": { + "table-schema": [ + "Source-Image", + "Negative-Prompt", + "Generated-Image", + ], + "kwarg-logging": ["image", "negative_prompt"], + "kwarg-actions": [wandb.Image, None], + }, + "VersatileDiffusionDualGuidedPipeline": { + "table-schema": [ + "Source-Image", + "Prompt", + "Negative-Prompt", + "Generated-Image", + ], + "kwarg-logging": ["image", "prompt", "negative_prompt"], + "kwarg-actions": [wandb.Image, None, None], + }, + "LDMPipeline": { + "table-schema": [ + "Batch-Size", + "Number-of-Inference-Steps", + "Generated-Image", + ], + "kwarg-logging": ["batch_size", "num_inference_steps"], + "kwarg-actions": [None, None], + }, + "TextToVideoSDPipeline": { + "table-schema": [ + "Prompt", + "Negative-Prompt", + "Number-of-Frames", + "Generated-Video", + ], + "kwarg-logging": ["prompt", "negative_prompt", "num_frames"], + "output-type": "video", + }, + "TextToVideoZeroPipeline": { + "table-schema": [ + "Prompt", + "Negative-Prompt", + "Number-of-Frames", + "Generated-Video", + ], + "kwarg-logging": ["prompt", "negative_prompt", "video_length"], + }, + "AmusedPipeline": { + "table-schema": [ + "Prompt", + "Guidance Scale", + "Generated-Image", + ], + "kwarg-logging": [ + "prompt", + "guidance_scale", + ], + "kwarg-actions": [None, None], + }, + "StableDiffusionXLControlNetPipeline": { + "table-schema": [ + "Prompt-1", + "Prompt-2", + "Control-Image", + "Negative-Prompt-1", + "Negative-Prompt-2", + "Generated-Image", + ], + "kwarg-logging": [ + "prompt", + "prompt_2", + "image", + "negative_prompt", + "negative_prompt_2", + ], + "kwarg-actions": [None, None, wandb.Image, None, None], + }, + "StableDiffusionXLControlNetImg2ImgPipeline": { + "table-schema": [ + "Prompt-1", + "Prompt-2", + "Input-Image", + "Control-Image", + "Negative-Prompt-1", + "Negative-Prompt-2", + "Generated-Image", + ], + "kwarg-logging": [ + "prompt", + "prompt_2", + "image", + "control_image", + "negative_prompt", + "negative_prompt_2", + ], + "kwarg-actions": [None, None, wandb.Image, wandb.Image, None, None], + }, + "Kandinsky3Pipeline": { + "table-schema": [ + "Prompt", + "Negative-Prompt", + "Generated-Image", + ], + "kwarg-logging": [ + "prompt", + "negative_prompt", + ], + "kwarg-actions": [None, None], + }, + "Kandinsky3Img2ImgPipeline": { + "table-schema": [ + "Prompt", + "Negative-Prompt", + "Input-Image", + "Generated-Image", + ], + "kwarg-logging": [ + "prompt", + "negative_prompt", + "image", + ], + "kwarg-actions": [None, None, wandb.Image], + }, + "StableDiffusionXLPipeline": { + "table-schema": [ + "Prompt", + "Negative-Prompt", + "Prompt-2", + "Negative-Prompt-2", + "Generated-Image", + ], + "kwarg-logging": [ + "prompt", + "negative_prompt", + "prompt_2", + "negative_prompt_2", + ], + "kwarg-actions": [None, None, None, None], + }, + "StableDiffusionXLImg2ImgPipeline": { + "table-schema": [ + "Prompt", + "Negative-Prompt", + "Prompt-2", + "Negative-Prompt-2", + "Input-Image", + "Generated-Image", + ], + "kwarg-logging": [ + "prompt", + "negative_prompt", + "prompt_2", + "negative_prompt_2", + "image", + ], + "kwarg-actions": [None, None, None, None, wandb.Image], + }, +} + + +class DiffusersMultiModalPipelineResolver: + """Resolver for request and responses from [HuggingFace Diffusers](https://huggingface.co/docs/diffusers/index) multi-modal Diffusion Pipelines, providing necessary data transformations, formatting, and logging. + + This resolver is internally involved in the + `__call__` for `wandb.integration.diffusers.pipeline_resolver.DiffusersPipelineResolver`. + This is based on `wandb.sdk.integration_utils.auto_logging.RequestResponseResolver`. + + Args: + pipeline_name: (str) The name of the Diffusion Pipeline. + """ + + def __init__(self, pipeline_name: str, pipeline_call_count: int) -> None: + self.pipeline_name = pipeline_name + self.pipeline_call_count = pipeline_call_count + columns = [] + if pipeline_name in SUPPORTED_MULTIMODAL_PIPELINES: + columns += SUPPORTED_MULTIMODAL_PIPELINES[pipeline_name]["table-schema"] + else: + wandb.Error("Pipeline not supported for logging") + self.wandb_table = wandb.Table(columns=columns) + + def __call__( + self, + args: Sequence[Any], + kwargs: Dict[str, Any], + response: Response, + start_time: float, + time_elapsed: float, + ) -> Any: + """Main call method for the `DiffusersPipelineResolver` class. + + Args: + args: (Sequence[Any]) List of arguments. + kwargs: (Dict[str, Any]) Dictionary of keyword arguments. + response: (wandb.sdk.integration_utils.auto_logging.Response) The response from + the request. + start_time: (float) Time when request started. + time_elapsed: (float) Time elapsed for the request. + + Returns: + Packed data as a dictionary for logging to wandb, None if an exception occurred. + """ + try: + # Get the pipeline and the args + pipeline, args = args[0], args[1:] + + # Update the Kwargs so that they can be logged easily + kwargs = get_updated_kwargs(pipeline, args, kwargs) + + # Get the pipeline configs + pipeline_configs = dict(pipeline.config) + pipeline_configs["pipeline-name"] = self.pipeline_name + + if "workflow" not in wandb.config: + wandb.config.update( + { + "workflow": [ + { + "pipeline": pipeline_configs, + "params": kwargs, + "stage": f"Pipeline-Call-{self.pipeline_call_count}", + } + ] + } + ) + else: + existing_workflow = wandb.config.workflow + updated_workflow = existing_workflow + [ + { + "pipeline": pipeline_configs, + "params": kwargs, + "stage": f"Pipeline-Call-{self.pipeline_call_count}", + } + ] + wandb.config.update( + {"workflow": updated_workflow}, allow_val_change=True + ) + + # Return the WandB loggable dict + return self.prepare_loggable_dict(pipeline, response, kwargs) + except Exception as e: + logger.warning(e) + return None + + def get_output_images(self, response: Response) -> List: + """Unpack the generated images, audio, video, etc. from the Diffusion Pipeline's response. + + Args: + response: (wandb.sdk.integration_utils.auto_logging.Response) The response from + the request. + + Returns: + List of generated images, audio, video, etc. + """ + if "output-type" not in SUPPORTED_MULTIMODAL_PIPELINES[self.pipeline_name]: + return response.images + else: + if ( + SUPPORTED_MULTIMODAL_PIPELINES[self.pipeline_name]["output-type"] + == "video" + ): + if self.pipeline_name in ["ShapEPipeline"]: + return response.images + return response.frames + elif ( + SUPPORTED_MULTIMODAL_PIPELINES[self.pipeline_name]["output-type"] + == "audio" + ): + return response.audios + + def log_media(self, image: Any, loggable_kwarg_chunks: List, idx: int) -> None: + """Log the generated images, audio, video, etc. from the Diffusion Pipeline's response along with an optional caption to a media panel in the run. + + Args: + image: (Any) The generated images, audio, video, etc. from the Diffusion + Pipeline's response. + loggable_kwarg_chunks: (List) Loggable chunks of kwargs. + """ + if "output-type" not in SUPPORTED_MULTIMODAL_PIPELINES[self.pipeline_name]: + try: + caption = "" + if self.pipeline_name in [ + "StableDiffusionXLPipeline", + "StableDiffusionXLImg2ImgPipeline", + ]: + prompt_index = SUPPORTED_MULTIMODAL_PIPELINES[self.pipeline_name][ + "kwarg-logging" + ].index("prompt") + prompt2_index = SUPPORTED_MULTIMODAL_PIPELINES[self.pipeline_name][ + "kwarg-logging" + ].index("prompt_2") + caption = f"Prompt-1: {loggable_kwarg_chunks[prompt_index][idx]}\nPrompt-2: {loggable_kwarg_chunks[prompt2_index][idx]}" + else: + prompt_index = SUPPORTED_MULTIMODAL_PIPELINES[self.pipeline_name][ + "kwarg-logging" + ].index("prompt") + caption = loggable_kwarg_chunks[prompt_index][idx] + except ValueError: + caption = None + wandb.log( + { + f"Generated-Image/Pipeline-Call-{self.pipeline_call_count}": wandb.Image( + image, caption=caption + ) + } + ) + else: + if ( + SUPPORTED_MULTIMODAL_PIPELINES[self.pipeline_name]["output-type"] + == "video" + ): + try: + prompt_index = SUPPORTED_MULTIMODAL_PIPELINES[self.pipeline_name][ + "kwarg-logging" + ].index("prompt") + caption = loggable_kwarg_chunks[prompt_index][idx] + except ValueError: + caption = None + wandb.log( + { + f"Generated-Video/Pipeline-Call-{self.pipeline_call_count}": wandb.Video( + postprocess_pils_to_np(image), fps=4, caption=caption + ) + } + ) + elif ( + SUPPORTED_MULTIMODAL_PIPELINES[self.pipeline_name]["output-type"] + == "audio" + ): + try: + prompt_index = SUPPORTED_MULTIMODAL_PIPELINES[self.pipeline_name][ + "kwarg-logging" + ].index("prompt") + caption = loggable_kwarg_chunks[prompt_index][idx] + except ValueError: + caption = None + wandb.log( + { + f"Generated-Audio/Pipeline-Call-{self.pipeline_call_count}": wandb.Audio( + image, sample_rate=16000, caption=caption + ) + } + ) + + def add_data_to_table( + self, image: Any, loggable_kwarg_chunks: List, idx: int + ) -> None: + """Populate the row of the `wandb.Table`. + + Args: + image: (Any) The generated images, audio, video, etc. from the Diffusion + Pipeline's response. + loggable_kwarg_chunks: (List) Loggable chunks of kwargs. + idx: (int) Chunk index. + """ + table_row = [] + kwarg_actions = SUPPORTED_MULTIMODAL_PIPELINES[self.pipeline_name][ + "kwarg-actions" + ] + for column_idx, loggable_kwarg_chunk in enumerate(loggable_kwarg_chunks): + if kwarg_actions[column_idx] is None: + table_row.append( + loggable_kwarg_chunk[idx] + if loggable_kwarg_chunk[idx] is not None + else "" + ) + else: + table_row.append(kwarg_actions[column_idx](loggable_kwarg_chunk[idx])) + if "output-type" not in SUPPORTED_MULTIMODAL_PIPELINES[self.pipeline_name]: + table_row.append(wandb.Image(image)) + else: + if ( + SUPPORTED_MULTIMODAL_PIPELINES[self.pipeline_name]["output-type"] + == "video" + ): + table_row.append(wandb.Video(postprocess_pils_to_np(image), fps=4)) + elif ( + SUPPORTED_MULTIMODAL_PIPELINES[self.pipeline_name]["output-type"] + == "audio" + ): + table_row.append(wandb.Audio(image, sample_rate=16000)) + self.wandb_table.add_data(*table_row) + + def prepare_loggable_dict( + self, pipeline: Any, response: Response, kwargs: Dict[str, Any] + ) -> Dict[str, Any]: + """Prepare the loggable dictionary, which is the packed data as a dictionary for logging to wandb, None if an exception occurred. + + Args: + pipeline: (Any) The Diffusion Pipeline. + response: (wandb.sdk.integration_utils.auto_logging.Response) The response from + the request. + kwargs: (Dict[str, Any]) Dictionary of keyword arguments. + + Returns: + Packed data as a dictionary for logging to wandb, None if an exception occurred. + """ + # Unpack the generated images, audio, video, etc. from the Diffusion Pipeline's response. + images = self.get_output_images(response) + if ( + self.pipeline_name == "StableDiffusionXLPipeline" + and kwargs["output_type"] == "latent" + ): + images = decode_sdxl_t2i_latents(pipeline, response.images) + + # Account for exception pipelines for text-to-video + if self.pipeline_name in ["TextToVideoSDPipeline", "TextToVideoZeroPipeline"]: + video = postprocess_np_arrays_for_video( + images, normalize=self.pipeline_name == "TextToVideoZeroPipeline" + ) + wandb.log( + { + f"Generated-Video/Pipeline-Call-{self.pipeline_call_count}": wandb.Video( + video, fps=4, caption=kwargs["prompt"] + ) + } + ) + loggable_kwarg_ids = SUPPORTED_MULTIMODAL_PIPELINES[self.pipeline_name][ + "kwarg-logging" + ] + table_row = [ + kwargs[loggable_kwarg_ids[idx]] + for idx in range(len(loggable_kwarg_ids)) + ] + table_row.append(wandb.Video(video, fps=4)) + self.wandb_table.add_data(*table_row) + else: + loggable_kwarg_ids = SUPPORTED_MULTIMODAL_PIPELINES[self.pipeline_name][ + "kwarg-logging" + ] + # chunkify loggable kwargs + loggable_kwarg_chunks = [] + for loggable_kwarg_id in loggable_kwarg_ids: + loggable_kwarg_chunks.append( + kwargs[loggable_kwarg_id] + if isinstance(kwargs[loggable_kwarg_id], list) + else [kwargs[loggable_kwarg_id]] + ) + # chunkify the generated media + images = chunkify(images, len(loggable_kwarg_chunks[0])) + for idx in range(len(loggable_kwarg_chunks[0])): + for image in images[idx]: + # Log media to media panel + self.log_media(image, loggable_kwarg_chunks, idx) + # Populate the row of the wandb_table + self.add_data_to_table(image, loggable_kwarg_chunks, idx) + return { + f"Result-Table/Pipeline-Call-{self.pipeline_call_count}": self.wandb_table + } diff --git a/lib/python3.12/site-packages/wandb/integration/diffusers/resolvers/utils.py b/lib/python3.12/site-packages/wandb/integration/diffusers/resolvers/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..82a6aac045ba9739af9d0ab342a63e33252b63e0 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/diffusers/resolvers/utils.py @@ -0,0 +1,102 @@ +import inspect +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence + +import wandb +from wandb.util import get_module + +if TYPE_CHECKING: + np_array = get_module("numpy.array") + torch_float_tensor = get_module("torch.FloatTensor") + + +def chunkify(input_list, chunk_size) -> List: + chunk_size = max(1, chunk_size) + return [ + input_list[i : i + chunk_size] for i in range(0, len(input_list), chunk_size) + ] + + +def get_updated_kwargs( + pipeline: Any, args: Sequence[Any], kwargs: Dict[str, Any] +) -> Dict[str, Any]: + pipeline_call_parameters = list( + inspect.signature(pipeline.__call__).parameters.items() + ) + for idx, arg in enumerate(args): + kwargs[pipeline_call_parameters[idx][0]] = arg + for pipeline_parameter in pipeline_call_parameters: + if pipeline_parameter[0] not in kwargs: + kwargs[pipeline_parameter[0]] = pipeline_parameter[1].default + if "generator" in kwargs: + generator = kwargs["generator"] + kwargs["generator"] = ( + { + "seed": generator.initial_seed(), + "device": generator.device, + "random_state": generator.get_state().cpu().numpy().tolist(), + } + if generator is not None + else None + ) + if "ip_adapter_image" in kwargs: + if kwargs["ip_adapter_image"] is not None: + wandb.log({"IP-Adapter-Image": wandb.Image(kwargs["ip_adapter_image"])}) + return kwargs + + +def postprocess_pils_to_np(image: List) -> "np_array": + np = get_module( + "numpy", + required="Please ensure NumPy is installed. You can run `pip install numpy` to install it.", + ) + return np.stack( + [np.transpose(np.array(img).astype("uint8"), axes=(2, 0, 1)) for img in image], + axis=0, + ) + + +def postprocess_np_arrays_for_video( + images: List["np_array"], normalize: Optional[bool] = False +) -> "np_array": + np = get_module( + "numpy", + required="Please ensure NumPy is installed. You can run `pip install numpy` to install it.", + ) + images = [(img * 255).astype("uint8") for img in images] if normalize else images + return np.transpose(np.stack((images), axis=0), axes=(0, 3, 1, 2)) + + +def decode_sdxl_t2i_latents(pipeline: Any, latents: "torch_float_tensor") -> List: + """Decode latents generated by [`diffusers.StableDiffusionXLPipeline`](https://huggingface.co/docs/diffusers/main/en/api/pipelines/stable_diffusion/stable_diffusion_xl#stable-diffusion-xl). + + Args: + pipeline: (diffusers.DiffusionPipeline) The Diffusion Pipeline from + [`diffusers`](https://huggingface.co/docs/diffusers). + latents (torch.FloatTensor): The generated latents. + + Returns: + List of `PIL` images corresponding to the generated latents. + """ + torch = get_module( + "torch", + required="Please ensure PyTorch is installed. You can check out https://pytorch.org/get-started/locally/#start-locally for installation instructions.", + ) + with torch.no_grad(): + needs_upcasting = ( + pipeline.vae.dtype == torch.float16 and pipeline.vae.config.force_upcast + ) + if needs_upcasting: + pipeline.upcast_vae() + latents = latents.to( + next(iter(pipeline.vae.post_quant_conv.parameters())).dtype + ) + images = pipeline.vae.decode( + latents / pipeline.vae.config.scaling_factor, return_dict=False + )[0] + if needs_upcasting: + pipeline.vae.to(dtype=torch.float16) + if pipeline.watermark is not None: + images = pipeline.watermark.apply_watermark(images) + images = pipeline.image_processor.postprocess(images, output_type="pil") + pipeline.maybe_free_model_hooks() + return images diff --git a/lib/python3.12/site-packages/wandb/integration/fastai/__init__.py b/lib/python3.12/site-packages/wandb/integration/fastai/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d12d71d58183b0ce679c2bbb0e4e917fed791e2c --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/fastai/__init__.py @@ -0,0 +1,243 @@ +"""Hooks that add fast.ai v1 Learners to Weights & Biases through a callback. + +Requested logged data can be configured through the callback constructor. + +Examples: + WandbCallback can be used when initializing the Learner:: + + ``` + from wandb.fastai import WandbCallback + [...] + learn = Learner(data, ..., callback_fns=WandbCallback) + learn.fit(epochs) + ``` + + Custom parameters can be given using functools.partial:: + + ``` + from wandb.fastai import WandbCallback + from functools import partial + [...] + learn = Learner(data, ..., callback_fns=partial(WandbCallback, ...)) + learn.fit(epochs) + ``` + + Finally, it is possible to use WandbCallback only when starting + training. In this case it must be instantiated:: + + ``` + learn.fit(..., callbacks=WandbCallback(learn)) + ``` + + or, with custom parameters:: + + ``` + learn.fit(..., callbacks=WandbCallback(learn, ...)) + ``` +""" + +import random +import sys +from pathlib import Path +from typing import Any, Literal, Optional + +import fastai +from fastai.callbacks import TrackerCallback + +import wandb +from wandb.sdk.lib import ipython + +try: + import matplotlib + + if not ipython.in_jupyter(): + matplotlib.use("Agg") # non-interactive backend (avoid tkinter issues) + import matplotlib.pyplot as plt +except ImportError: + wandb.termwarn("matplotlib required if logging sample image predictions") + + +class WandbCallback(TrackerCallback): + """Callback for saving model topology, losses & metrics. + + Optionally logs weights, gradients, sample predictions and best trained model. + + Args: + learn (fastai.basic_train.Learner): the fast.ai learner to hook. + log (str): "gradients", "parameters", "all", or None. Losses & metrics are always logged. + save_model (bool): save model at the end of each epoch. It will also load best model at the end of training. + monitor (str): metric to monitor for saving best model. None uses default TrackerCallback monitor value. + mode (str): "auto", "min" or "max" to compare "monitor" values and define best model. + input_type (str): "images" or None. Used to display sample predictions. + validation_data (list): data used for sample predictions if input_type is set. + predictions (int): number of predictions to make if input_type is set and validation_data is None. + seed (int): initialize random generator for sample predictions if input_type is set and validation_data is None. + """ + + # Record if watch has been called previously (even in another instance) + _watch_called = False + + def __init__( + self, + learn: "fastai.basic_train.Learner", + log: Optional[Literal["gradients", "parameters", "all"]] = "gradients", + save_model: bool = True, + monitor: Optional[str] = None, + mode: Literal["auto", "min", "max"] = "auto", + input_type: Optional[Literal["images"]] = None, + validation_data: Optional[list] = None, + predictions: int = 36, + seed: int = 12345, + ) -> None: + # Check if wandb.init has been called + if wandb.run is None: + raise ValueError("You must call wandb.init() before WandbCallback()") + + # Adapted from fast.ai "SaveModelCallback" + if monitor is None: + # use default TrackerCallback monitor value + super().__init__(learn, mode=mode) + else: + super().__init__(learn, monitor=monitor, mode=mode) + self.save_model = save_model + self.model_path = Path(wandb.run.dir) / "bestmodel.pth" + + self.log = log + self.input_type = input_type + self.best = None + + # Select items for sample predictions to see evolution along training + self.validation_data = validation_data + if input_type and not self.validation_data: + wandb_random = random.Random(seed) # For repeatability + predictions = min(predictions, len(learn.data.valid_ds)) + indices = wandb_random.sample(range(len(learn.data.valid_ds)), predictions) + self.validation_data = [learn.data.valid_ds[i] for i in indices] + + def on_train_begin(self, **kwargs: Any) -> None: + """Call watch method to log model topology, gradients & weights.""" + # Set self.best, method inherited from "TrackerCallback" by "SaveModelCallback" + super().on_train_begin() + + # Ensure we don't call "watch" multiple times + if not WandbCallback._watch_called: + WandbCallback._watch_called = True + + # Logs model topology and optionally gradients and weights + wandb.watch(self.learn.model, log=self.log) + + def on_epoch_end( + self, epoch: int, smooth_loss: float, last_metrics: list, **kwargs: Any + ) -> None: + """Log training loss, validation loss and custom metrics & log prediction samples & save model.""" + if self.save_model: + # Adapted from fast.ai "SaveModelCallback" + current = self.get_monitor_value() + if current is not None and self.operator(current, self.best): + wandb.termlog( + f"Better model found at epoch {epoch} with {self.monitor} value: {current}." + ) + self.best = current + + # Save within wandb folder + with self.model_path.open("wb") as model_file: + self.learn.save(model_file) + + # Log sample predictions if learn.predict is available + if self.validation_data: + try: + self._wandb_log_predictions() + except FastaiError as e: + wandb.termwarn(e.message) + self.validation_data = None # prevent from trying again on next loop + except Exception as e: + wandb.termwarn(f"Unable to log prediction samples.\n{e}") + self.validation_data = None # prevent from trying again on next loop + + # Log losses & metrics + # Adapted from fast.ai "CSVLogger" + logs = { + name: stat + for name, stat in list( + zip(self.learn.recorder.names, [epoch, smooth_loss] + last_metrics) + ) + } + wandb.log(logs) + + def on_train_end(self, **kwargs: Any) -> None: + """Load the best model.""" + if self.save_model: + # Adapted from fast.ai "SaveModelCallback" + if self.model_path.is_file(): + with self.model_path.open("rb") as model_file: + self.learn.load(model_file, purge=False) + wandb.termlog(f"Loaded best saved model from {self.model_path}") + + def _wandb_log_predictions(self) -> None: + """Log prediction samples.""" + pred_log = [] + + if self.validation_data is None: + return + + for x, y in self.validation_data: + try: + pred = self.learn.predict(x) + except Exception: + raise FastaiError( + 'Unable to run "predict" method from Learner to log prediction samples.' + ) + + # scalar -> likely to be a category + # tensor of dim 1 -> likely to be multicategory + if not pred[1].shape or pred[1].dim() == 1: + pred_log.append( + wandb.Image( + x.data, + caption=f"Ground Truth: {y}\nPrediction: {pred[0]}", + ) + ) + + # most vision datasets have a "show" function we can use + elif hasattr(x, "show"): + # log input data + pred_log.append(wandb.Image(x.data, caption="Input data", grouping=3)) + + # log label and prediction + for im, capt in ((pred[0], "Prediction"), (y, "Ground Truth")): + # Resize plot to image resolution + # from https://stackoverflow.com/a/13714915 + my_dpi = 100 + fig = plt.figure(frameon=False, dpi=my_dpi) + h, w = x.size + fig.set_size_inches(w / my_dpi, h / my_dpi) + ax = plt.Axes(fig, [0.0, 0.0, 1.0, 1.0]) + ax.set_axis_off() + fig.add_axes(ax) + + # Superpose label or prediction to input image + x.show(ax=ax, y=im) + pred_log.append(wandb.Image(fig, caption=capt)) + plt.close(fig) + + # likely to be an image + elif hasattr(y, "shape") and ( + (len(y.shape) == 2) or (len(y.shape) == 3 and y.shape[0] in [1, 3, 4]) + ): + pred_log.extend( + [ + wandb.Image(x.data, caption="Input data", grouping=3), + wandb.Image(pred[0].data, caption="Prediction"), + wandb.Image(y.data, caption="Ground Truth"), + ] + ) + + # we just log input data + else: + pred_log.append(wandb.Image(x.data, caption="Input data")) + + wandb.log({"Prediction Samples": pred_log}, commit=False) + + +class FastaiError(wandb.Error): + pass diff --git a/lib/python3.12/site-packages/wandb/integration/fastai/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/fastai/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c38afd53fcb2b1860fc4cc306686fccfd7a93c1b Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/fastai/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/gym/__init__.py b/lib/python3.12/site-packages/wandb/integration/gym/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b0624e4def25ab1e51044221f577066c7939108b --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/gym/__init__.py @@ -0,0 +1,98 @@ +import re +import sys +from typing import Literal, Optional + +import wandb +import wandb.util + +_gym_version_lt_0_26: Optional[bool] = None +_gymnasium_version_lt_1_0_0: Optional[bool] = None + +_required_error_msg = ( + "Couldn't import the gymnasium python package, install with `pip install gymnasium`" +) +GymLib = Literal["gym", "gymnasium"] + + +def monitor(): + """Monitor a gym environment. + + Supports both gym and gymnasium. + """ + gym_lib: Optional[GymLib] = None + + # gym is not maintained anymore, gymnasium is the drop-in replacement - prefer it + if wandb.util.get_module("gymnasium") is not None: + gym_lib = "gymnasium" + elif wandb.util.get_module("gym") is not None: + gym_lib = "gym" + + if gym_lib is None: + raise wandb.Error(_required_error_msg) + + global _gym_version_lt_0_26 + global _gymnasium_version_lt_1_0_0 + + if _gym_version_lt_0_26 is None or _gymnasium_version_lt_1_0_0 is None: + if gym_lib == "gym": + import gym + else: + import gymnasium as gym # type: ignore + + from packaging.version import parse + + gym_lib_version = parse(gym.__version__) + _gym_version_lt_0_26 = gym_lib_version < parse("0.26.0") + _gymnasium_version_lt_1_0_0 = gym_lib_version < parse("1.0.0a1") + + path = "path" # Default path + if gym_lib == "gymnasium" and not _gymnasium_version_lt_1_0_0: + vcr_recorder_attribute = "RecordVideo" + wrappers = wandb.util.get_module( + f"{gym_lib}.wrappers", + required=_required_error_msg, + ) + recorder = getattr(wrappers, vcr_recorder_attribute) + else: + vcr = wandb.util.get_module( + f"{gym_lib}.wrappers.monitoring.video_recorder", + required=_required_error_msg, + ) + # Breaking change in gym 0.26.0 + if _gym_version_lt_0_26: + vcr_recorder_attribute = "ImageEncoder" + recorder = getattr(vcr, vcr_recorder_attribute) + path = "output_path" # Override path for older gym versions + else: + vcr_recorder_attribute = "VideoRecorder" + recorder = getattr(vcr, vcr_recorder_attribute) + + recorder.orig_close = recorder.close + + def close(self): + recorder.orig_close(self) + if not self.enabled: + return + if wandb.run: + m = re.match(r".+(video\.\d+).+", getattr(self, path)) + key = m.group(1) if m else "videos" + wandb.log({key: wandb.Video(getattr(self, path))}) + + def del_(self): + self.orig_close() + + if not _gym_version_lt_0_26: + recorder.__del__ = del_ + recorder.close = close + + if gym_lib == "gymnasium" and not _gymnasium_version_lt_1_0_0: + wrapper_name = vcr_recorder_attribute + else: + wrapper_name = f"monitoring.video_recorder.{vcr_recorder_attribute}" + + wandb.patched["gym"].append( + [ + f"{gym_lib}.wrappers.{wrapper_name}", + "close", + ] + ) diff --git a/lib/python3.12/site-packages/wandb/integration/gym/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/gym/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fe64b411bba3b7776ec0019e833a464ba82da347 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/gym/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/huggingface/__init__.py b/lib/python3.12/site-packages/wandb/integration/huggingface/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..943249ebe22686fcbdfaa6708c33e15ec875bccc --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/huggingface/__init__.py @@ -0,0 +1,3 @@ +__all__ = ("autolog",) + +from .huggingface import autolog diff --git a/lib/python3.12/site-packages/wandb/integration/huggingface/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/huggingface/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ca2dfd0b028f1a188370aca57b7cee5fb37a0fb2 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/huggingface/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/huggingface/__pycache__/huggingface.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/huggingface/__pycache__/huggingface.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..60f5e574bc5662e2f50fea3c619252e1221c8f5d Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/huggingface/__pycache__/huggingface.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/huggingface/__pycache__/resolver.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/huggingface/__pycache__/resolver.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d73c743ad5b53219f0a2105f8a2b6c689acae940 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/huggingface/__pycache__/resolver.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/huggingface/huggingface.py b/lib/python3.12/site-packages/wandb/integration/huggingface/huggingface.py new file mode 100644 index 0000000000000000000000000000000000000000..d44cf2cbb51d009821e26924ab944a953a1ff866 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/huggingface/huggingface.py @@ -0,0 +1,18 @@ +import logging + +from wandb.sdk.integration_utils.auto_logging import AutologAPI + +from .resolver import HuggingFacePipelineRequestResponseResolver + +logger = logging.getLogger(__name__) + +resolver = HuggingFacePipelineRequestResponseResolver() + +autolog = AutologAPI( + name="transformers", + symbols=("Pipeline.__call__",), + resolver=resolver, + telemetry_feature="hf_pipeline_autolog", +) + +autolog.get_latest_id = resolver.get_latest_id diff --git a/lib/python3.12/site-packages/wandb/integration/huggingface/resolver.py b/lib/python3.12/site-packages/wandb/integration/huggingface/resolver.py new file mode 100644 index 0000000000000000000000000000000000000000..2acbdabe56e5e86c1d4442c99b7e9232d5621a48 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/huggingface/resolver.py @@ -0,0 +1,213 @@ +import logging +import os +from datetime import datetime +from typing import Any, Dict, List, Optional, Sequence, Tuple, Union + +import pytz + +import wandb +from wandb.sdk.integration_utils.auto_logging import Response +from wandb.sdk.lib.runid import generate_id + +logger = logging.getLogger(__name__) + +SUPPORTED_PIPELINE_TASKS = [ + "text-classification", + "sentiment-analysis", + "question-answering", + "summarization", + "translation", + "text2text-generation", + "text-generation", + # "conversational", +] + +PIPELINES_WITH_TOP_K = [ + "text-classification", + "sentiment-analysis", + "question-answering", +] + + +class HuggingFacePipelineRequestResponseResolver: + """Resolver for HuggingFace's pipeline request and responses, providing necessary data transformations and formatting. + + This is based off (from wandb.sdk.integration_utils.auto_logging import RequestResponseResolver) + """ + + autolog_id = None + + def __call__( + self, + args: Sequence[Any], + kwargs: Dict[str, Any], + response: Response, + start_time: float, + time_elapsed: float, + ) -> Optional[Dict[str, Any]]: + """Main call method for this class. + + :param args: list of arguments + :param kwargs: dictionary of keyword arguments + :param response: the response from the request + :param start_time: time when request started + :param time_elapsed: time elapsed for the request + :returns: packed data as a dictionary for logging to wandb, None if an exception occurred + """ + try: + pipe, input_data = args[:2] + task = pipe.task + + # Translation tasks are in the form of `translation_x_to_y` + if task in SUPPORTED_PIPELINE_TASKS or task.startswith("translation"): + model = self._get_model(pipe) + if model is None: + return None + model_alias = model.name_or_path + timestamp = datetime.now(pytz.utc) + + input_data, response = self._transform_task_specific_data( + task, input_data, response + ) + formatted_data = self._format_data(task, input_data, response, kwargs) + packed_data = self._create_table( + formatted_data, model_alias, timestamp, time_elapsed + ) + table_name = os.environ.get("WANDB_AUTOLOG_TABLE_NAME", f"{task}") + # TODO: Let users decide the name in a way that does not use an environment variable + + return { + table_name: wandb.Table( + columns=packed_data[0], data=packed_data[1:] + ) + } + + logger.warning( + f"The task: `{task}` is not yet supported.\nPlease contact `wandb` to notify us if you would like support for this task" + ) + except Exception as e: + logger.warning(e) + return None + + # TODO: This should have a dependency on PreTrainedModel. i.e. isinstance(PreTrainedModel) + # from transformers.modeling_utils import PreTrainedModel + # We do not want this dependency explicitly in our codebase so we make a very general + # assumption about the structure of the pipeline which may have unintended consequences + def _get_model(self, pipe) -> Optional[Any]: + """Extracts model from the pipeline. + + :param pipe: the HuggingFace pipeline + :returns: Model if available, None otherwise + """ + model = pipe.model + try: + return model.model + except AttributeError: + logger.info( + "Model does not have a `.model` attribute. Assuming `pipe.model` is the correct model." + ) + return model + + @staticmethod + def _transform_task_specific_data( + task: str, input_data: Union[List[Any], Any], response: Union[List[Any], Any] + ) -> Tuple[Union[List[Any], Any], Union[List[Any], Any]]: + """Transform input and response data based on specific tasks. + + :param task: the task name + :param input_data: the input data + :param response: the response data + :returns: tuple of transformed input_data and response + """ + if task == "question-answering": + input_data = input_data if isinstance(input_data, list) else [input_data] + input_data = [data.__dict__ for data in input_data] + elif task == "conversational": + # We only grab the latest input/output pair from the conversation + # Logging the whole conversation renders strangely. + input_data = input_data if isinstance(input_data, list) else [input_data] + input_data = [data.__dict__["past_user_inputs"][-1] for data in input_data] + + response = response if isinstance(response, list) else [response] + response = [data.__dict__["generated_responses"][-1] for data in response] + return input_data, response + + def _format_data( + self, + task: str, + input_data: Union[List[Any], Any], + response: Union[List[Any], Any], + kwargs: Dict[str, Any], + ) -> List[Dict[str, Any]]: + """Formats input data, response, and kwargs into a list of dictionaries. + + :param task: the task name + :param input_data: the input data + :param response: the response data + :param kwargs: dictionary of keyword arguments + :returns: list of dictionaries containing formatted data + """ + input_data = input_data if isinstance(input_data, list) else [input_data] + response = response if isinstance(response, list) else [response] + + formatted_data = [] + for i_text, r_text in zip(input_data, response): + # Unpack single element responses for better rendering in wandb UI when it is a task without top_k + # top_k = 1 would unpack the response into a single element while top_k > 1 would be a list + # this would cause the UI to not properly concatenate the tables of the same task by omitting the elements past the first + if ( + (isinstance(r_text, list)) + and (len(r_text) == 1) + and task not in PIPELINES_WITH_TOP_K + ): + r_text = r_text[0] + formatted_data.append( + {"input": i_text, "response": r_text, "kwargs": kwargs} + ) + return formatted_data + + def _create_table( + self, + formatted_data: List[Dict[str, Any]], + model_alias: str, + timestamp: float, + time_elapsed: float, + ) -> List[List[Any]]: + """Creates a table from formatted data, model alias, timestamp, and elapsed time. + + :param formatted_data: list of dictionaries containing formatted data + :param model_alias: alias of the model + :param timestamp: timestamp of the data + :param time_elapsed: time elapsed from the beginning + :returns: list of lists, representing a table of data. [0]th element = columns. [1]st element = data + """ + header = [ + "ID", + "Model Alias", + "Timestamp", + "Elapsed Time", + "Input", + "Response", + "Kwargs", + ] + table = [header] + autolog_id = generate_id(length=16) + + for data in formatted_data: + row = [ + autolog_id, + model_alias, + timestamp, + time_elapsed, + data["input"], + data["response"], + data["kwargs"], + ] + table.append(row) + + self.autolog_id = autolog_id + + return table + + def get_latest_id(self): + return self.autolog_id diff --git a/lib/python3.12/site-packages/wandb/integration/keras/__init__.py b/lib/python3.12/site-packages/wandb/integration/keras/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8699d6805b92dab3381456270a6449b2078d983b --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/keras/__init__.py @@ -0,0 +1,11 @@ +"""Tools for integrating `wandb` with [`Keras`](https://keras.io/).""" + +__all__ = ( + "WandbCallback", + "WandbMetricsLogger", + "WandbModelCheckpoint", + "WandbEvalCallback", +) + +from .callbacks import WandbEvalCallback, WandbMetricsLogger, WandbModelCheckpoint +from .keras import WandbCallback # TODO: legacy callback to be deprecated diff --git a/lib/python3.12/site-packages/wandb/integration/keras/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/keras/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cbcdf48bf030272e36cdf693f8879e71e172a38f Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/keras/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/keras/__pycache__/keras.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/keras/__pycache__/keras.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e01c65a8b042051dfea203a2409762c27c4f2bba Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/keras/__pycache__/keras.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/keras/callbacks/__init__.py b/lib/python3.12/site-packages/wandb/integration/keras/callbacks/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6035bd4c9c5067e553404a5a4bb818ef582a36a8 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/keras/callbacks/__init__.py @@ -0,0 +1,5 @@ +__all__ = ("WandbMetricsLogger", "WandbModelCheckpoint", "WandbEvalCallback") + +from .metrics_logger import WandbMetricsLogger +from .model_checkpoint import WandbModelCheckpoint +from .tables_builder import WandbEvalCallback diff --git a/lib/python3.12/site-packages/wandb/integration/keras/callbacks/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/keras/callbacks/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2d5110d88256558a365eb5711330b6f4ab2a978a Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/keras/callbacks/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/keras/callbacks/__pycache__/metrics_logger.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/keras/callbacks/__pycache__/metrics_logger.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5fa0a71396c39182fdb5c1f627275ac5320d434f Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/keras/callbacks/__pycache__/metrics_logger.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/keras/callbacks/__pycache__/model_checkpoint.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/keras/callbacks/__pycache__/model_checkpoint.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7e58134ba61356f2f807af1dc5a5989998e7511f Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/keras/callbacks/__pycache__/model_checkpoint.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/keras/callbacks/__pycache__/tables_builder.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/keras/callbacks/__pycache__/tables_builder.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d848b6b9e4ae8f16540dc465b2f568f374429a2c Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/keras/callbacks/__pycache__/tables_builder.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/keras/callbacks/metrics_logger.py b/lib/python3.12/site-packages/wandb/integration/keras/callbacks/metrics_logger.py new file mode 100644 index 0000000000000000000000000000000000000000..b51f5bfd7af7a134268529d617fd68daf5adcce9 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/keras/callbacks/metrics_logger.py @@ -0,0 +1,129 @@ +from typing import Any, Dict, Literal, Optional, Union + +import tensorflow as tf # type: ignore +from tensorflow.keras import callbacks + +import wandb +from wandb.integration.keras.keras import patch_tf_keras +from wandb.sdk.lib import telemetry + +LogStrategy = Literal["epoch", "batch"] + + +patch_tf_keras() + + +class WandbMetricsLogger(callbacks.Callback): + """Logger that sends system metrics to W&B. + + `WandbMetricsLogger` automatically logs the `logs` dictionary that callback methods + take as argument to wandb. + + This callback automatically logs the following to a W&B run page: + * system (CPU/GPU/TPU) metrics, + * train and validation metrics defined in `model.compile`, + * learning rate (both for a fixed value or a learning rate scheduler) + + Notes: + If you resume training by passing `initial_epoch` to `model.fit` and you are using a + learning rate scheduler, make sure to pass `initial_global_step` to + `WandbMetricsLogger`. The `initial_global_step` is `step_size * initial_step`, where + `step_size` is number of training steps per epoch. `step_size` can be calculated as + the product of the cardinality of the training dataset and the batch size. + + Args: + log_freq: ("epoch", "batch", or int) if "epoch", logs metrics + at the end of each epoch. If "batch", logs metrics at the end + of each batch. If an integer, logs metrics at the end of that + many batches. Defaults to "epoch". + initial_global_step: (int) Use this argument to correctly log the + learning rate when you resume training from some `initial_epoch`, + and a learning rate scheduler is used. This can be computed as + `step_size * initial_step`. Defaults to 0. + """ + + def __init__( + self, + log_freq: Union[LogStrategy, int] = "epoch", + initial_global_step: int = 0, + *args: Any, + **kwargs: Any, + ) -> None: + super().__init__(*args, **kwargs) + + if wandb.run is None: + raise wandb.Error( + "You must call `wandb.init()` before WandbMetricsLogger()" + ) + + with telemetry.context(run=wandb.run) as tel: + tel.feature.keras_metrics_logger = True + + if log_freq == "batch": + log_freq = 1 + + self.logging_batch_wise = isinstance(log_freq, int) + self.log_freq: Any = log_freq if self.logging_batch_wise else None + self.global_batch = 0 + self.global_step = initial_global_step + + if self.logging_batch_wise: + # define custom x-axis for batch logging. + wandb.define_metric("batch/batch_step") + # set all batch metrics to be logged against batch_step. + wandb.define_metric("batch/*", step_metric="batch/batch_step") + else: + # define custom x-axis for epoch-wise logging. + wandb.define_metric("epoch/epoch") + # set all epoch-wise metrics to be logged against epoch. + wandb.define_metric("epoch/*", step_metric="epoch/epoch") + + def _get_lr(self) -> Union[float, None]: + if isinstance( + self.model.optimizer.learning_rate, + (tf.Variable, tf.Tensor), + ) or ( + hasattr(self.model.optimizer.learning_rate, "shape") + and self.model.optimizer.learning_rate.shape == () + ): + return float(self.model.optimizer.learning_rate.numpy().item()) + try: + return float( + self.model.optimizer.learning_rate(step=self.global_step).numpy().item() + ) + except Exception as e: + wandb.termerror(f"Unable to log learning rate: {e}", repeat=False) + return None + + def on_epoch_end(self, epoch: int, logs: Optional[Dict[str, Any]] = None) -> None: + """Called at the end of an epoch.""" + logs = dict() if logs is None else {f"epoch/{k}": v for k, v in logs.items()} + + logs["epoch/epoch"] = epoch + + lr = self._get_lr() + if lr is not None: + logs["epoch/learning_rate"] = lr + + wandb.log(logs) + + def on_batch_end(self, batch: int, logs: Optional[Dict[str, Any]] = None) -> None: + self.global_step += 1 + """An alias for `on_train_batch_end` for backwards compatibility.""" + if self.logging_batch_wise and batch % self.log_freq == 0: + logs = {f"batch/{k}": v for k, v in logs.items()} if logs else {} + logs["batch/batch_step"] = self.global_batch + + lr = self._get_lr() + if lr is not None: + logs["batch/learning_rate"] = lr + + wandb.log(logs) + + self.global_batch += self.log_freq + + def on_train_batch_end( + self, batch: int, logs: Optional[Dict[str, Any]] = None + ) -> None: + """Called at the end of a training batch in `fit` methods.""" + self.on_batch_end(batch, logs if logs else {}) diff --git a/lib/python3.12/site-packages/wandb/integration/keras/callbacks/model_checkpoint.py b/lib/python3.12/site-packages/wandb/integration/keras/callbacks/model_checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..9990b4ffd19ad40d4f1204d5c6bc133c352c0835 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/keras/callbacks/model_checkpoint.py @@ -0,0 +1,188 @@ +import os +import string +from typing import Any, Dict, List, Literal, Optional, Union + +import tensorflow as tf # type: ignore +from tensorflow.keras import callbacks # type: ignore + +import wandb +from wandb.sdk.lib import telemetry +from wandb.sdk.lib.paths import StrPath + +from ..keras import patch_tf_keras + +Mode = Literal["auto", "min", "max"] +SaveStrategy = Literal["epoch"] + +patch_tf_keras() + + +class WandbModelCheckpoint(callbacks.ModelCheckpoint): + """A checkpoint that periodically saves a Keras model or model weights. + + Saved weights are uploaded to W&B as a `wandb.Artifact`. + + Since this callback is subclassed from `tf.keras.callbacks.ModelCheckpoint`, the + checkpointing logic is taken care of by the parent callback. You can learn more + here: https://www.tensorflow.org/api_docs/python/tf/keras/callbacks/ModelCheckpoint + + This callback is to be used in conjunction with training using `model.fit()` to save + a model or weights (in a checkpoint file) at some interval. The model checkpoints + will be logged as W&B Artifacts. You can learn more here: + https://docs.wandb.ai/guides/artifacts + + This callback provides the following features: + - Save the model that has achieved "best performance" based on "monitor". + - Save the model at the end of every epoch regardless of the performance. + - Save the model at the end of epoch or after a fixed number of training batches. + - Save only model weights, or save the whole model. + - Save the model either in SavedModel format or in `.h5` format. + + Args: + filepath: (Union[str, os.PathLike]) path to save the model file. `filepath` + can contain named formatting options, which will be filled by the value + of `epoch` and keys in `logs` (passed in `on_epoch_end`). For example: + if `filepath` is `model-{epoch:02d}-{val_loss:.2f}`, then the + model checkpoints will be saved with the epoch number and the + validation loss in the filename. + monitor: (str) The metric name to monitor. Default to "val_loss". + verbose: (int) Verbosity mode, 0 or 1. Mode 0 is silent, and mode 1 + displays messages when the callback takes an action. + save_best_only: (bool) if `save_best_only=True`, it only saves when the model + is considered the "best" and the latest best model according to the + quantity monitored will not be overwritten. If `filepath` doesn't contain + formatting options like `{epoch}` then `filepath` will be overwritten by + each new better model locally. The model logged as an artifact will still be + associated with the correct `monitor`. Artifacts will be uploaded + continuously and versioned separately as a new best model is found. + save_weights_only: (bool) if True, then only the model's weights will be saved. + mode: (Mode) one of {'auto', 'min', 'max'}. For `val_acc`, this should be `max`, + for `val_loss` this should be `min`, etc. + save_freq: (Union[SaveStrategy, int]) `epoch` or integer. When using `'epoch'`, + the callback saves the model after each epoch. When using an integer, the + callback saves the model at end of this many batches. + Note that when monitoring validation metrics such as `val_acc` or `val_loss`, + save_freq must be set to "epoch" as those metrics are only available at the + end of an epoch. + initial_value_threshold: (Optional[float]) Floating point initial "best" value of the metric + to be monitored. + """ + + def __init__( + self, + filepath: StrPath, + monitor: str = "val_loss", + verbose: int = 0, + save_best_only: bool = False, + save_weights_only: bool = False, + mode: Mode = "auto", + save_freq: Union[SaveStrategy, int] = "epoch", + initial_value_threshold: Optional[float] = None, + **kwargs: Any, + ) -> None: + super().__init__( + filepath=filepath, + monitor=monitor, + verbose=verbose, + save_best_only=save_best_only, + save_weights_only=save_weights_only, + mode=mode, + save_freq=save_freq, + initial_value_threshold=initial_value_threshold, + **kwargs, + ) + if wandb.run is None: + raise wandb.Error( + "You must call `wandb.init()` before `WandbModelCheckpoint()`" + ) + with telemetry.context(run=wandb.run) as tel: + tel.feature.keras_model_checkpoint = True + + self.save_weights_only = save_weights_only + + # User-friendly warning when trying to save the best model. + if self.save_best_only: + self._check_filepath() + + self._is_old_tf_keras_version: Optional[bool] = None + + def on_train_batch_end( + self, batch: int, logs: Optional[Dict[str, float]] = None + ) -> None: + if self._should_save_on_batch(batch): + if self.is_old_tf_keras_version: + # Save the model and get filepath + self._save_model(epoch=self._current_epoch, logs=logs) + filepath = self._get_file_path(epoch=self._current_epoch, logs=logs) + else: + # Save the model and get filepath + self._save_model(epoch=self._current_epoch, batch=batch, logs=logs) + filepath = self._get_file_path( + epoch=self._current_epoch, batch=batch, logs=logs + ) + # Log the model as artifact + aliases = ["latest", f"epoch_{self._current_epoch}_batch_{batch}"] + self._log_ckpt_as_artifact(filepath, aliases=aliases) + + def on_epoch_end(self, epoch: int, logs: Optional[Dict[str, float]] = None) -> None: + super().on_epoch_end(epoch, logs) + # Check if model checkpoint is created at the end of epoch. + if self.save_freq == "epoch": + # Get filepath where the model checkpoint is saved. + if self.is_old_tf_keras_version: + filepath = self._get_file_path(epoch=epoch, logs=logs) + else: + filepath = self._get_file_path(epoch=epoch, batch=None, logs=logs) + # Log the model as artifact + aliases = ["latest", f"epoch_{epoch}"] + self._log_ckpt_as_artifact(filepath, aliases=aliases) + + def _log_ckpt_as_artifact( + self, filepath: str, aliases: Optional[List[str]] = None + ) -> None: + """Log model checkpoint as W&B Artifact.""" + try: + assert wandb.run is not None + model_checkpoint_artifact = wandb.Artifact( + f"run_{wandb.run.id}_model", type="model" + ) + if os.path.isfile(filepath): + model_checkpoint_artifact.add_file(filepath) + elif os.path.isdir(filepath): + model_checkpoint_artifact.add_dir(filepath) + else: + raise FileNotFoundError(f"No such file or directory {filepath}") + wandb.log_artifact(model_checkpoint_artifact, aliases=aliases or []) + except ValueError: + # This error occurs when `save_best_only=True` and the model + # checkpoint is not saved for that epoch/batch. Since TF/Keras + # is giving friendly log, we can avoid clustering the stdout. + pass + + def _check_filepath(self) -> None: + placeholders = [] + for tup in string.Formatter().parse(self.filepath): + if tup[1] is not None: + placeholders.append(tup[1]) + if len(placeholders) == 0: + wandb.termwarn( + "When using `save_best_only`, ensure that the `filepath` argument " + "contains formatting placeholders like `{epoch:02d}` or `{batch:02d}`. " + "This ensures correct interpretation of the logged artifacts.", + repeat=False, + ) + + @property + def is_old_tf_keras_version(self) -> Optional[bool]: + if self._is_old_tf_keras_version is None: + from packaging.version import parse + + try: + if parse(tf.keras.__version__) < parse("2.6.0"): + self._is_old_tf_keras_version = True + else: + self._is_old_tf_keras_version = False + except AttributeError: + self._is_old_tf_keras_version = False + + return self._is_old_tf_keras_version diff --git a/lib/python3.12/site-packages/wandb/integration/keras/callbacks/tables_builder.py b/lib/python3.12/site-packages/wandb/integration/keras/callbacks/tables_builder.py new file mode 100644 index 0000000000000000000000000000000000000000..bd19bfb314f4c06c703858aed6fe47ca12bbcefc --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/keras/callbacks/tables_builder.py @@ -0,0 +1,228 @@ +import abc +from typing import Any, Dict, List, Optional + +from tensorflow.keras.callbacks import Callback # type: ignore + +import wandb +from wandb.sdk.lib import telemetry + + +class WandbEvalCallback(Callback, abc.ABC): + """Abstract base class to build Keras callbacks for model prediction visualization. + + You can build callbacks for visualizing model predictions `on_epoch_end` + that can be passed to `model.fit()` for classification, object detection, + segmentation, etc. tasks. + + To use this, inherit from this base callback class and implement the + `add_ground_truth` and `add_model_prediction` methods. + + The base class will take care of the following: + - Initialize `data_table` for logging the ground truth and + `pred_table` for predictions. + - The data uploaded to `data_table` is used as a reference for the + `pred_table`. This is to reduce the memory footprint. The `data_table_ref` + is a list that can be used to access the referenced data. + Check out the example below to see how it's done. + - Log the tables to W&B as W&B Artifacts. + - Each new `pred_table` is logged as a new version with aliases. + + Example: + ```python + class WandbClfEvalCallback(WandbEvalCallback): + def __init__(self, validation_data, data_table_columns, pred_table_columns): + super().__init__(data_table_columns, pred_table_columns) + + self.x = validation_data[0] + self.y = validation_data[1] + + def add_ground_truth(self): + for idx, (image, label) in enumerate(zip(self.x, self.y)): + self.data_table.add_data(idx, wandb.Image(image), label) + + def add_model_predictions(self, epoch): + preds = self.model.predict(self.x, verbose=0) + preds = tf.argmax(preds, axis=-1) + + data_table_ref = self.data_table_ref + table_idxs = data_table_ref.get_index() + + for idx in table_idxs: + pred = preds[idx] + self.pred_table.add_data( + epoch, + data_table_ref.data[idx][0], + data_table_ref.data[idx][1], + data_table_ref.data[idx][2], + pred, + ) + + + model.fit( + x, + y, + epochs=2, + validation_data=(x, y), + callbacks=[ + WandbClfEvalCallback( + validation_data=(x, y), + data_table_columns=["idx", "image", "label"], + pred_table_columns=["epoch", "idx", "image", "label", "pred"], + ) + ], + ) + ``` + + To have more fine-grained control, you can override the `on_train_begin` and + `on_epoch_end` methods. If you want to log the samples after N batched, you + can implement `on_train_batch_end` method. + """ + + def __init__( + self, + data_table_columns: List[str], + pred_table_columns: List[str], + *args: Any, + **kwargs: Any, + ) -> None: + super().__init__(*args, **kwargs) + + if wandb.run is None: + raise wandb.Error( + "You must call `wandb.init()` first before using this callback." + ) + + with telemetry.context(run=wandb.run) as tel: + tel.feature.keras_wandb_eval_callback = True + + self.data_table_columns = data_table_columns + self.pred_table_columns = pred_table_columns + + def on_train_begin(self, logs: Optional[Dict[str, float]] = None) -> None: + # Initialize the data_table + self.init_data_table(column_names=self.data_table_columns) + # Log the ground truth data + self.add_ground_truth(logs) + # Log the data_table as W&B Artifacts + self.log_data_table() + + def on_epoch_end(self, epoch: int, logs: Optional[Dict[str, float]] = None) -> None: + # Initialize the pred_table + self.init_pred_table(column_names=self.pred_table_columns) + # Log the model prediction + self.add_model_predictions(epoch, logs) + # Log the pred_table as W&B Artifacts + self.log_pred_table() + + @abc.abstractmethod + def add_ground_truth(self, logs: Optional[Dict[str, float]] = None) -> None: + """Add ground truth data to `data_table`. + + Use this method to write the logic for adding validation/training data to + `data_table` initialized using `init_data_table` method. + + Example: + ```python + for idx, data in enumerate(dataloader): + self.data_table.add_data(idx, data) + ``` + This method is called once `on_train_begin` or equivalent hook. + """ + raise NotImplementedError(f"{self.__class__.__name__}.add_ground_truth") + + @abc.abstractmethod + def add_model_predictions( + self, epoch: int, logs: Optional[Dict[str, float]] = None + ) -> None: + """Add a prediction from a model to `pred_table`. + + Use this method to write the logic for adding model prediction for validation/ + training data to `pred_table` initialized using `init_pred_table` method. + + Example: + ```python + # Assuming the dataloader is not shuffling the samples. + for idx, data in enumerate(dataloader): + preds = model.predict(data) + self.pred_table.add_data( + self.data_table_ref.data[idx][0], + self.data_table_ref.data[idx][1], + preds, + ) + ``` + This method is called `on_epoch_end` or equivalent hook. + """ + raise NotImplementedError(f"{self.__class__.__name__}.add_model_predictions") + + def init_data_table(self, column_names: List[str]) -> None: + """Initialize the W&B Tables for validation data. + + Call this method `on_train_begin` or equivalent hook. This is followed by adding + data to the table row or column wise. + + Args: + column_names: (list) Column names for W&B Tables. + """ + self.data_table = wandb.Table(columns=column_names, allow_mixed_types=True) + + def init_pred_table(self, column_names: List[str]) -> None: + """Initialize the W&B Tables for model evaluation. + + Call this method `on_epoch_end` or equivalent hook. This is followed by adding + data to the table row or column wise. + + Args: + column_names: (list) Column names for W&B Tables. + """ + self.pred_table = wandb.Table(columns=column_names) + + def log_data_table( + self, name: str = "val", type: str = "dataset", table_name: str = "val_data" + ) -> None: + """Log the `data_table` as W&B artifact and call `use_artifact` on it. + + This lets the evaluation table use the reference of already uploaded data + (images, text, scalar, etc.) without re-uploading. + + Args: + name: (str) A human-readable name for this artifact, which is how you can + identify this artifact in the UI or reference it in use_artifact calls. + (default is 'val') + type: (str) The type of the artifact, which is used to organize and + differentiate artifacts. (default is 'dataset') + table_name: (str) The name of the table as will be displayed in the UI. + (default is 'val_data'). + """ + data_artifact = wandb.Artifact(name, type=type) + data_artifact.add(self.data_table, table_name) + + # Calling `use_artifact` uploads the data to W&B. + assert wandb.run is not None + wandb.run.use_artifact(data_artifact) + data_artifact.wait() + + # We get the reference table. + self.data_table_ref = data_artifact.get(table_name) + + def log_pred_table( + self, + type: str = "evaluation", + table_name: str = "eval_data", + aliases: Optional[List[str]] = None, + ) -> None: + """Log the W&B Tables for model evaluation. + + The table will be logged multiple times creating new version. Use this + to compare models at different intervals interactively. + + Args: + type: (str) The type of the artifact, which is used to organize and + differentiate artifacts. (default is 'evaluation') + table_name: (str) The name of the table as will be displayed in the UI. + (default is 'eval_data') + aliases: (List[str]) List of aliases for the prediction table. + """ + assert wandb.run is not None + pred_artifact = wandb.Artifact(f"run_{wandb.run.id}_pred", type=type) + pred_artifact.add(self.pred_table, table_name) + wandb.run.log_artifact(pred_artifact, aliases=aliases or ["latest"]) diff --git a/lib/python3.12/site-packages/wandb/integration/keras/keras.py b/lib/python3.12/site-packages/wandb/integration/keras/keras.py new file mode 100644 index 0000000000000000000000000000000000000000..3dde8f71ed6a64c24790704ad4f3bf97e83ddd83 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/keras/keras.py @@ -0,0 +1,1086 @@ +"""keras init.""" + +import logging +import operator +import os +import shutil +import sys +from itertools import chain + +import numpy as np +import tensorflow as tf +import tensorflow.keras.backend as K # noqa: N812 + +import wandb +from wandb.proto.wandb_deprecated import Deprecated +from wandb.sdk.integration_utils.data_logging import ValidationDataLogger +from wandb.sdk.lib.deprecate import deprecate +from wandb.util import add_import_hook + + +def _check_keras_version(): + from keras import __version__ as keras_version + from packaging.version import parse + + if parse(keras_version) < parse("2.4.0"): + wandb.termwarn( + f"Keras version {keras_version} is not fully supported. Required keras >= 2.4.0" + ) + + +def _can_compute_flops() -> bool: + """FLOPS computation is restricted to TF 2.x as it requires tf.compat.v1.""" + from packaging.version import parse + + if parse(tf.__version__) >= parse("2.0.0"): + return True + + return False + + +if "keras" in sys.modules: + _check_keras_version() +else: + add_import_hook("keras", _check_keras_version) + + +logger = logging.getLogger(__name__) + + +def is_dataset(data): + dataset_ops = wandb.util.get_module("tensorflow.python.data.ops.dataset_ops") + if dataset_ops and hasattr(dataset_ops, "DatasetV2"): + dataset_types = (dataset_ops.DatasetV2,) + if hasattr(dataset_ops, "DatasetV1"): + dataset_types = dataset_types + (dataset_ops.DatasetV1,) + return isinstance(data, dataset_types) + else: + return False + + +def is_generator_like(data): + # Checks if data is a generator, Sequence, or Iterator. + + types = (tf.keras.utils.Sequence,) + iterator_ops = wandb.util.get_module("tensorflow.python.data.ops.iterator_ops") + if iterator_ops: + types = types + (iterator_ops.Iterator,) + # EagerIterator was in tensorflow < 2 + if hasattr(iterator_ops, "EagerIterator"): + types = types + (iterator_ops.EagerIterator,) + elif hasattr(iterator_ops, "IteratorV2"): + types = types + (iterator_ops.IteratorV2,) + return hasattr(data, "next") or hasattr(data, "__next__") or isinstance(data, types) + + +def patch_tf_keras(): # noqa: C901 + from packaging.version import parse + from tensorflow.python.eager import context + + if parse("2.6.0") <= parse(tf.__version__) < parse("2.13.0"): + keras_engine = "keras.engine" + try: + from keras.engine import training + from keras.engine import training_arrays_v1 as training_arrays + from keras.engine import training_generator_v1 as training_generator + except (ImportError, AttributeError): + wandb.termerror("Unable to patch Tensorflow/Keras") + logger.exception("exception while trying to patch_tf_keras") + return + else: + keras_engine = "tensorflow.python.keras.engine" + + from tensorflow.python.keras.engine import training + + try: + from tensorflow.python.keras.engine import ( + training_arrays_v1 as training_arrays, + ) + from tensorflow.python.keras.engine import ( + training_generator_v1 as training_generator, + ) + except (ImportError, AttributeError): + try: + from tensorflow.python.keras.engine import ( + training_arrays, + training_generator, + ) + except (ImportError, AttributeError): + wandb.termerror("Unable to patch Tensorflow/Keras") + logger.exception("exception while trying to patch_tf_keras") + return + + # Tensorflow 2.1 + training_v2_1 = wandb.util.get_module("tensorflow.python.keras.engine.training_v2") + # Tensorflow 2.2 + training_v2_2 = wandb.util.get_module(f"{keras_engine}.training_v1") + + if training_v2_1: + old_v2 = training_v2_1.Loop.fit + elif training_v2_2: + old_v2 = training.Model.fit + + old_arrays = training_arrays.fit_loop + old_generator = training_generator.fit_generator + + def set_wandb_attrs(cbk, val_data): + if isinstance(cbk, WandbCallback): + if is_generator_like(val_data): + cbk.generator = val_data + elif is_dataset(val_data): + if context.executing_eagerly(): + cbk.generator = iter(val_data) + else: + wandb.termwarn( + "Found a validation dataset in graph mode, can't patch Keras." + ) + elif isinstance(val_data, tuple) and isinstance(val_data[0], tf.Tensor): + # Graph mode dataset generator + def gen(): + while True: + yield K.get_session().run(val_data) + + cbk.generator = gen() + else: + cbk.validation_data = val_data + + def new_arrays(*args, **kwargs): + cbks = kwargs.get("callbacks", []) + val_inputs = kwargs.get("val_inputs") + val_targets = kwargs.get("val_targets") + # TODO: these could be generators, why index 0? + if val_inputs and val_targets: + for cbk in cbks: + set_wandb_attrs(cbk, (val_inputs[0], val_targets[0])) + return old_arrays(*args, **kwargs) + + def new_generator(*args, **kwargs): + cbks = kwargs.get("callbacks", []) + val_data = kwargs.get("validation_data") + if val_data: + for cbk in cbks: + set_wandb_attrs(cbk, val_data) + return old_generator(*args, **kwargs) + + def new_v2(*args, **kwargs): + cbks = kwargs.get("callbacks", []) + val_data = kwargs.get("validation_data") + if val_data: + for cbk in cbks: + set_wandb_attrs(cbk, val_data) + return old_v2(*args, **kwargs) + + training_arrays.orig_fit_loop = old_arrays + training_arrays.fit_loop = new_arrays + training_generator.orig_fit_generator = old_generator + training_generator.fit_generator = new_generator + wandb.patched["keras"].append([f"{keras_engine}.training_arrays", "fit_loop"]) + wandb.patched["keras"].append( + [f"{keras_engine}.training_generator", "fit_generator"] + ) + + if training_v2_1: + training_v2_1.Loop.fit = new_v2 + wandb.patched["keras"].append( + ["tensorflow.python.keras.engine.training_v2.Loop", "fit"] + ) + elif training_v2_2: + training.Model.fit = new_v2 + wandb.patched["keras"].append([f"{keras_engine}.training.Model", "fit"]) + + +def _array_has_dtype(array): + return hasattr(array, "dtype") + + +def _update_if_numeric(metrics, key, values): + if not _array_has_dtype(values): + _warn_not_logging(key) + return + + if not is_numeric_array(values): + _warn_not_logging_non_numeric(key) + return + + metrics[key] = wandb.Histogram(values) + + +def is_numeric_array(array): + return np.issubdtype(array.dtype, np.number) + + +def _warn_not_logging_non_numeric(name): + wandb.termwarn( + f"Non-numeric values found in layer: {name}, not logging this layer", + repeat=False, + ) + + +def _warn_not_logging(name): + wandb.termwarn( + f"Layer {name} has undetermined datatype not logging this layer", + repeat=False, + ) + + +tf_logger = tf.get_logger() + +patch_tf_keras() + + +### For gradient logging ### + + +def _get_custom_optimizer_parent_class(): + from packaging.version import parse + + if parse(tf.__version__) >= parse("2.9.0"): + custom_optimizer_parent_class = tf.keras.optimizers.legacy.Optimizer + else: + custom_optimizer_parent_class = tf.keras.optimizers.Optimizer + + return custom_optimizer_parent_class + + +_custom_optimizer_parent_class = _get_custom_optimizer_parent_class() + + +class _CustomOptimizer(_custom_optimizer_parent_class): + def __init__(self): + super().__init__(name="CustomOptimizer") + self._resource_apply_dense = tf.function(self._resource_apply_dense) + self._resource_apply_sparse = tf.function(self._resource_apply_sparse) + + def _resource_apply_dense(self, grad, var): + var.assign(grad) + + # this needs to be implemented to prevent a NotImplementedError when + # using Lookup layers. + def _resource_apply_sparse(self, grad, var, indices): + pass + + def get_config(self): + return super().get_config() + + +class _GradAccumulatorCallback(tf.keras.callbacks.Callback): + """Accumulates gradients during a fit() call when used in conjunction with the CustomOptimizer above.""" + + def set_model(self, model): + super().set_model(model) + self.og_weights = model.get_weights() + self.grads = [np.zeros(tuple(w.shape)) for w in model.trainable_weights] + + def on_batch_end(self, batch, logs=None): + for g, w in zip(self.grads, self.model.trainable_weights): + g += w.numpy() + self.model.set_weights(self.og_weights) + + def get_grads(self): + return [g.copy() for g in self.grads] + + +### + + +class WandbCallback(tf.keras.callbacks.Callback): + """`WandbCallback` automatically integrates keras with wandb. + + Example: + ```python + model.fit( + X_train, + y_train, + validation_data=(X_test, y_test), + callbacks=[WandbCallback()], + ) + ``` + + `WandbCallback` will automatically log history data from any + metrics collected by keras: loss and anything passed into `keras_model.compile()`. + + `WandbCallback` will set summary metrics for the run associated with the "best" training + step, where "best" is defined by the `monitor` and `mode` attributes. This defaults + to the epoch with the minimum `val_loss`. `WandbCallback` will by default save the model + associated with the best `epoch`. + + `WandbCallback` can optionally log gradient and parameter histograms. + + `WandbCallback` can optionally save training and validation data for wandb to visualize. + + Args: + monitor: (str) name of metric to monitor. Defaults to `val_loss`. + mode: (str) one of {`auto`, `min`, `max`}. + `min` - save model when monitor is minimized + `max` - save model when monitor is maximized + `auto` - try to guess when to save the model (default). + save_model: + True - save a model when monitor beats all previous epochs + False - don't save models + save_graph: (boolean) if True save model graph to wandb (default to True). + save_weights_only: (boolean) if True, then only the model's weights will be + saved (`model.save_weights(filepath)`), else the full model + is saved (`model.save(filepath)`). + log_weights: (boolean) if True save histograms of the model's layer's weights. + log_gradients: (boolean) if True log histograms of the training gradients + training_data: (tuple) Same format `(X,y)` as passed to `model.fit`. This is needed + for calculating gradients - this is mandatory if `log_gradients` is `True`. + validation_data: (tuple) Same format `(X,y)` as passed to `model.fit`. A set of data + for wandb to visualize. If this is set, every epoch, wandb will + make a small number of predictions and save the results for later visualization. In case + you are working with image data, please also set `input_type` and `output_type` in order + to log correctly. + generator: (generator) a generator that returns validation data for wandb to visualize. This + generator should return tuples `(X,y)`. Either `validate_data` or generator should + be set for wandb to visualize specific data examples. In case you are working with image data, + please also set `input_type` and `output_type` in order to log correctly. + validation_steps: (int) if `validation_data` is a generator, how many + steps to run the generator for the full validation set. + labels: (list) If you are visualizing your data with wandb this list of labels + will convert numeric output to understandable string if you are building a + multiclass classifier. If you are making a binary classifier you can pass in + a list of two labels ["label for false", "label for true"]. If `validate_data` + and generator are both false, this won't do anything. + predictions: (int) the number of predictions to make for visualization each epoch, max + is 100. + input_type: (string) type of the model input to help visualization. can be one of: + (`image`, `images`, `segmentation_mask`, `auto`). + output_type: (string) type of the model output to help visualization. can be one of: + (`image`, `images`, `segmentation_mask`, `label`). + log_evaluation: (boolean) if True, save a Table containing validation data and the + model's predictions at each epoch. See `validation_indexes`, + `validation_row_processor`, and `output_row_processor` for additional details. + class_colors: ([float, float, float]) if the input or output is a segmentation mask, + an array containing an rgb tuple (range 0-1) for each class. + log_batch_frequency: (integer) if None, callback will log every epoch. + If set to integer, callback will log training metrics every `log_batch_frequency` + batches. + log_best_prefix: (string) if None, no extra summary metrics will be saved. + If set to a string, the monitored metric and epoch will be prepended with this value + and stored as summary metrics. + validation_indexes: ([wandb.data_types._TableLinkMixin]) an ordered list of index keys to associate + with each validation example. If log_evaluation is True and `validation_indexes` is provided, + then a Table of validation data will not be created and instead each prediction will + be associated with the row represented by the `TableLinkMixin`. The most common way to obtain + such keys are is use `Table.get_index()` which will return a list of row keys. + validation_row_processor: (Callable) a function to apply to the validation data, commonly used to visualize the data. + The function will receive an `ndx` (int) and a `row` (dict). If your model has a single input, + then `row["input"]` will be the input data for the row. Else, it will be keyed based on the name of the + input slot. If your fit function takes a single target, then `row["target"]` will be the target data for the row. Else, + it will be keyed based on the name of the output slots. For example, if your input data is a single ndarray, + but you wish to visualize the data as an Image, then you can provide `lambda ndx, row: {"img": wandb.Image(row["input"])}` + as the processor. Ignored if log_evaluation is False or `validation_indexes` are present. + output_row_processor: (Callable) same as `validation_row_processor`, but applied to the model's output. `row["output"]` will contain + the results of the model output. + infer_missing_processors: (bool) Determines if `validation_row_processor` and `output_row_processor` + should be inferred if missing. Defaults to True. If `labels` are provided, we will attempt to infer classification-type + processors where appropriate. + log_evaluation_frequency: (int) Determines the frequency which evaluation results will be logged. Default 0 (only at the end of training). + Set to 1 to log every epoch, 2 to log every other epoch, and so on. Has no effect when log_evaluation is False. + compute_flops: (bool) Compute the FLOPs of your Keras Sequential or Functional model in GigaFLOPs unit. + """ + + def __init__( + self, + monitor="val_loss", + verbose=0, + mode="auto", + save_weights_only=False, + log_weights=False, + log_gradients=False, + save_model=True, + training_data=None, + validation_data=None, + labels=None, + predictions=36, + generator=None, + input_type=None, + output_type=None, + log_evaluation=False, + validation_steps=None, + class_colors=None, + log_batch_frequency=None, + log_best_prefix="best_", + save_graph=True, + validation_indexes=None, + validation_row_processor=None, + prediction_row_processor=None, + infer_missing_processors=True, + log_evaluation_frequency=0, + compute_flops=False, + **kwargs, + ): + if wandb.run is None: + raise wandb.Error("You must call wandb.init() before WandbCallback()") + + deprecate( + field_name=Deprecated.keras_callback, + warning_message=( + "WandbCallback is deprecated and will be removed in a future release. " + "Please use the WandbMetricsLogger, WandbModelCheckpoint, and WandbEvalCallback " + "callbacks instead. " + "See https://docs.wandb.ai/guides/integrations/keras for more information." + ), + ) + + with wandb.wandb_lib.telemetry.context(run=wandb.run) as tel: + tel.feature.keras = True + self.validation_data = None + # This is kept around for legacy reasons + if validation_data is not None: + if is_generator_like(validation_data): + generator = validation_data + else: + self.validation_data = validation_data + if labels is None: + labels = [] + self.labels = labels + self.predictions = min(predictions, 100) + + self.monitor = monitor + self.verbose = verbose + self.save_weights_only = save_weights_only + self.save_graph = save_graph + + wandb.save("model-best.h5") + self.filepath = os.path.join(wandb.run.dir, "model-best.h5") + self.save_model = save_model + if save_model: + deprecate( + field_name=Deprecated.keras_callback__save_model, + warning_message=( + "The save_model argument by default saves the model in the HDF5 format that cannot save " + "custom objects like subclassed models and custom layers. This behavior will be deprecated " + "in a future release in favor of the SavedModel format. Meanwhile, the HDF5 model is saved " + "as W&B files and the SavedModel as W&B Artifacts." + ), + ) + + self.save_model_as_artifact = True + self.log_weights = log_weights + self.log_gradients = log_gradients + self.training_data = training_data + self.generator = generator + self._graph_rendered = False + + data_type = kwargs.get("data_type", None) + if data_type is not None: + deprecate( + field_name=Deprecated.keras_callback__data_type, + warning_message=( + "The data_type argument of wandb.keras.WandbCallback is deprecated " + "and will be removed in a future release. Please use input_type instead.\n" + "Setting input_type = data_type." + ), + ) + input_type = data_type + self.input_type = input_type + self.output_type = output_type + self.log_evaluation = log_evaluation + self.validation_steps = validation_steps + self.class_colors = np.array(class_colors) if class_colors is not None else None + self.log_batch_frequency = log_batch_frequency + self.log_best_prefix = log_best_prefix + self.compute_flops = compute_flops + + self._prediction_batch_size = None + + if self.log_gradients: + if int(tf.__version__.split(".")[0]) < 2: + raise Exception("Gradient logging requires tensorflow 2.0 or higher.") + if self.training_data is None: + raise ValueError( + "training_data argument is required for gradient logging." + ) + if isinstance(self.training_data, (list, tuple)): + if len(self.training_data) != 2: + raise ValueError("training data must be a tuple of length two") + self._training_data_x, self._training_data_y = self.training_data + else: + self._training_data_x = ( + self.training_data + ) # generator, tf.data.Dataset etc + self._training_data_y = None + + # From Keras + if mode not in ["auto", "min", "max"]: + wandb.termwarn( + f"WandbCallback mode {mode} is unknown, fallback to auto mode." + ) + mode = "auto" + + if mode == "min": + self.monitor_op = operator.lt + self.best = float("inf") + elif mode == "max": + self.monitor_op = operator.gt + self.best = float("-inf") + else: + if "acc" in self.monitor or self.monitor.startswith("fmeasure"): + self.monitor_op = operator.gt + self.best = float("-inf") + else: + self.monitor_op = operator.lt + self.best = float("inf") + # Get the previous best metric for resumed runs + previous_best = wandb.run.summary.get(f"{self.log_best_prefix}{self.monitor}") + if previous_best is not None: + self.best = previous_best + + self._validation_data_logger = None + self._validation_indexes = validation_indexes + self._validation_row_processor = validation_row_processor + self._prediction_row_processor = prediction_row_processor + self._infer_missing_processors = infer_missing_processors + self._log_evaluation_frequency = log_evaluation_frequency + self._model_trained_since_last_eval = False + + def _build_grad_accumulator_model(self): + inputs = self.model.inputs + outputs = self.model(inputs) + grad_acc_model = tf.keras.models.Model(inputs, outputs) + grad_acc_model.compile(loss=self.model.loss, optimizer=_CustomOptimizer()) + + # make sure magic doesn't think this is a user model + grad_acc_model._wandb_internal_model = True + + self._grad_accumulator_model = grad_acc_model + self._grad_accumulator_callback = _GradAccumulatorCallback() + + def _implements_train_batch_hooks(self): + return self.log_batch_frequency is not None + + def _implements_test_batch_hooks(self): + return self.log_batch_frequency is not None + + def _implements_predict_batch_hooks(self): + return self.log_batch_frequency is not None + + def set_params(self, params): + self.params = params + + def set_model(self, model): + super().set_model(model) + if self.input_type == "auto" and len(model.inputs) == 1: + self.input_type = wandb.util.guess_data_type( + model.inputs[0].shape, risky=True + ) + if self.input_type and self.output_type is None and len(model.outputs) == 1: + self.output_type = wandb.util.guess_data_type(model.outputs[0].shape) + if self.log_gradients: + self._build_grad_accumulator_model() + + def _attempt_evaluation_log(self, commit=True): + if self.log_evaluation and self._validation_data_logger: + try: + if not self.model: + wandb.termwarn("WandbCallback unable to read model from trainer") + else: + self._validation_data_logger.log_predictions( + predictions=self._validation_data_logger.make_predictions( + self.model.predict + ), + commit=commit, + ) + self._model_trained_since_last_eval = False + except Exception as e: + wandb.termwarn("Error during prediction logging for epoch: " + str(e)) + + def on_epoch_end(self, epoch, logs=None): + if logs is None: + logs = {} + if self.log_weights: + wandb.log(self._log_weights(), commit=False) + + if self.log_gradients: + wandb.log(self._log_gradients(), commit=False) + + if self.input_type in ( + "image", + "images", + "segmentation_mask", + ) or self.output_type in ("image", "images", "segmentation_mask"): + if self.generator: + self.validation_data = next(self.generator) + if self.validation_data is None: + wandb.termwarn( + "No validation_data set, pass a generator to the callback." + ) + elif self.validation_data and len(self.validation_data) > 0: + wandb.log( + {"examples": self._log_images(num_images=self.predictions)}, + commit=False, + ) + + if ( + self._log_evaluation_frequency > 0 + and epoch % self._log_evaluation_frequency == 0 + ): + self._attempt_evaluation_log(commit=False) + + wandb.log({"epoch": epoch}, commit=False) + wandb.log(logs, commit=True) + + self.current = logs.get(self.monitor) + if self.current and self.monitor_op(self.current, self.best): + if self.log_best_prefix: + wandb.run.summary[f"{self.log_best_prefix}{self.monitor}"] = ( + self.current + ) + wandb.run.summary["{}{}".format(self.log_best_prefix, "epoch")] = epoch + if self.verbose and not self.save_model: + wandb.termlog( + f"Epoch {epoch:05d}: {self.monitor} improved from {self.best:.5f} to {self.current:.5f}" + ) + if self.save_model: + self._save_model(epoch) + + if self.save_model and self.save_model_as_artifact: + self._save_model_as_artifact(epoch) + + self.best = self.current + + # This is what keras used pre tensorflow.keras + def on_batch_begin(self, batch, logs=None): + pass + + # This is what keras used pre tensorflow.keras + def on_batch_end(self, batch, logs=None): + if self.save_graph and not self._graph_rendered: + # Couldn't do this in train_begin because keras may still not be built + wandb.run.summary["graph"] = wandb.Graph.from_keras(self.model) + self._graph_rendered = True + + if self.log_batch_frequency and batch % self.log_batch_frequency == 0: + wandb.log(logs, commit=True) + + def on_train_batch_begin(self, batch, logs=None): + self._model_trained_since_last_eval = True + + def on_train_batch_end(self, batch, logs=None): + if self.save_graph and not self._graph_rendered: + # Couldn't do this in train_begin because keras may still not be built + wandb.run.summary["graph"] = wandb.Graph.from_keras(self.model) + self._graph_rendered = True + + if self.log_batch_frequency and batch % self.log_batch_frequency == 0: + wandb.log(logs, commit=True) + + def on_test_begin(self, logs=None): + pass + + def on_test_end(self, logs=None): + pass + + def on_test_batch_begin(self, batch, logs=None): + pass + + def on_test_batch_end(self, batch, logs=None): + pass + + def on_train_begin(self, logs=None): + if self.log_evaluation: + try: + validation_data = None + if self.validation_data: + validation_data = self.validation_data + elif self.generator: + if not self.validation_steps: + wandb.termwarn( + "WandbCallback is unable to log validation data. " + "When using a generator for validation_data, you must pass validation_steps" + ) + else: + x = None + y_true = None + for _ in range(self.validation_steps): + bx, by_true = next(self.generator) + if x is None: + x, y_true = bx, by_true + else: + x, y_true = ( + np.append(x, bx, axis=0), + np.append(y_true, by_true, axis=0), + ) + validation_data = (x, y_true) + else: + wandb.termwarn( + "WandbCallback is unable to read validation_data from trainer " + "and therefore cannot log validation data. Ensure Keras is properly " + "patched by calling `from wandb.keras import WandbCallback` at the top of your script." + ) + if validation_data: + self._validation_data_logger = ValidationDataLogger( + inputs=validation_data[0], + targets=validation_data[1], + indexes=self._validation_indexes, + validation_row_processor=self._validation_row_processor, + prediction_row_processor=self._prediction_row_processor, + class_labels=self.labels, + infer_missing_processors=self._infer_missing_processors, + ) + except Exception as e: + wandb.termwarn( + "Error initializing ValidationDataLogger in WandbCallback. " + f"Skipping logging validation data. Error: {str(e)}" + ) + + if self.compute_flops and _can_compute_flops(): + try: + wandb.summary["GFLOPs"] = self.get_flops() + except Exception: + logger.exception("Error computing FLOPs") + wandb.termwarn("Unable to compute FLOPs for this model.") + + def on_train_end(self, logs=None): + if self._model_trained_since_last_eval: + self._attempt_evaluation_log() + + def on_predict_begin(self, logs=None): + pass + + def on_predict_end(self, logs=None): + pass + + def on_predict_batch_begin(self, batch, logs=None): + pass + + def on_predict_batch_end(self, batch, logs=None): + pass + + def _logits_to_captions(self, logits): + if logits[0].shape[-1] == 1: + # Scalar output from the model + # TODO: handle validation_y + if len(self.labels) == 2: + # User has named true and false + captions = [ + self.labels[1] if logits[0] > 0.5 else self.labels[0] + for logit in logits + ] + else: + if len(self.labels) != 0: + wandb.termwarn( + "keras model is producing a single output, " + 'so labels should be a length two array: ["False label", "True label"].' + ) + captions = [logit[0] for logit in logits] + else: + # Vector output from the model + # TODO: handle validation_y + labels = np.argmax(np.stack(logits), axis=1) + + if len(self.labels) > 0: + # User has named the categories in self.labels + captions = [] + for label in labels: + try: + captions.append(self.labels[label]) + except IndexError: + captions.append(label) + else: + captions = labels + return captions + + def _masks_to_pixels(self, masks): + # if its a binary mask, just return it as grayscale instead of picking the argmax + if len(masks[0].shape) == 2 or masks[0].shape[-1] == 1: + return masks + class_colors = ( + self.class_colors + if self.class_colors is not None + else np.array(wandb.util.class_colors(masks[0].shape[2])) + ) + imgs = class_colors[np.argmax(masks, axis=-1)] + return imgs + + def _log_images(self, num_images=36): + validation_X = self.validation_data[0] # noqa: N806 + validation_y = self.validation_data[1] + + validation_length = len(validation_X) + + if validation_length > num_images: + # pick some data at random + indices = np.random.choice(validation_length, num_images, replace=False) + else: + indices = range(validation_length) + + test_data = [] + test_output = [] + for i in indices: + test_example = validation_X[i] + test_data.append(test_example) + test_output.append(validation_y[i]) + + if self.model.stateful: + predictions = self.model.predict(np.stack(test_data), batch_size=1) + self.model.reset_states() + else: + predictions = self.model.predict( + np.stack(test_data), batch_size=self._prediction_batch_size + ) + if len(predictions) != len(test_data): + self._prediction_batch_size = 1 + predictions = self.model.predict( + np.stack(test_data), batch_size=self._prediction_batch_size + ) + + if self.input_type == "label": + if self.output_type in ("image", "images", "segmentation_mask"): + captions = self._logits_to_captions(test_data) + output_image_data = ( + self._masks_to_pixels(predictions) + if self.output_type == "segmentation_mask" + else predictions + ) + reference_image_data = ( + self._masks_to_pixels(test_output) + if self.output_type == "segmentation_mask" + else test_output + ) + output_images = [ + wandb.Image(data, caption=captions[i], grouping=2) + for i, data in enumerate(output_image_data) + ] + reference_images = [ + wandb.Image(data, caption=captions[i]) + for i, data in enumerate(reference_image_data) + ] + return list(chain.from_iterable(zip(output_images, reference_images))) + elif self.input_type in ("image", "images", "segmentation_mask"): + input_image_data = ( + self._masks_to_pixels(test_data) + if self.input_type == "segmentation_mask" + else test_data + ) + if self.output_type == "label": + # we just use the predicted label as the caption for now + captions = self._logits_to_captions(predictions) + return [ + wandb.Image(data, caption=captions[i]) + for i, data in enumerate(test_data) + ] + elif self.output_type in ("image", "images", "segmentation_mask"): + output_image_data = ( + self._masks_to_pixels(predictions) + if self.output_type == "segmentation_mask" + else predictions + ) + reference_image_data = ( + self._masks_to_pixels(test_output) + if self.output_type == "segmentation_mask" + else test_output + ) + input_images = [ + wandb.Image(data, grouping=3) + for i, data in enumerate(input_image_data) + ] + output_images = [ + wandb.Image(data) for i, data in enumerate(output_image_data) + ] + reference_images = [ + wandb.Image(data) for i, data in enumerate(reference_image_data) + ] + return list( + chain.from_iterable( + zip(input_images, output_images, reference_images) + ) + ) + else: + # unknown output, just log the input images + return [wandb.Image(img) for img in test_data] + elif self.output_type in ("image", "images", "segmentation_mask"): + # unknown input, just log the predicted and reference outputs without captions + output_image_data = ( + self._masks_to_pixels(predictions) + if self.output_type == "segmentation_mask" + else predictions + ) + reference_image_data = ( + self._masks_to_pixels(test_output) + if self.output_type == "segmentation_mask" + else test_output + ) + output_images = [ + wandb.Image(data, grouping=2) + for i, data in enumerate(output_image_data) + ] + reference_images = [ + wandb.Image(data) for i, data in enumerate(reference_image_data) + ] + return list(chain.from_iterable(zip(output_images, reference_images))) + + def _log_weights(self): + metrics = {} + for layer in self.model.layers: + weights = layer.get_weights() + if len(weights) == 1: + _update_if_numeric( + metrics, "parameters/" + layer.name + ".weights", weights[0] + ) + elif len(weights) == 2: + _update_if_numeric( + metrics, "parameters/" + layer.name + ".weights", weights[0] + ) + _update_if_numeric( + metrics, "parameters/" + layer.name + ".bias", weights[1] + ) + return metrics + + def _log_gradients(self): + # Suppress callback warnings grad accumulator + og_level = tf_logger.level + tf_logger.setLevel("ERROR") + + self._grad_accumulator_model.fit( + self._training_data_x, + self._training_data_y, + verbose=0, + callbacks=[self._grad_accumulator_callback], + ) + tf_logger.setLevel(og_level) + weights = self.model.trainable_weights + grads = self._grad_accumulator_callback.grads + metrics = {} + for weight, grad in zip(weights, grads): + metrics["gradients/" + weight.name.split(":")[0] + ".gradient"] = ( + wandb.Histogram(grad) + ) + return metrics + + def _log_dataframe(self): + x, y_true, y_pred = None, None, None + + if self.validation_data: + x, y_true = self.validation_data[0], self.validation_data[1] + y_pred = self.model.predict(x) + elif self.generator: + if not self.validation_steps: + wandb.termwarn( + "when using a generator for validation data with dataframes, " + "you must pass validation_steps. skipping" + ) + return None + + for _ in range(self.validation_steps): + bx, by_true = next(self.generator) + by_pred = self.model.predict(bx) + if x is None: + x, y_true, y_pred = bx, by_true, by_pred + else: + x, y_true, y_pred = ( + np.append(x, bx, axis=0), + np.append(y_true, by_true, axis=0), + np.append(y_pred, by_pred, axis=0), + ) + + if self.input_type in ("image", "images") and self.output_type == "label": + return wandb.image_categorizer_dataframe( + x=x, y_true=y_true, y_pred=y_pred, labels=self.labels + ) + elif ( + self.input_type in ("image", "images") + and self.output_type == "segmentation_mask" + ): + return wandb.image_segmentation_dataframe( + x=x, + y_true=y_true, + y_pred=y_pred, + labels=self.labels, + class_colors=self.class_colors, + ) + else: + wandb.termwarn( + f"unknown dataframe type for input_type={self.input_type} and output_type={self.output_type}" + ) + return None + + def _save_model(self, epoch): + if wandb.run.disabled: + return + if self.verbose > 0: + wandb.termlog( + f"Epoch {epoch:05d}: {self.monitor} improved from {self.best:.5f} to {self.current:.5f}, " + f"saving model to {self.filepath}" + ) + + try: + if self.save_weights_only: + self.model.save_weights(self.filepath, overwrite=True) + else: + self.model.save(self.filepath, overwrite=True) + # Was getting `RuntimeError: Unable to create link` in TF 1.13.1 + # also saw `TypeError: can't pickle _thread.RLock objects` + except (ImportError, RuntimeError, TypeError, AttributeError): + logger.exception("Error saving model in the h5py format") + wandb.termerror( + "Can't save model in the h5py format. The model will be saved as " + "as an W&B Artifact in the 'tf' format." + ) + + def _save_model_as_artifact(self, epoch): + if wandb.run.disabled: + return + + # Save the model in the SavedModel format. + # TODO: Replace this manual artifact creation with the `log_model` method + # after `log_model` is released from beta. + self.model.save(self.filepath[:-3], overwrite=True, save_format="tf") + + # Log the model as artifact. + name = wandb.util.make_artifact_name_safe(f"model-{wandb.run.name}") + model_artifact = wandb.Artifact(name, type="model") + model_artifact.add_dir(self.filepath[:-3]) + wandb.run.log_artifact(model_artifact, aliases=["latest", f"epoch_{epoch}"]) + + # Remove the SavedModel from wandb dir as we don't want to log it to save memory. + shutil.rmtree(self.filepath[:-3]) + + def get_flops(self) -> float: + """Calculate FLOPS [GFLOPs] for a tf.keras.Model or tf.keras.Sequential model in inference mode. + + It uses tf.compat.v1.profiler under the hood. + """ + if not hasattr(self, "model"): + raise wandb.Error("self.model must be set before using this method.") + + if not isinstance( + self.model, (tf.keras.models.Sequential, tf.keras.models.Model) + ): + raise TypeError( + "Calculating FLOPS is only supported for " + "`tf.keras.Model` and `tf.keras.Sequential` instances." + ) + + from tensorflow.python.framework.convert_to_constants import ( + convert_variables_to_constants_v2_as_graph, + ) + + # Compute FLOPs for one sample + batch_size = 1 + inputs = [ + tf.TensorSpec([batch_size] + inp.shape[1:], inp.dtype) + for inp in self.model.inputs + ] + + # convert tf.keras model into frozen graph to count FLOPs about operations used at inference + real_model = tf.function(self.model).get_concrete_function(inputs) + frozen_func, _ = convert_variables_to_constants_v2_as_graph(real_model) + + # Calculate FLOPs with tf.profiler + run_meta = tf.compat.v1.RunMetadata() + opts = ( + tf.compat.v1.profiler.ProfileOptionBuilder( + tf.compat.v1.profiler.ProfileOptionBuilder().float_operation() + ) + .with_empty_output() + .build() + ) + + flops = tf.compat.v1.profiler.profile( + graph=frozen_func.graph, run_meta=run_meta, cmd="scope", options=opts + ) + + # convert to GFLOPs + return (flops.total_float_ops / 1e9) / 2 diff --git a/lib/python3.12/site-packages/wandb/integration/kfp/__init__.py b/lib/python3.12/site-packages/wandb/integration/kfp/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1f3a362ca0fc895bce03dccdfeaa3274a9b5963d --- /dev/null +++ b/lib/python3.12/site-packages/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() diff --git a/lib/python3.12/site-packages/wandb/integration/kfp/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/kfp/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4c4649406c811d838ae39df73c0cb904decfbef3 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/kfp/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/kfp/__pycache__/helpers.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/kfp/__pycache__/helpers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e4a551ac88dfb6777c81d16c4ac22026a026b9c4 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/kfp/__pycache__/helpers.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/kfp/__pycache__/kfp_patch.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/kfp/__pycache__/kfp_patch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4878a02c52c7e959d88cbf50f10d49021dbb128e Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/kfp/__pycache__/kfp_patch.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/kfp/__pycache__/wandb_logging.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/kfp/__pycache__/wandb_logging.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5cb7ebdb18049c1f14c5675de288f781d9be0fac Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/kfp/__pycache__/wandb_logging.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/kfp/helpers.py b/lib/python3.12/site-packages/wandb/integration/kfp/helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..feebcc62ea3688a2a93818adb445f6454a6d9dc2 --- /dev/null +++ b/lib/python3.12/site-packages/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_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) diff --git a/lib/python3.12/site-packages/wandb/integration/kfp/kfp_patch.py b/lib/python3.12/site-packages/wandb/integration/kfp/kfp_patch.py new file mode 100644 index 0000000000000000000000000000000000000000..367b03afd0fd673250328fe3882fb8c2fef0550c --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/kfp/kfp_patch.py @@ -0,0 +1,335 @@ +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 packaging.version import parse + + MIN_KFP_VERSION = "1.6.1" + + if parse(kfp_version) < parse(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 diff --git a/lib/python3.12/site-packages/wandb/integration/kfp/wandb_logging.py b/lib/python3.12/site-packages/wandb/integration/kfp/wandb_logging.py new file mode 100644 index 0000000000000000000000000000000000000000..5d0edf3eac211b59908bc3aa4849a459c2da2148 --- /dev/null +++ b/lib/python3.12/site-packages/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'' + + 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) diff --git a/lib/python3.12/site-packages/wandb/integration/langchain/__init__.py b/lib/python3.12/site-packages/wandb/integration/langchain/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..aaec971a312e9369e1740f0b31e31782bb9b9e3f --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/langchain/__init__.py @@ -0,0 +1,3 @@ +__all__ = ("WandbTracer",) + +from .wandb_tracer import WandbTracer diff --git a/lib/python3.12/site-packages/wandb/integration/langchain/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/langchain/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ea8610ce498c2ee082374311857f39c7bf458de3 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/langchain/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/langchain/__pycache__/wandb_tracer.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/langchain/__pycache__/wandb_tracer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..04c35caee22c5161470e9ad4d40bb292bdfd6f5c Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/langchain/__pycache__/wandb_tracer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/langchain/wandb_tracer.py b/lib/python3.12/site-packages/wandb/integration/langchain/wandb_tracer.py new file mode 100644 index 0000000000000000000000000000000000000000..6c9a5ff31ce4d1effd6f2d28101b328257f0c577 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/langchain/wandb_tracer.py @@ -0,0 +1,49 @@ +"""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.proto.wandb_deprecated import Deprecated +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 + + +class WandbTracer(WandbTracer): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + deprecate.deprecate( + field_name=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`.", + ) diff --git a/lib/python3.12/site-packages/wandb/integration/lightgbm/__init__.py b/lib/python3.12/site-packages/wandb/integration/lightgbm/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..bc1029563baa721fbdbc2c2809974b59203008ee --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/lightgbm/__init__.py @@ -0,0 +1,239 @@ +"""W&B callback for lightgbm. + +Really simple callback to get logging for each tree + +Example usage: + +param_list = [("eta", 0.08), ("max_depth", 6), ("subsample", 0.8), ("colsample_bytree", 0.8), ("alpha", 8), ("num_class", 10)] +config.update(dict(param_list)) +lgb = lgb.train(param_list, d_train, callbacks=[wandb_callback()]) +""" + +from pathlib import Path +from typing import TYPE_CHECKING, Callable + +import lightgbm # type: ignore +from lightgbm import Booster + +import wandb +from wandb.sdk.lib import telemetry as wb_telemetry + +MINIMIZE_METRICS = [ + "l1", + "l2", + "rmse", + "mape", + "huber", + "fair", + "poisson", + "gamma", + "binary_logloss", +] + +MAXIMIZE_METRICS = ["map", "auc", "average_precision"] + + +if TYPE_CHECKING: + from typing import Any, Dict, List, NamedTuple, Tuple, Union + + # Note: upstream lightgbm has this defined incorrectly + _EvalResultTuple = Union[ + Tuple[str, str, float, bool], Tuple[str, str, float, bool, float] + ] + + class CallbackEnv(NamedTuple): + model: Any + params: Dict + iteration: int + begin_interation: int + end_iteration: int + evaluation_result_list: List[_EvalResultTuple] + + +def _define_metric(data: str, metric_name: str) -> None: + """Capture model performance at the best step. + + instead of the last step, of training in your `wandb.summary` + """ + if "loss" in str.lower(metric_name): + wandb.define_metric(f"{data}_{metric_name}", summary="min") + elif str.lower(metric_name) in MINIMIZE_METRICS: + wandb.define_metric(f"{data}_{metric_name}", summary="min") + elif str.lower(metric_name) in MAXIMIZE_METRICS: + wandb.define_metric(f"{data}_{metric_name}", summary="max") + + +def _checkpoint_artifact( + model: "Booster", iteration: int, aliases: "List[str]" +) -> None: + """Upload model checkpoint as W&B artifact.""" + # NOTE: type ignore required because wandb.run is improperly inferred as None type + model_name = f"model_{wandb.run.id}" # type: ignore + model_path = Path(wandb.run.dir) / f"model_ckpt_{iteration}.txt" # type: ignore + + model.save_model(model_path, num_iteration=iteration) + + model_artifact = wandb.Artifact(name=model_name, type="model") + model_artifact.add_file(str(model_path)) + wandb.log_artifact(model_artifact, aliases=aliases) + + +def _log_feature_importance(model: "Booster") -> None: + """Log feature importance.""" + feat_imps = model.feature_importance() + feats = model.feature_name() + fi_data = [[feat, feat_imp] for feat, feat_imp in zip(feats, feat_imps)] + table = wandb.Table(data=fi_data, columns=["Feature", "Importance"]) + wandb.log( + { + "Feature Importance": wandb.plot.bar( + table, "Feature", "Importance", title="Feature Importance" + ) + }, + commit=False, + ) + + +class _WandbCallback: + """Internal class to handle `wandb_callback` logic. + + This callback is adapted form the LightGBM's `_RecordEvaluationCallback`. + """ + + def __init__(self, log_params: bool = True, define_metric: bool = True) -> None: + self.order = 20 + self.before_iteration = False + self.log_params = log_params + self.define_metric_bool = define_metric + + def _init(self, env: "CallbackEnv") -> None: + with wb_telemetry.context() as tel: + tel.feature.lightgbm_wandb_callback = True + + # log the params as W&B config. + if self.log_params: + wandb.config.update(env.params) + + # use `define_metric` to set the wandb summary to the best metric value. + for item in env.evaluation_result_list: + if self.define_metric_bool: + if len(item) == 4: + data_name, eval_name = item[:2] + _define_metric(data_name, eval_name) + else: + data_name, eval_name = item[1].split() + _define_metric(data_name, f"{eval_name}-mean") + _define_metric(data_name, f"{eval_name}-stdv") + + def __call__(self, env: "CallbackEnv") -> None: + if env.iteration == env.begin_iteration: # type: ignore + self._init(env) + + for item in env.evaluation_result_list: + if len(item) == 4: + data_name, eval_name, result = item[:3] + wandb.log( + {data_name + "_" + eval_name: result}, + commit=False, + ) + else: + data_name, eval_name = item[1].split() + res_mean = item[2] + res_stdv = item[4] + wandb.log( + { + data_name + "_" + eval_name + "-mean": res_mean, + data_name + "_" + eval_name + "-stdv": res_stdv, + }, + commit=False, + ) + + # call `commit=True` to log the data as a single W&B step. + wandb.log({"iteration": env.iteration}, commit=True) + + +def wandb_callback(log_params: bool = True, define_metric: bool = True) -> Callable: + """Automatically integrates LightGBM with wandb. + + Args: + log_params: (boolean) if True (default) logs params passed to lightgbm.train as W&B config + define_metric: (boolean) if True (default) capture model performance at the best step, instead of the last step, of training in your `wandb.summary` + + Passing `wandb_callback` to LightGBM will: + - log params passed to lightgbm.train as W&B config (default). + - log evaluation metrics collected by LightGBM, such as rmse, accuracy etc to Weights & Biases + - Capture the best metric in `wandb.summary` when `define_metric=True` (default). + + Use `log_summary` as an extension of this callback. + + Example: + ```python + params = { + "boosting_type": "gbdt", + "objective": "regression", + } + gbm = lgb.train( + params, + lgb_train, + num_boost_round=10, + valid_sets=lgb_eval, + valid_names=("validation"), + callbacks=[wandb_callback()], + ) + ``` + """ + return _WandbCallback(log_params, define_metric) + + +def log_summary( + model: Booster, feature_importance: bool = True, save_model_checkpoint: bool = False +) -> None: + """Log useful metrics about lightgbm model after training is done. + + Args: + model: (Booster) is an instance of lightgbm.basic.Booster. + feature_importance: (boolean) if True (default), logs the feature importance plot. + save_model_checkpoint: (boolean) if True saves the best model and upload as W&B artifacts. + + Using this along with `wandb_callback` will: + + - log `best_iteration` and `best_score` as `wandb.summary`. + - log feature importance plot. + - save and upload your best trained model to Weights & Biases Artifacts (when `save_model_checkpoint = True`) + + Example: + ```python + params = { + "boosting_type": "gbdt", + "objective": "regression", + } + gbm = lgb.train( + params, + lgb_train, + num_boost_round=10, + valid_sets=lgb_eval, + valid_names=("validation"), + callbacks=[wandb_callback()], + ) + + log_summary(gbm) + ``` + """ + if wandb.run is None: + raise wandb.Error("You must call wandb.init() before WandbCallback()") + + if not isinstance(model, Booster): + raise wandb.Error("Model should be an instance of lightgbm.basic.Booster") + + wandb.run.summary["best_iteration"] = model.best_iteration + wandb.run.summary["best_score"] = model.best_score + + # Log feature importance + if feature_importance: + _log_feature_importance(model) + + if save_model_checkpoint: + _checkpoint_artifact(model, model.best_iteration, aliases=["best"]) + + with wb_telemetry.context() as tel: + tel.feature.lightgbm_log_summary = True diff --git a/lib/python3.12/site-packages/wandb/integration/lightgbm/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/lightgbm/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4f78de015ff13f2ebdd1bd4af214e826966af181 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/lightgbm/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/lightning/__init__.py b/lib/python3.12/site-packages/wandb/integration/lightning/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/wandb/integration/lightning/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/lightning/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..14d89c2a1a3cab6b8f104ba18bf9a12e97be2a4e Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/lightning/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/lightning/fabric/__init__.py b/lib/python3.12/site-packages/wandb/integration/lightning/fabric/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..eff7554e1131034c637c8cfef591756211b8a898 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/lightning/fabric/__init__.py @@ -0,0 +1,3 @@ +from wandb.integration.lightning.fabric.logger import WandbLogger + +__all__ = ("WandbLogger",) diff --git a/lib/python3.12/site-packages/wandb/integration/lightning/fabric/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/lightning/fabric/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..25f4fa24916b6fe6ac9e18ec9d9dbbd4629414ff Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/lightning/fabric/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/lightning/fabric/__pycache__/logger.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/lightning/fabric/__pycache__/logger.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f2873f7d30a396279dd84e22f0fb0cf48678f709 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/lightning/fabric/__pycache__/logger.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/lightning/fabric/logger.py b/lib/python3.12/site-packages/wandb/integration/lightning/fabric/logger.py new file mode 100644 index 0000000000000000000000000000000000000000..bc698bba511aac65a44837e0a893deb98838ae88 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/lightning/fabric/logger.py @@ -0,0 +1,764 @@ +import os +from argparse import Namespace +from pathlib import Path +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Mapping, Optional, Union + +from packaging import version +from typing_extensions import override + +import wandb +from wandb import Artifact +from wandb.sdk.lib import telemetry +from wandb.sdk.wandb_run import Run + +try: + import lightning + import torch.nn as nn + from lightning.fabric.loggers.logger import Logger, rank_zero_experiment + from lightning.fabric.utilities.exceptions import MisconfigurationException + from lightning.fabric.utilities.logger import ( + _add_prefix, + _convert_params, + _sanitize_callable_params, + ) + from lightning.fabric.utilities.rank_zero import rank_zero_only, rank_zero_warn + from lightning.fabric.utilities.types import _PATH + from torch import Tensor + from torch.nn import Module + + if version.parse(lightning.__version__) > version.parse("2.1.3"): + wandb.termwarn( + """This integration is tested and supported for lightning Fabric 2.1.3. + Please report any issues to https://github.com/wandb/wandb/issues with the tag `lightning-fabric`.""", + repeat=False, + ) + + if TYPE_CHECKING: + from lightning.pytorch.callbacks.model_checkpoint import ModelCheckpoint + +except ImportError as e: + wandb.Error(e) + + +class WandbLogger(Logger): + r"""Log using `Weights and Biases `_. + + **Installation and set-up** + + Install with pip: + + .. code-block:: bash + + pip install wandb + + Create a `WandbLogger` instance: + + .. code-block:: python + + from lightning.fabric.loggers import WandbLogger + + wandb_logger = WandbLogger(project="MNIST") + + Pass the logger instance to the `Trainer`: + + .. code-block:: python + + trainer = Trainer(logger=wandb_logger) + + A new W&B run will be created when training starts if you have not created one manually before with `wandb.init()`. + + **Log metrics** + + Log from :class:`~lightning.pytorch.core.LightningModule`: + + .. code-block:: python + + class LitModule(LightningModule): + def training_step(self, batch, batch_idx): + self.log("train/loss", loss) + + Use directly wandb module: + + .. code-block:: python + + wandb.log({"train/loss": loss}) + + **Log hyper-parameters** + + Save :class:`~lightning.pytorch.core.LightningModule` parameters: + + .. code-block:: python + + class LitModule(LightningModule): + def __init__(self, *args, **kwarg): + self.save_hyperparameters() + + Add other config parameters: + + .. code-block:: python + + # add one parameter + wandb_logger.experiment.config["key"] = value + + # add multiple parameters + wandb_logger.experiment.config.update({key1: val1, key2: val2}) + + # use directly wandb module + wandb.config["key"] = value + wandb.config.update() + + **Log gradients, parameters and model topology** + + Call the `watch` method for automatically tracking gradients: + + .. code-block:: python + + # log gradients and model topology + wandb_logger.watch(model) + + # log gradients, parameter histogram and model topology + wandb_logger.watch(model, log="all") + + # change log frequency of gradients and parameters (100 steps by default) + wandb_logger.watch(model, log_freq=500) + + # do not log graph (in case of errors) + wandb_logger.watch(model, log_graph=False) + + The `watch` method adds hooks to the model which can be removed at the end of training: + + .. code-block:: python + + wandb_logger.experiment.unwatch(model) + + **Log model checkpoints** + + Log model checkpoints at the end of training: + + .. code-block:: python + + wandb_logger = WandbLogger(log_model=True) + + Log model checkpoints as they get created during training: + + .. code-block:: python + + wandb_logger = WandbLogger(log_model="all") + + Custom checkpointing can be set up through :class:`~lightning.pytorch.callbacks.ModelCheckpoint`: + + .. code-block:: python + + # log model only if `val_accuracy` increases + wandb_logger = WandbLogger(log_model="all") + checkpoint_callback = ModelCheckpoint(monitor="val_accuracy", mode="max") + trainer = Trainer(logger=wandb_logger, callbacks=[checkpoint_callback]) + + `latest` and `best` aliases are automatically set to easily retrieve a model checkpoint: + + .. code-block:: python + + # reference can be retrieved in artifacts panel + # "VERSION" can be a version (ex: "v2") or an alias ("latest or "best") + checkpoint_reference = "USER/PROJECT/MODEL-RUN_ID:VERSION" + + # download checkpoint locally (if not already cached) + run = wandb.init(project="MNIST") + artifact = run.use_artifact(checkpoint_reference, type="model") + artifact_dir = artifact.download() + + # load checkpoint + model = LitModule.load_from_checkpoint(Path(artifact_dir) / "model.ckpt") + + **Log media** + + Log text with: + + .. code-block:: python + + # using columns and data + columns = ["input", "label", "prediction"] + data = [["cheese", "english", "english"], ["fromage", "french", "spanish"]] + wandb_logger.log_text(key="samples", columns=columns, data=data) + + # using a pandas DataFrame + wandb_logger.log_text(key="samples", dataframe=my_dataframe) + + Log images with: + + .. code-block:: python + + # using tensors, numpy arrays or PIL images + wandb_logger.log_image(key="samples", images=[img1, img2]) + + # adding captions + wandb_logger.log_image( + key="samples", images=[img1, img2], caption=["tree", "person"] + ) + + # using file path + wandb_logger.log_image(key="samples", images=["img_1.jpg", "img_2.jpg"]) + + More arguments can be passed for logging segmentation masks and bounding boxes. Refer to + `Image Overlays documentation `_. + + **Log Tables** + + `W&B Tables `_ can be used to log, + query and analyze tabular data. + + They support any type of media (text, image, video, audio, molecule, html, etc) and are great for storing, + understanding and sharing any form of data, from datasets to model predictions. + + .. code-block:: python + + columns = ["caption", "image", "sound"] + data = [ + ["cheese", wandb.Image(img_1), wandb.Audio(snd_1)], + ["wine", wandb.Image(img_2), wandb.Audio(snd_2)], + ] + wandb_logger.log_table(key="samples", columns=columns, data=data) + + + **Downloading and Using Artifacts** + + To download an artifact without starting a run, call the ``download_artifact`` + function on the class: + + .. code-block:: python + + artifact_dir = wandb_logger.download_artifact(artifact="path/to/artifact") + + To download an artifact and link it to an ongoing run call the ``download_artifact`` + function on the logger instance: + + .. code-block:: python + + class MyModule(LightningModule): + def any_lightning_module_function_or_hook(self): + self.logger.download_artifact(artifact="path/to/artifact") + + To link an artifact from a previous run you can use ``use_artifact`` function: + + .. code-block:: python + + wandb_logger.use_artifact(artifact="path/to/artifact") + + See Also: + - `Demo in Google Colab `__ with hyperparameter search and model logging + - `W&B Documentation `__ + + Args: + name: Display name for the run. + save_dir: Path where data is saved. + version: Sets the version, mainly used to resume a previous run. + offline: Run offline (data can be streamed later to wandb servers). + dir: Same as save_dir. + id: Same as version. + anonymous: Enables or explicitly disables anonymous logging. + project: The name of the project to which this run will belong. If not set, the environment variable + `WANDB_PROJECT` will be used as a fallback. If both are not set, it defaults to ``'lightning_logs'``. + log_model: Log checkpoints created by :class:`~lightning.pytorch.callbacks.ModelCheckpoint` + as W&B artifacts. `latest` and `best` aliases are automatically set. + + * if ``log_model == 'all'``, checkpoints are logged during training. + * if ``log_model == True``, checkpoints are logged at the end of training, except when + `~lightning.pytorch.callbacks.ModelCheckpoint.save_top_k` ``== -1`` + which also logs every checkpoint during training. + * if ``log_model == False`` (default), no checkpoint is logged. + + prefix: A string to put at the beginning of metric keys. + experiment: WandB experiment object. Automatically set when creating a run. + checkpoint_name: Name of the model checkpoint artifact being logged. + log_checkpoint_on: When to log model checkpoints as W&B artifacts. Only used if ``log_model`` is ``True``. + Options: ``"success"``, ``"all"``. Default: ``"success"``. + \**kwargs: Arguments passed to :func:`wandb.init` like `entity`, `group`, `tags`, etc. + + Raises: + ModuleNotFoundError: + If required WandB package is not installed on the device. + MisconfigurationException: + If both ``log_model`` and ``offline`` is set to ``True``. + + """ + + LOGGER_JOIN_CHAR = "-" + + def __init__( + self, + name: Optional[str] = None, + save_dir: _PATH = ".", + version: Optional[str] = None, + offline: bool = False, + dir: Optional[_PATH] = None, + id: Optional[str] = None, + anonymous: Optional[bool] = None, + project: Optional[str] = None, + log_model: Union[Literal["all"], bool] = False, + experiment: Optional["Run"] = None, + prefix: str = "", + checkpoint_name: Optional[str] = None, + log_checkpoint_on: Union[Literal["success"], Literal["all"]] = "success", + **kwargs: Any, + ) -> None: + if offline and log_model: + raise MisconfigurationException( + f"Providing log_model={log_model} and offline={offline} is an invalid configuration" + " since model checkpoints cannot be uploaded in offline mode.\n" + "Hint: Set `offline=False` to log your model." + ) + + super().__init__() + self._offline = offline + self._log_model = log_model + self._prefix = prefix + self._experiment = experiment + self._logged_model_time: Dict[str, float] = {} + self._checkpoint_callback: Optional[ModelCheckpoint] = None + + # paths are processed as strings + if save_dir is not None: + save_dir = os.fspath(save_dir) + elif dir is not None: + dir = os.fspath(dir) + + project = project or os.environ.get("WANDB_PROJECT", "lightning_fabric_logs") + + # set wandb init arguments + self._wandb_init: Dict[str, Any] = { + "name": name, + "project": project, + "dir": save_dir or dir, + "id": version or id, + "resume": "allow", + "anonymous": ("allow" if anonymous else None), + } + self._wandb_init.update(**kwargs) + # extract parameters + self._project = self._wandb_init.get("project") + self._save_dir = self._wandb_init.get("dir") + self._name = self._wandb_init.get("name") + self._id = self._wandb_init.get("id") + self._checkpoint_name = checkpoint_name + self._log_checkpoint_on = log_checkpoint_on + + def __getstate__(self) -> Dict[str, Any]: + # Hack: If the 'spawn' launch method is used, the logger will get pickled and this `__getstate__` gets called. + # We create an experiment here in the main process, and attach to it in the worker process. + # Using wandb-service, we persist the same experiment even if multiple `Trainer.fit/test/validate` calls + # are made. + _ = self.experiment + + state = self.__dict__.copy() + # args needed to reload correct experiment + if self._experiment is not None: + state["_id"] = getattr(self._experiment, "id", None) + state["_attach_id"] = getattr(self._experiment, "_attach_id", None) + state["_name"] = self._experiment.name + + # cannot be pickled + state["_experiment"] = None + return state + + @property + @rank_zero_experiment + def experiment(self) -> "Run": + r"""Actual wandb object. + + To use wandb features in your :class:`~lightning.pytorch.core.LightningModule`, do the + following. + + Example:: + + .. code-block:: python + + self.logger.experiment.some_wandb_function() + + """ + if self._experiment is None: + if self._offline: + os.environ["WANDB_MODE"] = "dryrun" + + attach_id = getattr(self, "_attach_id", None) + if wandb.run is not None: + # wandb process already created in this instance + rank_zero_warn( + "There is a wandb run already in progress and newly created instances of `WandbLogger` will reuse" + " this run. If this is not desired, call `wandb.finish()` before instantiating `WandbLogger`." + ) + self._experiment = wandb.run + elif attach_id is not None and hasattr(wandb, "_attach"): + # attach to wandb process referenced + self._experiment = wandb._attach(attach_id) + else: + # create new wandb process + self._experiment = wandb.init(**self._wandb_init) + + # define default x-axis + if isinstance(self._experiment, Run) and getattr( + self._experiment, "define_metric", None + ): + self._experiment.define_metric("trainer/global_step") + self._experiment.define_metric( + "*", step_metric="trainer/global_step", step_sync=True + ) + + self._experiment._label(repo="lightning_fabric_logger") # pylint: disable=protected-access + with telemetry.context(run=self._experiment) as tel: + tel.feature.lightning_fabric_logger = True + return self._experiment + + def watch( + self, + model: nn.Module, + log: str = "gradients", + log_freq: int = 100, + log_graph: bool = True, + ) -> None: + self.experiment.watch(model, log=log, log_freq=log_freq, log_graph=log_graph) + + @override + @rank_zero_only + def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None: # type: ignore[override] + params = _convert_params(params) + params = _sanitize_callable_params(params) + self.experiment.config.update(params, allow_val_change=True) + + @override + @rank_zero_only + def log_metrics( + self, metrics: Mapping[str, float], step: Optional[int] = None + ) -> None: + assert rank_zero_only.rank == 0, "experiment tried to log from global_rank != 0" + + metrics = _add_prefix(metrics, self._prefix, self.LOGGER_JOIN_CHAR) + if step is not None: + self.experiment.log(dict(metrics, **{"trainer/global_step": step})) + else: + self.experiment.log(metrics) + + @rank_zero_only + def log_table( + self, + key: str, + columns: Optional[List[str]] = None, + data: Optional[List[List[Any]]] = None, + dataframe: Any = None, + step: Optional[int] = None, + ) -> None: + """Log a Table containing any object type (text, image, audio, video, molecule, html, etc). + + Can be defined either with `columns` and `data` or with `dataframe`. + + """ + metrics = {key: wandb.Table(columns=columns, data=data, dataframe=dataframe)} + self.log_metrics(metrics, step) + + @rank_zero_only + def log_text( + self, + key: str, + columns: Optional[List[str]] = None, + data: Optional[List[List[str]]] = None, + dataframe: Any = None, + step: Optional[int] = None, + ) -> None: + """Log text as a Table. + + Can be defined either with `columns` and `data` or with `dataframe`. + + """ + self.log_table(key, columns, data, dataframe, step) + + @rank_zero_only + def log_html( + self, key: str, htmls: List[Any], step: Optional[int] = None, **kwargs: Any + ) -> None: + """Log html files. + + Optional kwargs are lists passed to each html (ex: inject). + + """ + if not isinstance(htmls, list): + raise TypeError(f'Expected a list as "htmls", found {type(htmls)}') + n = len(htmls) + for k, v in kwargs.items(): + if len(v) != n: + raise ValueError(f"Expected {n} items but only found {len(v)} for {k}") + kwarg_list = [{k: kwargs[k][i] for k in kwargs} for i in range(n)] + + metrics = { + key: [wandb.Html(html, **kwarg) for html, kwarg in zip(htmls, kwarg_list)] + } + self.log_metrics(metrics, step) # type: ignore[arg-type] + + @rank_zero_only + def log_image( + self, key: str, images: List[Any], step: Optional[int] = None, **kwargs: Any + ) -> None: + """Log images (tensors, numpy arrays, PIL Images or file paths). + + Optional kwargs are lists passed to each image (ex: caption, masks, boxes). + + """ + if not isinstance(images, list): + raise TypeError(f'Expected a list as "images", found {type(images)}') + n = len(images) + for k, v in kwargs.items(): + if len(v) != n: + raise ValueError(f"Expected {n} items but only found {len(v)} for {k}") + kwarg_list = [{k: kwargs[k][i] for k in kwargs} for i in range(n)] + + metrics = { + key: [wandb.Image(img, **kwarg) for img, kwarg in zip(images, kwarg_list)] + } + self.log_metrics(metrics, step) # type: ignore[arg-type] + + @rank_zero_only + def log_audio( + self, key: str, audios: List[Any], step: Optional[int] = None, **kwargs: Any + ) -> None: + r"""Log audios (numpy arrays, or file paths). + + Args: + key: The key to be used for logging the audio files + audios: The list of audio file paths, or numpy arrays to be logged + step: The step number to be used for logging the audio files + \**kwargs: Optional kwargs are lists passed to each ``Wandb.Audio`` instance (ex: caption, sample_rate). + + Optional kwargs are lists passed to each audio (ex: caption, sample_rate). + + """ + if not isinstance(audios, list): + raise TypeError(f'Expected a list as "audios", found {type(audios)}') + n = len(audios) + for k, v in kwargs.items(): + if len(v) != n: + raise ValueError(f"Expected {n} items but only found {len(v)} for {k}") + kwarg_list = [{k: kwargs[k][i] for k in kwargs} for i in range(n)] + + metrics = { + key: [ + wandb.Audio(audio, **kwarg) for audio, kwarg in zip(audios, kwarg_list) + ] + } + self.log_metrics(metrics, step) # type: ignore[arg-type] + + @rank_zero_only + def log_video( + self, key: str, videos: List[Any], step: Optional[int] = None, **kwargs: Any + ) -> None: + """Log videos (numpy arrays, or file paths). + + Args: + key: The key to be used for logging the video files + videos: The list of video file paths, or numpy arrays to be logged + step: The step number to be used for logging the video files + **kwargs: Optional kwargs are lists passed to each Wandb.Video instance (ex: caption, fps, format). + + Optional kwargs are lists passed to each video (ex: caption, fps, format). + + """ + if not isinstance(videos, list): + raise TypeError(f'Expected a list as "videos", found {type(videos)}') + n = len(videos) + for k, v in kwargs.items(): + if len(v) != n: + raise ValueError(f"Expected {n} items but only found {len(v)} for {k}") + kwarg_list = [{k: kwargs[k][i] for k in kwargs} for i in range(n)] + + metrics = { + key: [ + wandb.Video(video, **kwarg) for video, kwarg in zip(videos, kwarg_list) + ] + } + self.log_metrics(metrics, step) # type: ignore[arg-type] + + @property + @override + def save_dir(self) -> Optional[str]: + """Gets the save directory. + + Returns: + The path to the save directory. + + """ + return self._save_dir + + @property + @override + def name(self) -> Optional[str]: + """The project name of this experiment. + + Returns: + The name of the project the current experiment belongs to. This name is not the same as `wandb.Run`'s + name. To access wandb's internal experiment name, use ``logger.experiment.name`` instead. + + """ + return self._project + + @property + @override + def version(self) -> Optional[str]: + """Gets the id of the experiment. + + Returns: + The id of the experiment if the experiment exists else the id given to the constructor. + + """ + # don't create an experiment if we don't have one + return self._experiment.id if self._experiment else self._id + + @property + def log_dir(self) -> Optional[str]: + """Gets the save directory. + + Returns: + The path to the save directory. + + """ + return self.save_dir + + @property + def group_separator(self) -> str: + """Return the default separator used by the logger to group the data into subfolders.""" + return self.LOGGER_JOIN_CHAR + + @property + def root_dir(self) -> Optional[str]: + """Return the root directory. + + Return the root directory where all versions of an experiment get saved, or `None` if the logger does not + save data locally. + """ + return self.save_dir.parent if self.save_dir else None + + def log_graph(self, model: Module, input_array: Optional[Tensor] = None) -> None: + """Record model graph. + + Args: + model: the model with an implementation of ``forward``. + input_array: input passes to `model.forward` + + This is a noop function and does not perform any operation. + """ + return + + @override + def after_save_checkpoint(self, checkpoint_callback: "ModelCheckpoint") -> None: + # log checkpoints as artifacts + if ( + self._log_model == "all" + or self._log_model is True + and checkpoint_callback.save_top_k == -1 + ): + # TODO: Replace with new Fabric Checkpoints system + self._scan_and_log_pytorch_checkpoints(checkpoint_callback) + elif self._log_model is True: + self._checkpoint_callback = checkpoint_callback + + @staticmethod + @rank_zero_only + def download_artifact( + artifact: str, + save_dir: Optional[_PATH] = None, + artifact_type: Optional[str] = None, + use_artifact: Optional[bool] = True, + ) -> str: + """Downloads an artifact from the wandb server. + + Args: + artifact: The path of the artifact to download. + save_dir: The directory to save the artifact to. + artifact_type: The type of artifact to download. + use_artifact: Whether to add an edge between the artifact graph. + + Returns: + The path to the downloaded artifact. + + """ + if wandb.run is not None and use_artifact: + artifact = wandb.run.use_artifact(artifact) + else: + api = wandb.Api() + artifact = api.artifact(artifact, type=artifact_type) + + save_dir = None if save_dir is None else os.fspath(save_dir) + return artifact.download(root=save_dir) + + def use_artifact( + self, artifact: str, artifact_type: Optional[str] = None + ) -> "Artifact": + """Logs to the wandb dashboard that the mentioned artifact is used by the run. + + Args: + artifact: The path of the artifact. + artifact_type: The type of artifact being used. + + Returns: + wandb Artifact object for the artifact. + + """ + return self.experiment.use_artifact(artifact, type=artifact_type) + + @override + @rank_zero_only + def save(self) -> None: + """Save log data.""" + self.experiment.log({}, commit=True) + + @override + @rank_zero_only + def finalize(self, status: str) -> None: + if self._log_checkpoint_on == "success" and status != "success": + # Currently, checkpoints only get logged on success + return + # log checkpoints as artifacts + if ( + self._checkpoint_callback + and self._experiment is not None + and self._log_checkpoint_on in ["success", "all"] + ): + self._scan_and_log_pytorch_checkpoints(self._checkpoint_callback) + + def _scan_and_log_pytorch_checkpoints( + self, checkpoint_callback: "ModelCheckpoint" + ) -> None: + from lightning.pytorch.loggers.utilities import _scan_checkpoints + + # get checkpoints to be saved with associated score + checkpoints = _scan_checkpoints(checkpoint_callback, self._logged_model_time) + + # log iteratively all new checkpoints + for t, p, s, _ in checkpoints: + metadata = { + "score": s.item() if isinstance(s, Tensor) else s, + "original_filename": Path(p).name, + checkpoint_callback.__class__.__name__: { + k: getattr(checkpoint_callback, k) + for k in [ + "monitor", + "mode", + "save_last", + "save_top_k", + "save_weights_only", + "_every_n_train_steps", + ] + # ensure it does not break if `ModelCheckpoint` args change + if hasattr(checkpoint_callback, k) + }, + } + if not self._checkpoint_name: + self._checkpoint_name = f"model-{self.experiment.id}" + artifact = wandb.Artifact( + name=self._checkpoint_name, type="model", metadata=metadata + ) + artifact.add_file(p, name="model.ckpt") + aliases = ( + ["latest", "best"] + if p == checkpoint_callback.best_model_path + else ["latest"] + ) + self.experiment.log_model(artifact, aliases=aliases) + # remember logged models - timestamp needed in case filename didn't change (lastkckpt or custom name) + self._logged_model_time[p] = t diff --git a/lib/python3.12/site-packages/wandb/integration/metaflow/__init__.py b/lib/python3.12/site-packages/wandb/integration/metaflow/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..748a71eaebae9f844d881b0842251b102f6b127d --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/metaflow/__init__.py @@ -0,0 +1,3 @@ +from .metaflow import wandb_log, wandb_track, wandb_use + +__all__ = ["wandb_log", "wandb_track", "wandb_use"] diff --git a/lib/python3.12/site-packages/wandb/integration/metaflow/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/metaflow/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d34203aec0cd7ad783f5d7e2188b5e2a1867df7a Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/metaflow/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/metaflow/__pycache__/metaflow.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/metaflow/__pycache__/metaflow.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..880db885afe8f648b73fdd4df2fab8d54b562bf7 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/metaflow/__pycache__/metaflow.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/metaflow/metaflow.py b/lib/python3.12/site-packages/wandb/integration/metaflow/metaflow.py new file mode 100644 index 0000000000000000000000000000000000000000..9e801744d29434264577f8417a2d12dd5356455f --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/metaflow/metaflow.py @@ -0,0 +1,383 @@ +"""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 + +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 + +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 + + +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!" + ) + +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!" + ) + +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!" + ) + + +class ArtifactProxy: + def __init__(self, flow): + # do this to avoid recursion problem with __setattr__ + self.__dict__.update( + { + "flow": flow, + "inputs": {}, + "outputs": {}, + "base": set(dir(flow)), + "params": {p: getattr(flow, p) for p in current.parameter_names}, + } + ) + + def __setattr__(self, key, val): + self.outputs[key] = val + return setattr(self.flow, key, val) + + def __getattr__(self, key): + if key not in self.base and key not in self.outputs: + self.inputs[key] = getattr(self.flow, key) + return getattr(self.flow, key) + + +@dispatch +def wandb_track( + name: str, + data: Union[dict, list, set, str, int, float, bool], + run=None, + testing=False, + *args, + **kwargs, +): # type: ignore + if testing: + return "scalar" + + run.log({name: data}) + + +@dispatch +def wandb_track( + name: str, data: Path, datasets=False, run=None, testing=False, *args, **kwargs +): + 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)})") + + +# this is the base case +@dispatch +def wandb_track( + name: str, data, others=False, run=None, testing=False, *args, **kwargs +): + if testing: + return "generic" if others else None + + 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)})") + + +@dispatch +def wandb_use(name: str, data, *args, **kwargs): + try: + return _wandb_use(name, data, *args, **kwargs) + 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!" + ) + + +@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 + 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_use(name: str, data, others=False, run=None, testing=False, *args, **kwargs): # type: ignore + if testing: + return "others" if others else None + + if others: + run.use_artifact(f"{name}:latest") + wandb.termlog(f"Using artifact: {name} ({type(data)})") + + +def coalesce(*arg): + return next((a for a in arg if a is not None), None) + + +def wandb_log( + func=None, + # /, # py38 only + datasets=False, + models=False, + others=False, + settings=None, +): + """Automatically log parameters and artifacts to W&B by type dispatch. + + 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 + + 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}` + """ + + @wraps(func) + def decorator(func): + # If you decorate a class, apply the decoration to all methods in that class + if inspect.isclass(func): + cls = func + for attr in cls.__dict__: + if callable(getattr(cls, attr)): + if not hasattr(attr, "_base_func"): + setattr(cls, attr, decorator(getattr(cls, attr))) + return cls + + # prefer the earliest decoration (i.e. method decoration overrides class decoration) + if hasattr(func, "_base_func"): + return func + + @wraps(func) + def wrapper(self, *args, settings=settings, **kwargs): + if not isinstance(settings, wandb.sdk.wandb_settings.Settings): + settings = wandb.Settings() + + settings.update_from_dict( + { + "run_group": coalesce( + settings.run_group, f"{current.flow_name}/{current.run_id}" + ), + "run_job_type": coalesce(settings.run_job_type, current.step_name), + } + ) + + with wandb.init(settings=settings) as run: + with wb_telemetry.context(run=run) as tel: + tel.feature.metaflow = True + proxy = ArtifactProxy(self) + run.config.update(proxy.params) + func(proxy, *args, **kwargs) + + for name, data in proxy.inputs.items(): + wandb_use( + name, + data, + datasets=datasets, + models=models, + others=others, + run=run, + ) + + for name, data in proxy.outputs.items(): + wandb_track( + name, + data, + datasets=datasets, + models=models, + others=others, + run=run, + ) + + wrapper._base_func = func + + # Add for testing visibility + wrapper._kwargs = { + "datasets": datasets, + "models": models, + "others": others, + "settings": settings, + } + return wrapper + + if func is None: + return decorator + else: + return decorator(func) diff --git a/lib/python3.12/site-packages/wandb/integration/openai/__init__.py b/lib/python3.12/site-packages/wandb/integration/openai/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f2c216c5f72a7d0310b5cfab98b4c6ff5e0e75e6 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/openai/__init__.py @@ -0,0 +1,3 @@ +__all__ = ("autolog", "WandbLogger") + +from .openai import autolog diff --git a/lib/python3.12/site-packages/wandb/integration/openai/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/openai/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d1199eedfa632688ec8a91e0e248dbdc6732e857 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/openai/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/openai/__pycache__/fine_tuning.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/openai/__pycache__/fine_tuning.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..622c0435594e9816b0a9166589e418677bace75a Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/openai/__pycache__/fine_tuning.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/openai/__pycache__/openai.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/openai/__pycache__/openai.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2d57ffcccb952b82e2b30b59cfb0c8b7ae7f2ad9 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/openai/__pycache__/openai.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/openai/__pycache__/resolver.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/openai/__pycache__/resolver.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..58a7fd59973d45d33dc2c85ad42e8a3ec19fb5af Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/openai/__pycache__/resolver.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/openai/fine_tuning.py b/lib/python3.12/site-packages/wandb/integration/openai/fine_tuning.py new file mode 100644 index 0000000000000000000000000000000000000000..783cdcf77208e2fd342a29ef65b42c54d48a0261 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/openai/fine_tuning.py @@ -0,0 +1,481 @@ +import base64 +import datetime +import io +import json +import os +import re +import tempfile +import time +from typing import Any, Dict, List, Optional, Tuple, Union + +from packaging.version import parse + +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", + required="This integration requires `openai`. To install, please run `pip install openai`", + lazy=False, +) + +if parse(openai.__version__) < parse("1.12.0"): + raise wandb.Error( + f"This integration requires openai version 1.12.0 and above. Your current version is {openai.__version__} " + "To fix, please `pip install -U openai`" + ) + +from openai import OpenAI # noqa: E402 +from openai.types.fine_tuning import FineTuningJob # noqa: E402 +from openai.types.fine_tuning.fine_tuning_job import ( # noqa: E402 + Error, + Hyperparameters, +) + +np = util.get_module( + name="numpy", + required="`numpy` not installed >> This integration requires numpy! To fix, please `pip install numpy`", + lazy=False, +) + +pd = util.get_module( + name="pandas", + required="`pandas` not installed >> This integration requires pandas! To fix, please `pip install pandas`", + lazy=False, +) + + +class WandbLogger: + """Log OpenAI fine-tunes to [Weights & Biases](https://wandb.me/openai-docs).""" + + _wandb_api: Optional[wandb.Api] = None + _logged_in: bool = False + openai_client: Optional[OpenAI] = None + _run: Optional[Run] = None + + @classmethod + def sync( + cls, + fine_tune_job_id: Optional[str] = None, + openai_client: Optional[OpenAI] = None, + num_fine_tunes: Optional[int] = None, + project: str = "OpenAI-Fine-Tune", + entity: Optional[str] = None, + overwrite: bool = False, + wait_for_job_success: bool = True, + log_datasets: bool = True, + model_artifact_name: str = "model-metadata", + model_artifact_type: str = "model", + **kwargs_wandb_init: Dict[str, Any], + ) -> str: + """Sync fine-tunes to Weights & Biases. + + :param fine_tune_job_id: The id of the fine-tune (optional) + :param openai_client: Pass the `OpenAI()` client (optional) + :param num_fine_tunes: Number of most recent fine-tunes to log when an fine_tune_job_id is not provided. By default, every fine-tune is synced. + :param project: Name of the project where you're sending runs. By default, it is "GPT-3". + :param entity: Username or team name where you're sending runs. By default, your default entity is used, which is usually your username. + :param overwrite: Forces logging and overwrite existing wandb run of the same fine-tune. + :param wait_for_job_success: Waits for the fine-tune to be complete and then log metrics to W&B. By default, it is True. + :param model_artifact_name: Name of the model artifact that is logged + :param model_artifact_type: Type of the model artifact that is logged + """ + if openai_client is None: + openai_client = OpenAI() + cls.openai_client = openai_client + + if fine_tune_job_id: + wandb.termlog("Retrieving fine-tune job...") + fine_tune = openai_client.fine_tuning.jobs.retrieve( + fine_tuning_job_id=fine_tune_job_id + ) + fine_tunes = [fine_tune] + else: + # get list of fine_tune to log + fine_tunes = openai_client.fine_tuning.jobs.list() + if not fine_tunes or fine_tunes.data is None: + wandb.termwarn("No fine-tune has been retrieved") + return + # Select the `num_fine_tunes` from the `fine_tunes.data` list. + # If `num_fine_tunes` is None, it selects all items in the list (from start to end). + # If for example, `num_fine_tunes` is 5, it selects the last 5 items in the list. + # Note that the last items in the list are the latest fine-tune jobs. + fine_tunes = fine_tunes.data[ + -num_fine_tunes if num_fine_tunes is not None else None : + ] + + # log starting from oldest fine_tune + show_individual_warnings = ( + fine_tune_job_id is not None or num_fine_tunes is not None + ) + fine_tune_logged = [] + for fine_tune in fine_tunes: + fine_tune_id = fine_tune.id + # check run with the given `fine_tune_id` has not been logged already + run_path = f"{project}/{fine_tune_id}" + if entity is not None: + run_path = f"{entity}/{run_path}" + wandb_run = cls._get_wandb_run(run_path) + if wandb_run: + wandb_status = wandb_run.summary.get("status") + if show_individual_warnings: + if wandb_status == "succeeded" and not overwrite: + wandb.termwarn( + f"Fine-tune {fine_tune_id} has already been logged successfully at {wandb_run.url}. " + "Use `overwrite=True` if you want to overwrite previous run" + ) + elif wandb_status != "succeeded" or overwrite: + if wandb_status != "succeeded": + wandb.termwarn( + f"A run for fine-tune {fine_tune_id} was previously created but didn't end successfully" + ) + wandb.termlog( + f"A new wandb run will be created for fine-tune {fine_tune_id} and previous run will be overwritten" + ) + overwrite = True + if wandb_status == "succeeded" and not overwrite: + return + + # check if the user has not created a wandb run externally + if wandb.run is None: + cls._run = wandb.init( + job_type="fine-tune", + project=project, + entity=entity, + name=fine_tune_id, + id=fine_tune_id, + **kwargs_wandb_init, + ) + else: + # if a run exits - created externally + cls._run = wandb.run + + if wait_for_job_success: + fine_tune = cls._wait_for_job_success(fine_tune) + + cls._log_fine_tune( + fine_tune, + project, + entity, + overwrite, + show_individual_warnings, + log_datasets, + model_artifact_name, + model_artifact_type, + **kwargs_wandb_init, + ) + + if not show_individual_warnings and not any(fine_tune_logged): + wandb.termwarn("No new successful fine-tunes were found") + + return "🎉 wandb sync completed successfully" + + @classmethod + def _wait_for_job_success(cls, fine_tune: FineTuningJob) -> FineTuningJob: + wandb.termlog("Waiting for the OpenAI fine-tuning job to finish training...") + wandb.termlog( + "To avoid blocking, you can call `WandbLogger.sync` with `wait_for_job_success=False` after OpenAI training completes." + ) + while True: + if fine_tune.status == "succeeded": + wandb.termlog( + "Fine-tuning finished, logging metrics, model metadata, and run metadata to Weights & Biases" + ) + return fine_tune + if fine_tune.status == "failed": + wandb.termwarn( + f"Fine-tune {fine_tune.id} has failed and will not be logged" + ) + return fine_tune + if fine_tune.status == "cancelled": + wandb.termwarn( + f"Fine-tune {fine_tune.id} was cancelled and will not be logged" + ) + return fine_tune + time.sleep(10) + fine_tune = cls.openai_client.fine_tuning.jobs.retrieve( + fine_tuning_job_id=fine_tune.id + ) + + @classmethod + def _log_fine_tune( + cls, + fine_tune: FineTuningJob, + project: str, + entity: Optional[str], + overwrite: bool, + show_individual_warnings: bool, + log_datasets: bool, + model_artifact_name: str, + model_artifact_type: str, + **kwargs_wandb_init: Dict[str, Any], + ): + fine_tune_id = fine_tune.id + status = fine_tune.status + + with telemetry.context(run=cls._run) as tel: + tel.feature.openai_finetuning = True + + # check run completed successfully + if status != "succeeded": + if show_individual_warnings: + wandb.termwarn( + f'Fine-tune {fine_tune_id} has the status "{status}" and will not be logged' + ) + return + + # check results are present + try: + results_id = fine_tune.result_files[0] + try: + encoded_results = cls.openai_client.files.content( + file_id=results_id + ).read() + results = base64.b64decode(encoded_results).decode("utf-8") + except Exception: + # attempt to read as text, works for older jobs + results = cls.openai_client.files.content(file_id=results_id).text + except openai.NotFoundError: + if show_individual_warnings: + wandb.termwarn( + f"Fine-tune {fine_tune_id} has no results and will not be logged" + ) + return + + # update the config + cls._run.config.update(cls._get_config(fine_tune)) + + # log results + df_results = pd.read_csv(io.StringIO(results)) + for _, row in df_results.iterrows(): + metrics = {k: v for k, v in row.items() if not np.isnan(v)} + step = metrics.pop("step") + if step is not None: + step = int(step) + cls._run.log(metrics, step=step) + fine_tuned_model = fine_tune.fine_tuned_model + if fine_tuned_model is not None: + cls._run.summary["fine_tuned_model"] = fine_tuned_model + + # training/validation files and fine-tune details + cls._log_artifacts( + fine_tune, + project, + entity, + log_datasets, + overwrite, + model_artifact_name, + model_artifact_type, + ) + + # mark run as complete + cls._run.summary["status"] = "succeeded" + + cls._run.finish() + return True + + @classmethod + def _ensure_logged_in(cls): + if not cls._logged_in: + if wandb.login(): + cls._logged_in = True + else: + raise Exception( + "It appears you are not currently logged in to Weights & Biases. " + "Please run `wandb login` in your terminal or `wandb.login()` in a notebook." + "When prompted, you can obtain your API key by visiting wandb.ai/authorize." + ) + + @classmethod + def _get_wandb_run(cls, run_path: str): + cls._ensure_logged_in() + try: + if cls._wandb_api is None: + cls._wandb_api = wandb.Api() + return cls._wandb_api.run(run_path) + except Exception: + return None + + @classmethod + def _get_wandb_artifact(cls, artifact_path: str): + cls._ensure_logged_in() + try: + if cls._wandb_api is None: + cls._wandb_api = wandb.Api() + return cls._wandb_api.artifact(artifact_path) + except Exception: + return None + + @classmethod + def _get_config(cls, fine_tune: FineTuningJob) -> Dict[str, Any]: + config = dict(fine_tune) + config["result_files"] = config["result_files"][0] + if config.get("created_at"): + config["created_at"] = datetime.datetime.fromtimestamp( + config["created_at"] + ).strftime("%Y-%m-%d %H:%M:%S") + if config.get("finished_at"): + config["finished_at"] = datetime.datetime.fromtimestamp( + config["finished_at"] + ).strftime("%Y-%m-%d %H:%M:%S") + if config.get("hyperparameters"): + config["hyperparameters"] = cls.sanitize(config["hyperparameters"]) + if config.get("error"): + config["error"] = cls.sanitize(config["error"]) + return config + + @classmethod + def _unpack_hyperparameters(cls, hyperparameters: Hyperparameters): + # `Hyperparameters` object is not unpacking properly using `vars` or `__dict__`, + # vars(hyperparameters) return {n_epochs: n} only. + hyperparams = {} + try: + hyperparams["n_epochs"] = hyperparameters.n_epochs + hyperparams["batch_size"] = hyperparameters.batch_size + hyperparams["learning_rate_multiplier"] = ( + hyperparameters.learning_rate_multiplier + ) + except Exception: + # If unpacking fails, return the object to be logged as config + return None + + return hyperparams + + @staticmethod + def sanitize(input: Any) -> Union[Dict, List, str]: + valid_types = [bool, int, float, str] + if isinstance(input, (Hyperparameters, Error)): + return dict(input) + if isinstance(input, dict): + return { + k: v if type(v) in valid_types else str(v) for k, v in input.items() + } + elif isinstance(input, list): + return [v if type(v) in valid_types else str(v) for v in input] + else: + return str(input) + + @classmethod + def _log_artifacts( + cls, + fine_tune: FineTuningJob, + project: str, + entity: Optional[str], + log_datasets: bool, + overwrite: bool, + model_artifact_name: str, + model_artifact_type: str, + ) -> None: + if log_datasets: + wandb.termlog("Logging training/validation files...") + # training/validation files + training_file = fine_tune.training_file if fine_tune.training_file else None + validation_file = ( + fine_tune.validation_file if fine_tune.validation_file else None + ) + for file, prefix, artifact_type in ( + (training_file, "train", "training_files"), + (validation_file, "valid", "validation_files"), + ): + if file is not None: + cls._log_artifact_inputs( + file, prefix, artifact_type, project, entity, overwrite + ) + + # fine-tune details + fine_tune_id = fine_tune.id + artifact = wandb.Artifact( + model_artifact_name, + type=model_artifact_type, + metadata=dict(fine_tune), + ) + + with artifact.new_file("model_metadata.json", mode="w", encoding="utf-8") as f: + dict_fine_tune = dict(fine_tune) + dict_fine_tune["hyperparameters"] = cls.sanitize( + dict_fine_tune["hyperparameters"] + ) + dict_fine_tune["error"] = cls.sanitize(dict_fine_tune["error"]) + dict_fine_tune = cls.sanitize(dict_fine_tune) + json.dump(dict_fine_tune, f, indent=2) + cls._run.log_artifact( + artifact, + aliases=["latest", fine_tune_id], + ) + + @classmethod + def _log_artifact_inputs( + cls, + file_id: Optional[str], + prefix: str, + artifact_type: str, + project: str, + entity: Optional[str], + overwrite: bool, + ) -> None: + # get input artifact + artifact_name = f"{prefix}-{file_id}" + # sanitize name to valid wandb artifact name + artifact_name = re.sub(r"[^a-zA-Z0-9_\-.]", "_", artifact_name) + artifact_alias = file_id + artifact_path = f"{project}/{artifact_name}:{artifact_alias}" + if entity is not None: + artifact_path = f"{entity}/{artifact_path}" + artifact = cls._get_wandb_artifact(artifact_path) + + # create artifact if file not already logged previously + if artifact is None or overwrite: + # get file content + try: + file_content = cls.openai_client.files.content(file_id=file_id) + except openai.NotFoundError: + wandb.termerror( + f"File {file_id} could not be retrieved. Make sure you have OpenAI permissions to download training/validation files" + ) + return + + artifact = wandb.Artifact(artifact_name, type=artifact_type) + with tempfile.NamedTemporaryFile(delete=False) as tmp_file: + tmp_file.write(file_content.content) + tmp_file_path = tmp_file.name + artifact.add_file(tmp_file_path, file_id) + os.unlink(tmp_file_path) + + # create a Table + try: + table, n_items = cls._make_table(file_content.text) + # Add table to the artifact. + artifact.add(table, file_id) + # Add the same table to the workspace. + cls._run.log({f"{prefix}_data": table}) + # Update the run config and artifact metadata + cls._run.config.update({f"n_{prefix}": n_items}) + artifact.metadata["items"] = n_items + except Exception as e: + wandb.termerror( + f"Issue saving {file_id} as a Table to Artifacts, exception:\n '{e}'" + ) + else: + # log number of items + cls._run.config.update({f"n_{prefix}": artifact.metadata.get("items")}) + + cls._run.use_artifact(artifact, aliases=["latest", artifact_alias]) + + @classmethod + def _make_table(cls, file_content: str) -> Tuple[Table, int]: + table = wandb.Table(columns=["role: system", "role: user", "role: assistant"]) + + df = pd.read_json(io.StringIO(file_content), orient="records", lines=True) + for _idx, message in df.iterrows(): + messages = message.messages + assert len(messages) == 3 + table.add_data( + messages[0]["content"], + messages[1]["content"], + messages[2]["content"], + ) + + return table, len(df) diff --git a/lib/python3.12/site-packages/wandb/integration/openai/openai.py b/lib/python3.12/site-packages/wandb/integration/openai/openai.py new file mode 100644 index 0000000000000000000000000000000000000000..250d437d7239bbcd31cc950c80563519f87e5fde --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/openai/openai.py @@ -0,0 +1,22 @@ +import logging + +from wandb.sdk.integration_utils.auto_logging import AutologAPI + +from .resolver import OpenAIRequestResponseResolver + +logger = logging.getLogger(__name__) + + +autolog = AutologAPI( + name="OpenAI", + symbols=( + "Edit.create", + "Completion.create", + "ChatCompletion.create", + "Edit.acreate", + "Completion.acreate", + "ChatCompletion.acreate", + ), + resolver=OpenAIRequestResponseResolver(), + telemetry_feature="openai_autolog", +) diff --git a/lib/python3.12/site-packages/wandb/integration/openai/resolver.py b/lib/python3.12/site-packages/wandb/integration/openai/resolver.py new file mode 100644 index 0000000000000000000000000000000000000000..500c58ce2f4b33a6a987243169bc998cf387fef4 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/openai/resolver.py @@ -0,0 +1,240 @@ +import datetime +import io +import logging +from dataclasses import asdict, dataclass +from typing import Any, Dict, List, Optional, Sequence + +import wandb +from wandb.sdk.data_types import trace_tree +from wandb.sdk.integration_utils.auto_logging import Response + +logger = logging.getLogger(__name__) + + +@dataclass +class UsageMetrics: + elapsed_time: float = None + prompt_tokens: int = None + completion_tokens: int = None + total_tokens: int = None + + +@dataclass +class Metrics: + usage: UsageMetrics = None + stats: wandb.Table = None + trace: trace_tree.WBTraceTree = None + + +usage_metric_keys = {f"usage/{k}" for k in asdict(UsageMetrics())} + + +class OpenAIRequestResponseResolver: + def __init__(self): + self.define_metrics_called = False + + def __call__( + self, + args: Sequence[Any], + kwargs: Dict[str, Any], + response: Response, + start_time: float, # pass to comply with the protocol, but use response["created"] instead + time_elapsed: float, + ) -> Optional[Dict[str, Any]]: + request = kwargs + + if not self.define_metrics_called: + # define metrics on first call + for key in usage_metric_keys: + wandb.define_metric(key, step_metric="_timestamp") + self.define_metrics_called = True + + try: + if response.get("object") == "edit": + return self._resolve_edit(request, response, time_elapsed) + elif response.get("object") == "text_completion": + return self._resolve_completion(request, response, time_elapsed) + elif response.get("object") == "chat.completion": + return self._resolve_chat_completion(request, response, time_elapsed) + else: + # todo: properly treat failed requests + logger.info( + f"Unsupported OpenAI response object: {response.get('object')}" + ) + except Exception as e: + logger.warning(f"Failed to resolve request/response: {e}") + return None + + @staticmethod + def results_to_trace_tree( + request: Dict[str, Any], + response: Response, + results: List[trace_tree.Result], + time_elapsed: float, + ) -> trace_tree.WBTraceTree: + """Converts the request, response, and results into a trace tree. + + params: + request: The request dictionary + response: The response object + results: A list of results object + time_elapsed: The time elapsed in seconds + returns: + A wandb trace tree object. + """ + start_time_ms = int(round(response["created"] * 1000)) + end_time_ms = start_time_ms + int(round(time_elapsed * 1000)) + span = trace_tree.Span( + name=f"{response.get('model', 'openai')}_{response['object']}_{response.get('created')}", + attributes=dict(response), # type: ignore + start_time_ms=start_time_ms, + end_time_ms=end_time_ms, + span_kind=trace_tree.SpanKind.LLM, + results=results, + ) + model_obj = {"request": request, "response": response, "_kind": "openai"} + return trace_tree.WBTraceTree(root_span=span, model_dict=model_obj) + + def _resolve_edit( + self, + request: Dict[str, Any], + response: Response, + time_elapsed: float, + ) -> Dict[str, Any]: + """Resolves the request and response objects for `openai.Edit`.""" + request_str = ( + f"\n\n**Instruction**: {request['instruction']}\n\n" + f"**Input**: {request['input']}\n" + ) + choices = [ + f"\n\n**Edited**: {choice['text']}\n" for choice in response["choices"] + ] + + return self._resolve_metrics( + request=request, + response=response, + request_str=request_str, + choices=choices, + time_elapsed=time_elapsed, + ) + + def _resolve_completion( + self, + request: Dict[str, Any], + response: Response, + time_elapsed: float, + ) -> Dict[str, Any]: + """Resolves the request and response objects for `openai.Completion`.""" + request_str = f"\n\n**Prompt**: {request['prompt']}\n" + choices = [ + f"\n\n**Completion**: {choice['text']}\n" for choice in response["choices"] + ] + + return self._resolve_metrics( + request=request, + response=response, + request_str=request_str, + choices=choices, + time_elapsed=time_elapsed, + ) + + def _resolve_chat_completion( + self, + request: Dict[str, Any], + response: Response, + time_elapsed: float, + ) -> Dict[str, Any]: + """Resolves the request and response objects for `openai.Completion`.""" + prompt = io.StringIO() + for message in request["messages"]: + prompt.write(f"\n\n**{message['role']}**: {message['content']}\n") + request_str = prompt.getvalue() + + choices = [ + f"\n\n**{choice['message']['role']}**: {choice['message']['content']}\n" + for choice in response["choices"] + ] + + return self._resolve_metrics( + request=request, + response=response, + request_str=request_str, + choices=choices, + time_elapsed=time_elapsed, + ) + + def _resolve_metrics( + self, + request: Dict[str, Any], + response: Response, + request_str: str, + choices: List[str], + time_elapsed: float, + ) -> Dict[str, Any]: + """Resolves the request and response objects for `openai.Completion`.""" + results = [ + trace_tree.Result( + inputs={"request": request_str}, + outputs={"response": choice}, + ) + for choice in choices + ] + metrics = self._get_metrics_to_log(request, response, results, time_elapsed) + return self._convert_metrics_to_dict(metrics) + + @staticmethod + def _get_usage_metrics(response: Response, time_elapsed: float) -> UsageMetrics: + """Gets the usage stats from the response object.""" + if response.get("usage"): + usage_stats = UsageMetrics(**response["usage"]) + else: + usage_stats = UsageMetrics() + usage_stats.elapsed_time = time_elapsed + return usage_stats + + def _get_metrics_to_log( + self, + request: Dict[str, Any], + response: Response, + results: List[Any], + time_elapsed: float, + ) -> Metrics: + model = response.get("model") or request.get("model") + usage_metrics = self._get_usage_metrics(response, time_elapsed) + + usage = [] + for result in results: + row = { + "request": result.inputs["request"], + "response": result.outputs["response"], + "model": model, + "start_time": datetime.datetime.fromtimestamp(response["created"]), + "end_time": datetime.datetime.fromtimestamp( + response["created"] + time_elapsed + ), + "request_id": response.get("id", None), + "api_type": response.get("api_type", "openai"), + "session_id": wandb.run.id, + } + row.update(asdict(usage_metrics)) + usage.append(row) + usage_table = wandb.Table( + columns=list(usage[0].keys()), + data=[(item.values()) for item in usage], + ) + + trace = self.results_to_trace_tree(request, response, results, time_elapsed) + + metrics = Metrics(stats=usage_table, trace=trace, usage=usage_metrics) + return metrics + + @staticmethod + def _convert_metrics_to_dict(metrics: Metrics) -> Dict[str, Any]: + """Converts metrics to a dict.""" + metrics_dict = { + "stats": metrics.stats, + "trace": metrics.trace, + } + usage_stats = {f"usage/{k}": v for k, v in asdict(metrics.usage).items()} + metrics_dict.update(usage_stats) + return metrics_dict diff --git a/lib/python3.12/site-packages/wandb/integration/prodigy/__init__.py b/lib/python3.12/site-packages/wandb/integration/prodigy/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..94ed0ea26bc5f886cf90db26da37dc8f08f9b18e --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/prodigy/__init__.py @@ -0,0 +1,3 @@ +from .prodigy import upload_dataset + +__all__ = ["upload_dataset"] diff --git a/lib/python3.12/site-packages/wandb/integration/prodigy/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/prodigy/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1db5efa195b586aeab5e4c6ea2c6fe07978b3a27 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/prodigy/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/prodigy/__pycache__/prodigy.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/prodigy/__pycache__/prodigy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7380f8f5c7809b132775e9b769f46652deadb1ff Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/prodigy/__pycache__/prodigy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/prodigy/prodigy.py b/lib/python3.12/site-packages/wandb/integration/prodigy/prodigy.py new file mode 100644 index 0000000000000000000000000000000000000000..d25fc3cbab7a2c49aa585631e7162485487945e6 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/prodigy/prodigy.py @@ -0,0 +1,291 @@ +"""Prodigy integration for W&B. + +User can upload Prodigy annotated datasets directly +from the local database to W&B in Tables format. + +Example usage: + +```python +import wandb +from wandb.integration.prodigy import upload_dataset + +run = wandb.init(project="prodigy") +upload_dataset("name_of_dataset") +wandb.finish() +``` +""" + +import base64 +import collections.abc +import io +import urllib +from copy import deepcopy + +import pandas as pd +from PIL import Image + +import wandb +from wandb import util +from wandb.plot.utils import test_missing +from wandb.sdk.lib import telemetry as wb_telemetry + + +def named_entity(docs): + """Create a named entity visualization. + + Taken from https://github.com/wandb/wandb/blob/main/wandb/plots/named_entity.py. + """ + spacy = util.get_module( + "spacy", + required="part_of_speech requires the spacy library, install with `pip install spacy`", + ) + + util.get_module( + "en_core_web_md", + required="part_of_speech requires `en_core_web_md` library, install with `python -m spacy download en_core_web_md`", + ) + + # Test for required packages and missing & non-integer values in docs data + if test_missing(docs=docs): + html = spacy.displacy.render( + docs, style="ent", page=True, minify=True, jupyter=False + ) + wandb_html = wandb.Html(html) + return wandb_html + + +def merge(dict1, dict2): + """Return a new dictionary by merging two dictionaries recursively.""" + result = deepcopy(dict1) + + for key, value in dict2.items(): + if isinstance(value, collections.abc.Mapping): + result[key] = merge(result.get(key, {}), value) + else: + result[key] = deepcopy(dict2[key]) + + return result + + +def get_schema(list_data_dict, struct, array_dict_types): + """Get a schema of the dataset's structure and data types.""" + # Get the structure of the JSON objects in the database + # This is similar to getting a JSON schema but with slightly different format + for _i, item in enumerate(list_data_dict): + # If the list contains dict objects + for k, v in item.items(): + # Check if key already exists in template + if k not in struct.keys(): + if isinstance(v, list): + if len(v) > 0 and isinstance(v[0], list): + # nested list structure + struct[k] = type(v) # type list + elif len(v) > 0 and not ( + isinstance(v[0], list) or isinstance(v[0], dict) + ): + # list of singular values + struct[k] = type(v) # type list + else: + # list of dicts + array_dict_types.append( + k + ) # keep track of keys that are type list[dict] + struct[k] = {} + struct[k] = get_schema(v, struct[k], array_dict_types) + elif isinstance(v, dict): + struct[k] = {} + struct[k] = get_schema([v], struct[k], array_dict_types) + else: + struct[k] = type(v) + else: + # Get the value of struct[k] which is the current template + # Find new keys and then merge the two templates together + cur_struct = struct[k] + if isinstance(v, list): + if len(v) > 0 and isinstance(v[0], list): + # nested list coordinate structure + # if the value in the item is currently None, then update + if v is not None: + struct[k] = type(v) # type list + elif len(v) > 0 and not ( + isinstance(v[0], list) or isinstance(v[0], dict) + ): + # single list with values + # if the value in the item is currently None, then update + if v is not None: + struct[k] = type(v) # type list + else: + array_dict_types.append( + k + ) # keep track of keys that are type list[dict] + struct[k] = {} + struct[k] = get_schema(v, struct[k], array_dict_types) + # merge cur_struct and struct[k], remove duplicates + struct[k] = merge(struct[k], cur_struct) + elif isinstance(v, dict): + struct[k] = {} + struct[k] = get_schema([v], struct[k], array_dict_types) + # merge cur_struct and struct[k], remove duplicates + struct[k] = merge(struct[k], cur_struct) + else: + # if the value in the item is currently None, then update + if v is not None: + struct[k] = type(v) + + return struct + + +def standardize(item, structure, array_dict_types): + """Standardize all rows/entries in dataset to fit the schema. + + Will look for missing values and fill it in so all rows have + the same items and structure. + """ + for k, v in structure.items(): + if k not in item: + # If the structure/field does not exist + if isinstance(v, dict) and (k not in array_dict_types): + # If key k is of type dict, and not not a type list[dict] + item[k] = {} + standardize(item[k], v, array_dict_types) + elif isinstance(v, dict) and (k in array_dict_types): + # If key k is of type dict, and is actually of type list[dict], + # just treat as a list and set to None by default + item[k] = None + else: + # Assign a default type + item[k] = v() + else: + # If the structure/field already exists and is a list or dict + if isinstance(item[k], list): + # ignore if item is a nested list structure or list of non-dicts + condition = ( + not (len(item[k]) > 0 and isinstance(item[k][0], list)) + ) and ( + not ( + len(item[k]) > 0 + and not ( + isinstance(item[k][0], list) or isinstance(item[k][0], dict) + ) + ) + ) + if condition: + for sub_item in item[k]: + standardize(sub_item, v, array_dict_types) + elif isinstance(item[k], dict): + standardize(item[k], v, array_dict_types) + + +def create_table(data): + """Create a W&B Table. + + - Create/decode images from URL/Base64 + - Uses spacy to translate NER span data to visualizations. + """ + # create table object from columns + table_df = pd.DataFrame(data) + columns = list(table_df.columns) + if ("spans" in table_df.columns) and ("text" in table_df.columns): + columns.append("spans_visual") + if "image" in columns: + columns.append("image_visual") + main_table = wandb.Table(columns=columns) + + # Convert to dictionary format to maintain order during processing + matrix = table_df.to_dict(orient="records") + + # Import en_core_web_md if exists + en_core_web_md = util.get_module( + "en_core_web_md", + required="part_of_speech requires `en_core_web_md` library, install with `python -m spacy download en_core_web_md`", + ) + nlp = en_core_web_md.load(disable=["ner"]) + + # Go through each individual row + for _i, document in enumerate(matrix): + # Text NER span visualizations + if ("spans_visual" in columns) and ("text" in columns): + # Add visuals for spans + document["spans_visual"] = None + doc = nlp(document["text"]) + ents = [] + if ("spans" in document) and (document["spans"] is not None): + for span in document["spans"]: + if ("start" in span) and ("end" in span) and ("label" in span): + charspan = doc.char_span( + span["start"], span["end"], span["label"] + ) + ents.append(charspan) + doc.ents = ents + document["spans_visual"] = named_entity(docs=doc) + + # Convert image link to wandb Image + if "image" in columns: + # Turn into wandb image + document["image_visual"] = None + if ("image" in document) and (document["image"] is not None): + isurl = urllib.parse.urlparse(document["image"]).scheme in ( + "http", + "https", + ) + isbase64 = ("data:" in document["image"]) and ( + ";base64" in document["image"] + ) + if isurl: + # is url + try: + im = Image.open(urllib.request.urlopen(document["image"])) + document["image_visual"] = wandb.Image(im) + except urllib.error.URLError: + wandb.termwarn(f"Image URL {document['image']} is invalid.") + document["image_visual"] = None + elif isbase64: + # is base64 uri + imgb64 = document["image"].split("base64,")[1] + try: + msg = base64.b64decode(imgb64) + buf = io.BytesIO(msg) + im = Image.open(buf) + document["image_visual"] = wandb.Image(im) + except base64.binascii.Error: + wandb.termwarn(f"Base64 string {document['image']} is invalid.") + document["image_visual"] = None + else: + # is data path + document["image_visual"] = wandb.Image(document["image"]) + + # Create row and append to table + values_list = list(document.values()) + main_table.add_data(*values_list) + return main_table + + +def upload_dataset(dataset_name): + """Upload dataset from local database to Weights & Biases. + + Args: + dataset_name: The name of the dataset in the Prodigy database. + """ + # Check if wandb.init has been called + if wandb.run is None: + raise ValueError("You must call wandb.init() before upload_dataset()") + + with wb_telemetry.context(run=wandb.run) as tel: + tel.feature.prodigy = True + + prodigy_db = util.get_module( + "prodigy.components.db", + required="`prodigy` library is required but not installed. Please see https://prodi.gy/docs/install", + ) + # Retrieve and upload prodigy dataset + database = prodigy_db.connect() + data = database.get_dataset(dataset_name) + + array_dict_types = [] + schema = get_schema(data, {}, array_dict_types) + + for i, _d in enumerate(data): + standardize(data[i], schema, array_dict_types) + table = create_table(data) + wandb.log({dataset_name: table}) + wandb.termlog(f"Prodigy dataset `{dataset_name}` uploaded.") diff --git a/lib/python3.12/site-packages/wandb/integration/sacred/__init__.py b/lib/python3.12/site-packages/wandb/integration/sacred/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..dabcc47afd726c57cf269e698096e399a1574b3a --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/sacred/__init__.py @@ -0,0 +1,117 @@ +import warnings + +import numpy +from sacred.dependencies import get_digest +from sacred.observers import RunObserver + +import wandb + + +class WandbObserver(RunObserver): + """Log sacred experiment data to W&B. + + Args: + Accepts all the arguments accepted by wandb.init(). + + name — A display name for this run, which shows up in the UI and is editable, doesn't have to be unique + notes — A multiline string description associated with the run + config — a dictionary-like object to set as initial config + project — the name of the project to which this run will belong + tags — a list of strings to associate with this run as tags + dir — the path to a directory where artifacts will be written (default: ./wandb) + entity — the team posting this run (default: your username or your default team) + job_type — the type of job you are logging, e.g. eval, worker, ps (default: training) + save_code — save the main python or notebook file to wandb to enable diffing (default: editable from your settings page) + group — a string by which to group other runs; see Grouping + reinit — Shorthand for the reinit setting that defines what to do when `wandb.init()` is called while a run is active. See the setting's documentation. + id — A unique ID for this run primarily used for Resuming. It must be globally unique, and if you delete a run you can't reuse the ID. Use the name field for a descriptive, useful name for the run. The ID cannot contain special characters. + resume — if set to True, the run auto resumes; can also be a unique string for manual resuming; see Resuming (default: False) + anonymous — can be "allow", "never", or "must". This enables or explicitly disables anonymous logging. (default: never) + force — whether to force a user to be logged into wandb when running a script (default: False) + magic — (bool, dict, or str, optional): magic configuration as bool, dict, json string, yaml filename. If set to True will attempt to auto-instrument your script. (default: None) + sync_tensorboard — A boolean indicating whether or not copy all TensorBoard logs wandb; see Tensorboard (default: False) + monitor_gym — A boolean indicating whether or not to log videos generated by OpenAI Gym; see Ray Tune (default: False) + allow_val_change — whether to allow wandb.config values to change, by default we throw an exception if config values are overwritten. (default: False) + + Examples: + Create sacred experiment:: + from wandb.sacred import WandbObserver + ex.observers.append(WandbObserver(project='sacred_test', + name='test1')) + @ex.config + def cfg(): + C = 1.0 + gamma = 0.7 + @ex.automain + def run(C, gamma, _run): + iris = datasets.load_iris() + per = permutation(iris.target.size) + iris.data = iris.data[per] + iris.target = iris.target[per] + clf = svm.SVC(C, 'rbf', gamma=gamma) + clf.fit(iris.data[:90], + iris.target[:90]) + return clf.score(iris.data[90:], + iris.target[90:]) + """ + + def __init__(self, **kwargs): + self.run = wandb.init(**kwargs) + self.resources = {} + + def started_event( + self, ex_info, command, host_info, start_time, config, meta_info, _id + ): + # TODO: add the source code file + # TODO: add dependencies and metadata. + self.__update_config(config) + + def completed_event(self, stop_time, result): + if result: + if not isinstance(result, tuple): + result = ( + result, + ) # transform single result to tuple so that both single & multiple results use same code + + for i, r in enumerate(result): + if isinstance(r, float) or isinstance(r, int): + wandb.log({f"result_{i}": float(r)}) + elif isinstance(r, dict): + wandb.log(r) + elif isinstance(r, object): + artifact = wandb.Artifact(f"result_{i}.pkl", type="result") + artifact.add_file(r) + self.run.log_artifact(artifact) + elif isinstance(r, numpy.ndarray): + wandb.log({f"result_{i}": wandb.Image(r)}) + else: + warnings.warn( + f"logging results does not support type '{type(r)}' results. Ignoring this result", + stacklevel=2, + ) + + def artifact_event(self, name, filename, metadata=None, content_type=None): + if content_type is None: + content_type = "file" + artifact = wandb.Artifact(name, type=content_type) + artifact.add_file(filename) + self.run.log_artifact(artifact) + + def resource_event(self, filename): + """TODO: Maintain resources list.""" + if filename not in self.resources: + md5 = get_digest(filename) + self.resources[filename] = md5 + + def log_metrics(self, metrics_by_name, info): + for metric_name, metric_ptr in metrics_by_name.items(): + for _step, value in zip(metric_ptr["steps"], metric_ptr["values"]): + if isinstance(value, numpy.ndarray): + wandb.log({metric_name: wandb.Image(value)}) + else: + wandb.log({metric_name: value}) + + def __update_config(self, config): + for k, v in config.items(): + self.run.config[k] = v + self.run.config["resources"] = [] diff --git a/lib/python3.12/site-packages/wandb/integration/sacred/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/sacred/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b01cd6a52806168810dcd190244e9c9b52dc6e9e Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/sacred/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/sagemaker/__init__.py b/lib/python3.12/site-packages/wandb/integration/sagemaker/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5c8509c8fae0a795d4f17bedd708c87ee587e931 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/sagemaker/__init__.py @@ -0,0 +1,14 @@ +"""wandb integration sagemaker module.""" + +from .auth import sagemaker_auth +from .config import is_using_sagemaker, parse_sm_config +from .resources import parse_sm_secrets, set_global_settings, set_run_id + +__all__ = [ + "sagemaker_auth", + "is_using_sagemaker", + "parse_sm_config", + "parse_sm_secrets", + "set_global_settings", + "set_run_id", +] diff --git a/lib/python3.12/site-packages/wandb/integration/sagemaker/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/sagemaker/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6bad1b87df4c8981406307034adb9bc01997cc5e Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/sagemaker/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/sagemaker/__pycache__/auth.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/sagemaker/__pycache__/auth.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1bacc75a65eeadce8cac848e0353544c72c168d7 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/sagemaker/__pycache__/auth.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/sagemaker/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/sagemaker/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7d5249f24c8c315db4bf98287badb818b1751099 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/sagemaker/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/sagemaker/__pycache__/files.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/sagemaker/__pycache__/files.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..89eac384b3ff0e02a7c19f6ed9ee096847ef0331 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/sagemaker/__pycache__/files.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/sagemaker/__pycache__/resources.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/sagemaker/__pycache__/resources.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ef13bc594a29654414a38934cce6c284de3075ae Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/sagemaker/__pycache__/resources.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/sagemaker/auth.py b/lib/python3.12/site-packages/wandb/integration/sagemaker/auth.py new file mode 100644 index 0000000000000000000000000000000000000000..fc82952db6dd710c1e474640cf939476ca171c8d --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/sagemaker/auth.py @@ -0,0 +1,29 @@ +import os + +import wandb +from wandb import env +from wandb.sdk import wandb_setup + + +def sagemaker_auth(overrides=None, path=".", api_key=None): + """Write a secrets.env file with the W&B ApiKey and any additional secrets passed. + + Args: + overrides (dict, optional): Additional environment variables to write + to secrets.env + path (str, optional): The path to write the secrets file. + """ + settings = wandb_setup.singleton().settings + current_api_key = wandb.wandb_lib.apikey.api_key(settings=settings) + + overrides = overrides or dict() + api_key = overrides.get(env.API_KEY, api_key or current_api_key) + if api_key is None: + raise ValueError( + "Can't find W&B ApiKey, set the WANDB_API_KEY env variable " + "or run `wandb login`" + ) + overrides[env.API_KEY] = api_key + with open(os.path.join(path, "secrets.env"), "w") as file: + for k, v in overrides.items(): + file.write(f"{k}={v}\n") diff --git a/lib/python3.12/site-packages/wandb/integration/sagemaker/config.py b/lib/python3.12/site-packages/wandb/integration/sagemaker/config.py new file mode 100644 index 0000000000000000000000000000000000000000..be71b92c19d204b249c277a1c0c82c70bab3a044 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/sagemaker/config.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +import json +import os +import re +import warnings +from typing import Any + +from . import files as sm_files + + +def is_using_sagemaker() -> bool: + """Returns whether we're in a SageMaker environment.""" + return ( + os.path.exists(sm_files.SM_PARAM_CONFIG) # + or "SM_TRAINING_ENV" in os.environ + ) + + +def parse_sm_config() -> dict[str, Any]: + """Parses SageMaker configuration. + + Returns: + A dictionary of SageMaker config keys/values + or an empty dict if not found. + SM_TRAINING_ENV is a json string of the + training environment variables set by SageMaker + and is only available when running in SageMaker, + but not in local mode. + SM_TRAINING_ENV is set by the SageMaker container and + contains arguments such as hyperparameters + and arguments passed to the training job. + """ + conf = {} + + if os.path.exists(sm_files.SM_PARAM_CONFIG): + conf["sagemaker_training_job_name"] = os.getenv("TRAINING_JOB_NAME") + + # Hyperparameter searches quote configs... + with open(sm_files.SM_PARAM_CONFIG) as fid: + for key, val in json.load(fid).items(): + cast = val.strip('"') + if re.match(r"^-?[\d]+$", cast): + cast = int(cast) + elif re.match(r"^-?[.\d]+$", cast): + cast = float(cast) + conf[key] = cast + + if env := os.environ.get("SM_TRAINING_ENV"): + try: + conf.update(json.loads(env)) + except json.JSONDecodeError: + warnings.warn( + "Failed to parse SM_TRAINING_ENV not valid JSON string", + stacklevel=2, + ) + + return conf diff --git a/lib/python3.12/site-packages/wandb/integration/sagemaker/files.py b/lib/python3.12/site-packages/wandb/integration/sagemaker/files.py new file mode 100644 index 0000000000000000000000000000000000000000..1f91e72fb07b5febb7ee6cb89d06f18444fce826 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/sagemaker/files.py @@ -0,0 +1,2 @@ +SM_PARAM_CONFIG = "/opt/ml/input/config/hyperparameters.json" +SM_SECRETS = "secrets.env" diff --git a/lib/python3.12/site-packages/wandb/integration/sagemaker/resources.py b/lib/python3.12/site-packages/wandb/integration/sagemaker/resources.py new file mode 100644 index 0000000000000000000000000000000000000000..4410c755772eead22a846cf25a5beed6fca6cc45 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/sagemaker/resources.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +import os +import secrets +import socket +import string + +import wandb + +from . import config +from . import files as sm_files + + +def set_run_id(run_settings: wandb.Settings) -> bool: + """Set a run ID and group when using SageMaker. + + Returns whether the ID and group were updated. + """ + # Added in https://github.com/wandb/wandb/pull/3290. + # + # Prevents SageMaker from overriding the run ID configured + # in environment variables. Note, however, that it will still + # override a run ID passed explicitly to `wandb.init()`. + if os.getenv("WANDB_RUN_ID"): + return False + + run_group = os.getenv("TRAINING_JOB_NAME") + if not run_group: + return False + + alphanumeric = string.ascii_lowercase + string.digits + random = "".join(secrets.choice(alphanumeric) for _ in range(6)) + + host = os.getenv("CURRENT_HOST", socket.gethostname()) + + run_settings.run_id = f"{run_group}-{random}-{host}" + run_settings.run_group = run_group + return True + + +def set_global_settings(settings: wandb.Settings) -> None: + """Set global W&B settings based on the SageMaker environment.""" + if env := parse_sm_secrets(): + settings.update_from_env_vars(env) + + # The SageMaker config may contain an API key, in which case it + # takes precedence over the value in the secrets. It's unclear + # whether this is by design, or by accident; we keep it for + # backward compatibility for now. + sm_config = config.parse_sm_config() + if api_key := sm_config.get("wandb_api_key"): + settings.api_key = api_key + + +def parse_sm_secrets() -> dict[str, str]: + """We read our api_key from secrets.env in SageMaker.""" + env_dict = dict() + # Set secret variables + if os.path.exists(sm_files.SM_SECRETS): + for line in open(sm_files.SM_SECRETS): + key, val = line.strip().split("=", 1) + env_dict[key] = val + return env_dict diff --git a/lib/python3.12/site-packages/wandb/integration/sb3/__init__.py b/lib/python3.12/site-packages/wandb/integration/sb3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..29dd81a941cb1745cf82fecbcf12a857ce786f45 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/sb3/__init__.py @@ -0,0 +1,3 @@ +from .sb3 import WandbCallback + +__all__ = ["WandbCallback"] diff --git a/lib/python3.12/site-packages/wandb/integration/sb3/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/sb3/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..66d21c0e2547e73816510843282c02970b33a8f6 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/sb3/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/sb3/__pycache__/sb3.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/sb3/__pycache__/sb3.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8e0e217ae41b264111bed7eda4425de89e025151 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/sb3/__pycache__/sb3.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/sb3/sb3.py b/lib/python3.12/site-packages/wandb/integration/sb3/sb3.py new file mode 100644 index 0000000000000000000000000000000000000000..2eec145d0fc052575ca6ba65d68283dcaa0ea69b --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/sb3/sb3.py @@ -0,0 +1,147 @@ +"""W&B callback for sb3. + +Really simple callback to get logging for each tree + +Example usage: + +```python +import gym +from stable_baselines3 import PPO +from stable_baselines3.common.monitor import Monitor +from stable_baselines3.common.vec_env import DummyVecEnv, VecVideoRecorder +import wandb +from wandb.integration.sb3 import WandbCallback + + +config = { + "policy_type": "MlpPolicy", + "total_timesteps": 25000, + "env_name": "CartPole-v1", +} +run = wandb.init( + project="sb3", + config=config, + sync_tensorboard=True, # auto-upload sb3's tensorboard metrics + monitor_gym=True, # auto-upload the videos of agents playing the game + save_code=True, # optional +) + + +def make_env(): + env = gym.make(config["env_name"]) + env = Monitor(env) # record stats such as returns + return env + + +env = DummyVecEnv([make_env]) +env = VecVideoRecorder( + env, "videos", record_video_trigger=lambda x: x % 2000 == 0, video_length=200 +) +model = PPO(config["policy_type"], env, verbose=1, tensorboard_log=f"runs") +model.learn( + total_timesteps=config["total_timesteps"], + callback=WandbCallback( + model_save_path=f"models/{run.id}", + gradient_save_freq=100, + log="all", + ), +) +``` +""" + +import logging +import os +from typing import Literal, Optional + +from stable_baselines3.common.callbacks import BaseCallback # type: ignore + +import wandb +from wandb.sdk.lib import telemetry as wb_telemetry + +logger = logging.getLogger(__name__) + + +class WandbCallback(BaseCallback): + """Callback for logging experiments to Weights and Biases. + + Log SB3 experiments to Weights and Biases + - Added model tracking and uploading + - Added complete hyperparameters recording + - Added gradient logging + - Note that `wandb.init(...)` must be called before the WandbCallback can be used. + + Args: + verbose: The verbosity of sb3 output + model_save_path: Path to the folder where the model will be saved, The default value is `None` so the model is not logged + model_save_freq: Frequency to save the model + gradient_save_freq: Frequency to log gradient. The default value is 0 so the gradients are not logged + log: What to log. One of "gradients", "parameters", or "all". + """ + + def __init__( + self, + verbose: int = 0, + model_save_path: Optional[str] = None, + model_save_freq: int = 0, + gradient_save_freq: int = 0, + log: Optional[Literal["gradients", "parameters", "all"]] = "all", + ) -> None: + super().__init__(verbose) + if wandb.run is None: + raise wandb.Error("You must call wandb.init() before WandbCallback()") + with wb_telemetry.context() as tel: + tel.feature.sb3 = True + self.model_save_freq = model_save_freq + self.model_save_path = model_save_path + self.gradient_save_freq = gradient_save_freq + if log not in ["gradients", "parameters", "all", None]: + wandb.termwarn( + "`log` must be one of `None`, 'gradients', 'parameters', or 'all', " + "falling back to 'all'" + ) + log = "all" + self.log = log + # Create folder if needed + if self.model_save_path is not None: + os.makedirs(self.model_save_path, exist_ok=True) + self.path = os.path.join(self.model_save_path, "model.zip") + else: + assert self.model_save_freq == 0, ( + "to use the `model_save_freq` you have to set the `model_save_path` parameter" + ) + + def _init_callback(self) -> None: + d = {} + if "algo" not in d: + d["algo"] = type(self.model).__name__ + for key in self.model.__dict__: + if key in wandb.config: + continue + if type(self.model.__dict__[key]) in [float, int, str]: + d[key] = self.model.__dict__[key] + else: + d[key] = str(self.model.__dict__[key]) + if self.gradient_save_freq > 0: + wandb.watch( + self.model.policy, + log_freq=self.gradient_save_freq, + log=self.log, + ) + wandb.config.setdefaults(d) + + def _on_step(self) -> bool: + if self.model_save_freq > 0: + if self.model_save_path is not None: + if self.n_calls % self.model_save_freq == 0: + self.save_model() + return True + + def _on_training_end(self) -> None: + if self.model_save_path is not None: + self.save_model() + + def save_model(self) -> None: + self.model.save(self.path) + wandb.save(self.path, base_path=self.model_save_path) + if self.verbose > 1: + logger.info(f"Saving model checkpoint to {self.path}") diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/__init__.py b/lib/python3.12/site-packages/wandb/integration/sklearn/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fb3cb14c923c289c8075c296066f840c3c4e0858 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/sklearn/__init__.py @@ -0,0 +1,37 @@ +"""Create informative charts for scikit-learn models and log them to W&B.""" + +from .plot import ( + plot_calibration_curve, + plot_class_proportions, + plot_classifier, + plot_clusterer, + plot_confusion_matrix, + plot_elbow_curve, + plot_feature_importances, + plot_learning_curve, + plot_outlier_candidates, + plot_precision_recall, + plot_regressor, + plot_residuals, + plot_roc, + plot_silhouette, + plot_summary_metrics, +) + +__all__ = [ + "plot_classifier", + "plot_clusterer", + "plot_regressor", + "plot_summary_metrics", + "plot_learning_curve", + "plot_feature_importances", + "plot_class_proportions", + "plot_calibration_curve", + "plot_roc", + "plot_precision_recall", + "plot_confusion_matrix", + "plot_elbow_curve", + "plot_silhouette", + "plot_residuals", + "plot_outlier_candidates", +] diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/sklearn/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a7d7c4be9ffcbb90e36736cdc338e3500bc6258b Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/sklearn/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/sklearn/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..604967e9d083352ec00c50cf2d105767c8853173 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/sklearn/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__init__.py b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0d22d629bbc7dced6da390f885a4d327049491ff --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__init__.py @@ -0,0 +1,32 @@ +"""Calculates and formats metrics and charts for introspecting sklearn models. + +The functions in these modules are designed to be called by functions from the +plot submodule that have been exported into the namespace of the wandb.sklearn +submodule, rather than being called directly. +""" + +from .calibration_curves import calibration_curves +from .class_proportions import class_proportions +from .confusion_matrix import confusion_matrix +from .decision_boundaries import decision_boundaries +from .elbow_curve import elbow_curve +from .feature_importances import feature_importances +from .learning_curve import learning_curve +from .outlier_candidates import outlier_candidates +from .residuals import residuals +from .silhouette import silhouette +from .summary_metrics import summary_metrics + +__all__ = [ + "calibration_curves", + "class_proportions", + "confusion_matrix", + "decision_boundaries", + "elbow_curve", + "feature_importances", + "learning_curve", + "outlier_candidates", + "residuals", + "silhouette", + "summary_metrics", +] diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..98e75c423c93e2d5b25e5e97009ca557e3b27f3d Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/calibration_curves.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/calibration_curves.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d230db732f08a1c5dd8d05eb8ee42239ac6434e5 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/calibration_curves.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/class_proportions.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/class_proportions.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a90e3b96194f48297149499b68018909e6eabe93 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/class_proportions.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/confusion_matrix.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/confusion_matrix.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..804bc55999688fdc672aec3e70b83d9788b35884 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/confusion_matrix.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/decision_boundaries.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/decision_boundaries.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6940cf34f94077b8ca74f2f274370dd5c20e5653 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/decision_boundaries.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/elbow_curve.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/elbow_curve.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..be04054ad428ce27b88253d1f339c391bcf7ac22 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/elbow_curve.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/feature_importances.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/feature_importances.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..880382c07c2719cb51f8e40b7ab0bf65fbe8b4f9 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/feature_importances.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/learning_curve.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/learning_curve.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6d5a6a1e2edf9d739f2d7fc16527890aebaba119 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/learning_curve.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/outlier_candidates.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/outlier_candidates.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6219074557ea6b3efdae6e6844f5860530905a6c Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/outlier_candidates.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/residuals.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/residuals.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..953bd2de8dddb666d9f2da25d126d49a24e13702 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/residuals.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/silhouette.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/silhouette.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0ba4101cc0b3364cba22b83732eed242f94405bd Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/silhouette.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/summary_metrics.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/summary_metrics.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3e3858f6071c08848030ebfda0d3d3c33413df33 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/__pycache__/summary_metrics.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/calibration_curves.py b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/calibration_curves.py new file mode 100644 index 0000000000000000000000000000000000000000..b59aa25d7d8c984ca30529d497eb0e263570c365 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/calibration_curves.py @@ -0,0 +1,125 @@ +from warnings import simplefilter + +import numpy as np +import sklearn +from sklearn import model_selection, naive_bayes +from sklearn.calibration import CalibratedClassifierCV +from sklearn.linear_model import LogisticRegression + +import wandb +from wandb.integration.sklearn import utils + +# ignore all future warnings +simplefilter(action="ignore", category=FutureWarning) + + +def calibration_curves(clf, X, y, clf_name): # noqa: N803 + # ComplementNB (introduced in 0.20.0) requires non-negative features + if int(sklearn.__version__.split(".")[1]) >= 20 and isinstance( + clf, naive_bayes.ComplementNB + ): + X = X - X.min() # noqa:N806 + + # Calibrated with isotonic calibration + isotonic = CalibratedClassifierCV(clf, cv=2, method="isotonic") + + # Calibrated with sigmoid calibration + sigmoid = CalibratedClassifierCV(clf, cv=2, method="sigmoid") + + # Logistic regression with no calibration as baseline + lr = LogisticRegression(C=1.0) + + model_column = [] # color + frac_positives_column = [] # y axis + mean_pred_value_column = [] # x axis + hist_column = [] # barchart y + edge_column = [] # barchart x + + # Add curve for perfectly calibrated model + # format: model, fraction_of_positives, mean_predicted_value + model_column.append("Perfectly calibrated") + frac_positives_column.append(0) + mean_pred_value_column.append(0) + hist_column.append(0) + edge_column.append(0) + model_column.append("Perfectly calibrated") + hist_column.append(0) + edge_column.append(0) + frac_positives_column.append(1) + mean_pred_value_column.append(1) + + x_train, x_test, y_train, y_test = model_selection.train_test_split( + X, y, test_size=0.9, random_state=42 + ) + + # Add curve for LogisticRegression baseline and other models + + models = [lr, isotonic, sigmoid] + names = ["Logistic", f"{clf_name} Isotonic", f"{clf_name} Sigmoid"] + + for model, name in zip(models, names): + model.fit(x_train, y_train) + if hasattr(model, "predict_proba"): + prob_pos = model.predict_proba(x_test)[:, 1] + else: # use decision function + prob_pos = model.decision_function(x_test) + prob_pos = (prob_pos - prob_pos.min()) / (prob_pos.max() - prob_pos.min()) + + hist, edges = np.histogram(prob_pos, bins=10, density=False) + frac_positives, mean_pred_value = sklearn.calibration.calibration_curve( + y_test, prob_pos, n_bins=10 + ) + + # format: model, fraction_of_positives, mean_predicted_value + num_entries = len(frac_positives) + for i in range(num_entries): + hist_column.append(hist[i]) + edge_column.append(edges[i]) + model_column.append(name) + frac_positives_column.append(utils.round_3(frac_positives[i])) + mean_pred_value_column.append(utils.round_3(mean_pred_value[i])) + if utils.check_against_limit( + i, + "calibration_curve", + utils.chart_limit - 2, + ): + break + + table = make_table( + model_column, + frac_positives_column, + mean_pred_value_column, + hist_column, + edge_column, + ) + chart = wandb.visualize("wandb/calibration/v1", table) + + return chart + + +def make_table( + model_column, + frac_positives_column, + mean_pred_value_column, + hist_column, + edge_column, +): + columns = [ + "model", + "fraction_of_positives", + "mean_predicted_value", + "hist_dict", + "edge_dict", + ] + + data = list( + zip( + model_column, + frac_positives_column, + mean_pred_value_column, + hist_column, + edge_column, + ) + ) + + return wandb.Table(columns=columns, data=data) diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/class_proportions.py b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/class_proportions.py new file mode 100644 index 0000000000000000000000000000000000000000..183bf785a2fb62e2efe20fbdfde68f7a0fc45024 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/class_proportions.py @@ -0,0 +1,68 @@ +from warnings import simplefilter + +import numpy as np +from sklearn.utils.multiclass import unique_labels + +import wandb +from wandb.integration.sklearn import utils + +# ignore all future warnings +simplefilter(action="ignore", category=FutureWarning) + + +def class_proportions(y_train, y_test, labels): + # Get the unique values from the dataset + targets = (y_train,) if y_test is None else (y_train, y_test) + class_ids = np.array(unique_labels(*targets)) + + # Compute the class counts + counts_train = np.array([(y_train == c).sum() for c in class_ids]) + counts_test = np.array([(y_test == c).sum() for c in class_ids]) + + class_column, dataset_column, count_column = make_columns( + class_ids, counts_train, counts_test + ) + + if labels is not None and ( + isinstance(class_column[0], int) or isinstance(class_column[0], np.integer) + ): + class_column = get_named_labels(labels, class_column) + + table = make_table(class_column, dataset_column, count_column) + chart = wandb.visualize("wandb/class_proportions/v1", table) + + return chart + + +def make_table(class_column, dataset_column, count_column): + columns = ["class", "dataset", "count"] + data = list(zip(class_column, dataset_column, count_column)) + + return wandb.Table(data=data, columns=columns) + + +def make_columns(class_ids, counts_train, counts_test): + class_column, dataset_column, count_column = [], [], [] + + for i in range(len(class_ids)): + # add class counts from training set + class_column.append(class_ids[i]) + dataset_column.append("train") + count_column.append(counts_train[i]) + # add class counts from test set + class_column.append(class_ids[i]) + dataset_column.append("test") + count_column.append(counts_test[i]) + + if utils.check_against_limit( + i, + "class_proportions", + utils.chart_limit, + ): + break + + return class_column, dataset_column, count_column + + +def get_named_labels(labels, numeric_labels): + return np.array([labels[num_label] for num_label in numeric_labels]) diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/confusion_matrix.py b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/confusion_matrix.py new file mode 100644 index 0000000000000000000000000000000000000000..777c84ada92237222e0a058c5facd7b761313b12 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/confusion_matrix.py @@ -0,0 +1,93 @@ +import itertools +from warnings import simplefilter + +import numpy as np +from sklearn import metrics +from sklearn.utils.multiclass import unique_labels + +import wandb + +from .. import utils + +# ignore all future warnings +simplefilter(action="ignore", category=FutureWarning) + + +def validate_labels(*args, **kwargs): # FIXME + raise AssertionError() + + +def confusion_matrix( + y_true=None, + y_pred=None, + labels=None, + true_labels=None, + pred_labels=None, + normalize=False, +): + """Compute the confusion matrix to evaluate the performance of a classification. + + Called by plot_confusion_matrix to visualize roc curves. Please use the function + plot_confusion_matrix() if you wish to visualize your confusion matrix. + """ + cm = metrics.confusion_matrix(y_true, y_pred) + + if labels is None: + classes = unique_labels(y_true, y_pred) + else: + classes = np.asarray(labels) + + if normalize: + cm = cm.astype("float") / cm.sum(axis=1)[:, np.newaxis] + cm = np.around(cm, decimals=2) + cm[np.isnan(cm)] = 0.0 + + if true_labels is None: + true_classes = classes + else: + validate_labels(classes, true_labels, "true_labels") + + true_label_indexes = np.in1d(classes, true_labels) + + true_classes = classes[true_label_indexes] + cm = cm[true_label_indexes] + + if pred_labels is None: + pred_classes = classes + else: + validate_labels(classes, pred_labels, "pred_labels") + + pred_label_indexes = np.in1d(classes, pred_labels) + + pred_classes = classes[pred_label_indexes] + cm = cm[:, pred_label_indexes] + + table = make_table(cm, pred_classes, true_classes, labels) + chart = wandb.visualize("wandb/confusion_matrix/v1", table) + + return chart + + +def make_table(cm, pred_classes, true_classes, labels): + data, count = [], 0 + for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])): + if labels is not None and ( + isinstance(pred_classes[i], int) or isinstance(pred_classes[0], np.integer) + ): + pred = labels[pred_classes[i]] + true = labels[true_classes[j]] + else: + pred = pred_classes[i] + true = true_classes[j] + data.append([pred, true, cm[i, j]]) + count += 1 + if utils.check_against_limit( + count, + "confusion_matrix", + utils.chart_limit, + ): + break + + table = wandb.Table(columns=["Predicted", "Actual", "Count"], data=data) + + return table diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/decision_boundaries.py b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/decision_boundaries.py new file mode 100644 index 0000000000000000000000000000000000000000..a7a849b86b151270adf2437e4f2e72a03ce6c1a9 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/decision_boundaries.py @@ -0,0 +1,40 @@ +from warnings import simplefilter + +import wandb + +# ignore all future warnings +simplefilter(action="ignore", category=FutureWarning) + + +def decision_boundaries( + decision_boundary_x, + decision_boundary_y, + decision_boundary_color, + train_x, + train_y, + train_color, + test_x, + test_y, + test_color, +): + x_dict, y_dict, color_dict = [], [], [] + for i in range(min(len(decision_boundary_x), 100)): + x_dict.append(decision_boundary_x[i]) + y_dict.append(decision_boundary_y[i]) + color_dict.append(decision_boundary_color) + for i in range(300): + x_dict.append(test_x[i]) + y_dict.append(test_y[i]) + color_dict.append(test_color[i]) + for i in range(min(len(train_x), 600)): + x_dict.append(train_x[i]) + y_dict.append(train_y[i]) + color_dict.append(train_color[i]) + + return wandb.visualize( + "wandb/decision_boundaries/v1", + wandb.Table( + columns=["x", "y", "color"], + data=[[x_dict[i], y_dict[i], color_dict[i]] for i in range(len(x_dict))], + ), + ) diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/elbow_curve.py b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/elbow_curve.py new file mode 100644 index 0000000000000000000000000000000000000000..731102d75adabf238d44de00b8d79bca59f0a87a --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/elbow_curve.py @@ -0,0 +1,55 @@ +import time +from warnings import simplefilter + +import numpy as np +from joblib import Parallel, delayed +from sklearn.base import clone + +import wandb + +# ignore all future warnings +simplefilter(action="ignore", category=FutureWarning) + + +def elbow_curve(clusterer, X, cluster_ranges, n_jobs, show_cluster_time): # noqa: N803 + if cluster_ranges is None: + cluster_ranges = range(1, 10, 2) + else: + cluster_ranges = sorted(cluster_ranges) + + clfs, times = _compute_results_parallel(n_jobs, clusterer, X, cluster_ranges) + + clfs = np.absolute(clfs) + + table = make_table(cluster_ranges, clfs, times) + chart = wandb.visualize("wandb/elbow/v1", table) + + return chart + + +def make_table(cluster_ranges, clfs, times): + columns = ["cluster_ranges", "errors", "clustering_time"] + + data = list(zip(cluster_ranges, clfs, times)) + + table = wandb.Table(columns=columns, data=data) + + return table + + +def _compute_results_parallel(n_jobs, clusterer, x, cluster_ranges): + parallel_runner = Parallel(n_jobs=n_jobs) + _cluster_scorer = delayed(_clone_and_score_clusterer) + results = parallel_runner(_cluster_scorer(clusterer, x, i) for i in cluster_ranges) + + clfs, times = zip(*results) + + return clfs, times + + +def _clone_and_score_clusterer(clusterer, x, n_clusters): + start = time.time() + clusterer = clone(clusterer) + clusterer.n_clusters = n_clusters + + return clusterer.fit(x).score(x), time.time() - start diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/feature_importances.py b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/feature_importances.py new file mode 100644 index 0000000000000000000000000000000000000000..fac0452b944fd31d3bf611df054f4f6c0dbc0895 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/feature_importances.py @@ -0,0 +1,67 @@ +from warnings import simplefilter + +import numpy as np + +import wandb + +# ignore all future warnings +simplefilter(action="ignore", category=FutureWarning) + + +def feature_importances(model, feature_names): + attributes_to_check = ["feature_importances_", "feature_log_prob_", "coef_"] + found_attribute = check_for_attribute_on(model, attributes_to_check) + if found_attribute is None: + wandb.termwarn( + f"could not find any of attributes {', '.join(attributes_to_check)} on classifier. Cannot plot feature importances." + ) + return + elif found_attribute == "feature_importances_": + importances = model.feature_importances_ + elif found_attribute == "coef_": # ElasticNet-like models + importances = model.coef_ + elif found_attribute == "feature_log_prob_": + # coef_ was deprecated in sklearn 0.24, replaced with + # feature_log_prob_ + importances = model.feature_log_prob_ + + if len(importances.shape) > 1: + n_significant_dims = sum(i > 1 for i in importances.shape) + if n_significant_dims > 1: + nd = len(importances.shape) + wandb.termwarn( + f"{nd}-dimensional feature importances array passed to plot_feature_importances. " + f"{nd}-dimensional and higher feature importances arrays are not currently supported. " + f"These importances will not be plotted." + ) + return + else: + importances = np.squeeze(importances) + + indices = np.argsort(importances)[::-1] + importances = importances[indices] + + if feature_names is None: + feature_names = indices + else: + feature_names = np.array(feature_names)[indices] + + table = make_table(feature_names, importances) + chart = wandb.visualize("wandb/feature_importances/v1", table) + + return chart + + +def make_table(feature_names, importances): + table = wandb.Table( + columns=["feature_names", "importances"], + data=[[feature_names[i], importances[i]] for i in range(len(feature_names))], + ) + return table + + +def check_for_attribute_on(model, attributes_to_check): + for attr in attributes_to_check: + if hasattr(model, attr): + return attr + return None diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/learning_curve.py b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/learning_curve.py new file mode 100644 index 0000000000000000000000000000000000000000..296c7d7780fec21e5f9e3a7a346301c3c38b7de8 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/learning_curve.py @@ -0,0 +1,64 @@ +from warnings import simplefilter + +import numpy as np +from sklearn import model_selection + +import wandb +from wandb.integration.sklearn import utils + +# ignore all future warnings +simplefilter(action="ignore", category=FutureWarning) + + +def learning_curve( + model, + X, # noqa: N803 + y, + cv=None, + shuffle=False, + random_state=None, + train_sizes=None, + n_jobs=1, + scoring=None, +): + """Train model on datasets of varying size and generates plot of score vs size. + + Called by plot_learning_curve to visualize learning curve. Please use the function + plot_learning_curve() if you wish to visualize your learning curves. + """ + train_sizes, train_scores, test_scores = model_selection.learning_curve( + model, + X, + y, + cv=cv, + n_jobs=n_jobs, + train_sizes=train_sizes, + scoring=scoring, + shuffle=shuffle, + random_state=random_state, + ) + train_scores_mean = np.mean(train_scores, axis=1) + test_scores_mean = np.mean(test_scores, axis=1) + + table = make_table(train_scores_mean, test_scores_mean, train_sizes) + chart = wandb.visualize("wandb/learning_curve/v1", table) + + return chart + + +def make_table(train, test, train_sizes): + data = [] + for i in range(len(train)): + if utils.check_against_limit( + i, + "learning_curve", + utils.chart_limit / 2, + ): + break + train_set = ["train", utils.round_2(train[i]), train_sizes[i]] + test_set = ["test", utils.round_2(test[i]), train_sizes[i]] + data.append(train_set) + data.append(test_set) + + table = wandb.Table(columns=["dataset", "score", "train_size"], data=data) + return table diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/outlier_candidates.py b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/outlier_candidates.py new file mode 100644 index 0000000000000000000000000000000000000000..5fb0a29f4e7c49939e2d34420c02f9d82d08570b --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/outlier_candidates.py @@ -0,0 +1,69 @@ +from warnings import simplefilter + +import numpy as np + +import wandb +from wandb.integration.sklearn import utils + +# ignore all future warnings +simplefilter(action="ignore", category=FutureWarning) + + +def outlier_candidates(regressor, X, y): # noqa: N803 + # Fit a linear model to X and y to compute MSE + regressor.fit(X, y) + + # Leverage is computed as the diagonal of the projection matrix of X + leverage = (X * np.linalg.pinv(X).T).sum(1) + + # Compute the rank and the degrees of freedom of the OLS model + rank = np.linalg.matrix_rank(X) + df = X.shape[0] - rank + + # Compute the MSE from the residuals + residuals = y - regressor.predict(X) + mse = np.dot(residuals, residuals) / df + + # Compute Cook's distance + residuals_studentized = residuals / np.sqrt(mse) / np.sqrt(1 - leverage) + distance_ = residuals_studentized**2 / X.shape[1] + distance_ *= leverage / (1 - leverage) + + # Compute the influence threshold rule of thumb + influence_threshold_ = 4 / X.shape[0] + outlier_percentage_ = sum(distance_ >= influence_threshold_) / X.shape[0] + outlier_percentage_ *= 100.0 + + distance_dict, count = [], 0 + for d in distance_: + distance_dict.append(d) + count += 1 + if utils.check_against_limit( + count, + "outlier_candidates", + utils.chart_limit, + ): + break + + table = make_table(distance_dict, outlier_percentage_, influence_threshold_) + chart = wandb.visualize("wandb/outliers/v1", table) + + return chart + + +def make_table(distance, outlier_percentage, influence_threshold): + columns = [ + "distance", + "instance_indicies", + "outlier_percentage", + "influence_threshold", + ] + + data = [ + [distance[i], i, utils.round_3(outlier_percentage), influence_threshold] + for i in range(len(distance)) + ] + + table = wandb.Table(columns=columns, data=data) + + return table diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/residuals.py b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/residuals.py new file mode 100644 index 0000000000000000000000000000000000000000..b45df7b84fa5b3824f99425800bbf0039f7cacd8 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/residuals.py @@ -0,0 +1,86 @@ +from warnings import simplefilter + +from sklearn import model_selection + +import wandb +from wandb.integration.sklearn import utils + +# ignore all future warnings +simplefilter(action="ignore", category=FutureWarning) + + +def residuals(regressor, X, y): # noqa: N803 + # Create the train and test splits + x_train, x_test, y_train, y_test = model_selection.train_test_split( + X, y, test_size=0.2 + ) + + # Store labels and colors for the legend ordered by call + regressor.fit(x_train, y_train) + train_score_ = regressor.score(x_train, y_train) + test_score_ = regressor.score(x_test, y_test) + + y_pred_train = regressor.predict(x_train) + residuals_train = y_pred_train - y_train + + y_pred_test = regressor.predict(x_test) + residuals_test = y_pred_test - y_test + + table = make_table( + y_pred_train, + residuals_train, + y_pred_test, + residuals_test, + train_score_, + test_score_, + ) + chart = wandb.visualize("wandb/residuals_plot/v1", table) + + return chart + + +def make_table( + y_pred_train, + residuals_train, + y_pred_test, + residuals_test, + train_score_, + test_score_, +): + y_pred_column, dataset_column, residuals_column = [], [], [] + + datapoints, max_datapoints_train = 0, 100 + for pred, residual in zip(y_pred_train, residuals_train): + # add class counts from training set + y_pred_column.append(pred) + dataset_column.append("train") + residuals_column.append(residual) + datapoints += 1 + if utils.check_against_limit(datapoints, "residuals", max_datapoints_train): + break + + datapoints = 0 + for pred, residual in zip(y_pred_test, residuals_test): + # add class counts from training set + y_pred_column.append(pred) + dataset_column.append("test") + residuals_column.append(residual) + datapoints += 1 + if utils.check_against_limit(datapoints, "residuals", max_datapoints_train): + break + + columns = ["dataset", "y_pred", "residuals", "train_score", "test_score"] + data = [ + [ + dataset_column[i], + y_pred_column[i], + residuals_column[i], + train_score_, + test_score_, + ] + for i in range(len(y_pred_column)) + ] + + table = wandb.Table(columns=columns, data=data) + + return table diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/silhouette.py b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/silhouette.py new file mode 100644 index 0000000000000000000000000000000000000000..a71ba88c8f03d99e6908e4f886ba5d0be4f4dd41 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/silhouette.py @@ -0,0 +1,118 @@ +from warnings import simplefilter + +import numpy as np +from sklearn.metrics import silhouette_samples, silhouette_score +from sklearn.preprocessing import LabelEncoder + +import wandb +from wandb.integration.sklearn import utils + +# ignore all future warnings +simplefilter(action="ignore", category=FutureWarning) + + +def silhouette(clusterer, X, cluster_labels, labels, metric, kmeans): # noqa: N803 + # Run clusterer for n_clusters in range(len(cluster_ranges), get cluster labels + # TODO - keep/delete once we decide if we should train clusterers + # or ask for trained models + # clusterer.set_params(n_clusters=n_clusters, random_state=42) + # cluster_labels = clusterer.fit_predict(X) + cluster_labels = np.asarray(cluster_labels) + labels = np.asarray(labels) + + le = LabelEncoder() + _ = le.fit_transform(cluster_labels) + n_clusters = len(np.unique(cluster_labels)) + + # The silhouette_score gives the average value for all the samples. + # This gives a perspective into the density and separation of the formed + # clusters + silhouette_avg = silhouette_score(X, cluster_labels, metric=metric) + + # Compute the silhouette scores for each sample + sample_silhouette_values = silhouette_samples(X, cluster_labels, metric=metric) + + x_sil, y_sil, color_sil = [], [], [] + + count, y_lower = 0, 10 + for i in range(n_clusters): + # Aggregate the silhouette scores for samples belonging to + # cluster i, and sort them + ith_cluster_silhouette_values = sample_silhouette_values[cluster_labels == i] + + ith_cluster_silhouette_values.sort() + + size_cluster_i = ith_cluster_silhouette_values.shape[0] + y_upper = y_lower + size_cluster_i + + y_values = np.arange(y_lower, y_upper) + + for j in range(len(y_values)): + y_sil.append(y_values[j]) + x_sil.append(ith_cluster_silhouette_values[j]) + color_sil.append(i) + count += 1 + if utils.check_against_limit(count, "silhouette", utils.chart_limit): + break + + # Compute the new y_lower for next plot + y_lower = y_upper + 10 # 10 for the 0 samples + + if kmeans: + centers = clusterer.cluster_centers_ + centerx = centers[:, 0] + centery = centers[:, 1] + + else: + centerx = [None] * len(color_sil) + centery = [None] * len(color_sil) + + table = make_table( + X[:, 0], + X[:, 1], + cluster_labels, + centerx, + centery, + y_sil, + x_sil, + color_sil, + silhouette_avg, + ) + chart = wandb.visualize("wandb/silhouette_/v1", table) + + return chart + + +def make_table(x, y, colors, centerx, centery, y_sil, x_sil, color_sil, silhouette_avg): + columns = [ + "x", + "y", + "colors", + "centerx", + "centery", + "y_sil", + "x1", + "x2", + "color_sil", + "silhouette_avg", + ] + + data = [ + [ + x[i], + y[i], + colors[i], + centerx[colors[i]], + centery[colors[i]], + y_sil[i], + 0, + x_sil[i], + color_sil[i], + silhouette_avg, + ] + for i in range(len(color_sil)) + ] + + table = wandb.Table(data=data, columns=columns) + + return table diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/summary_metrics.py b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/summary_metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..e4b6f25ead2fd3c94c302db7693ca3fd529812d9 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/sklearn/calculate/summary_metrics.py @@ -0,0 +1,62 @@ +from warnings import simplefilter + +import numpy as np +import sklearn + +import wandb +from wandb.integration.sklearn import utils + +# ignore all future warnings +simplefilter(action="ignore", category=FutureWarning) + + +def summary_metrics(model=None, X=None, y=None, X_test=None, y_test=None): # noqa: N803 + """Calculate summary metrics for both regressors and classifiers. + + Called by plot_summary_metrics to visualize metrics. Please use the function + plot_summary_metrics() if you wish to visualize your summary metrics. + """ + y, y_test = np.asarray(y), np.asarray(y_test) + metrics = {} + model_name = model.__class__.__name__ + + y_pred = model.predict(X_test) + + if sklearn.base.is_classifier(model): + accuracy_score = sklearn.metrics.accuracy_score(y_test, y_pred) + metrics["accuracy_score"] = accuracy_score + + precision = sklearn.metrics.precision_score(y_test, y_pred, average="weighted") + metrics["precision"] = precision + + recall = sklearn.metrics.recall_score(y_test, y_pred, average="weighted") + metrics["recall"] = recall + + f1_score = sklearn.metrics.f1_score(y_test, y_pred, average="weighted") + metrics["f1_score"] = f1_score + + elif sklearn.base.is_regressor(model): + mae = sklearn.metrics.mean_absolute_error(y_test, y_pred) + metrics["mae"] = mae + + mse = sklearn.metrics.mean_squared_error(y_test, y_pred) + metrics["mse"] = mse + + r2_score = sklearn.metrics.r2_score(y_test, y_pred) + metrics["r2_score"] = r2_score + + metrics = {name: utils.round_2(metric) for name, metric in metrics.items()} + + table = make_table(metrics, model_name) + chart = wandb.visualize("wandb/metrics/v1", table) + + return chart + + +def make_table(metrics, model_name): + columns = ["metric_name", "metric_value", "model_name"] + table_content = [[name, value, model_name] for name, value in metrics.items()] + + table = wandb.Table(columns=columns, data=table_content) + + return table diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/plot/__init__.py b/lib/python3.12/site-packages/wandb/integration/sklearn/plot/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..710e24c010e108bb4f25d361092a55182ad26917 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/sklearn/plot/__init__.py @@ -0,0 +1,35 @@ +"""Create and logs charts introspecting models built with scikit-learn to W&B.""" + +from .classifier import calibration_curve as plot_calibration_curve +from .classifier import class_proportions as plot_class_proportions +from .classifier import classifier as plot_classifier +from .classifier import confusion_matrix as plot_confusion_matrix +from .classifier import feature_importances as plot_feature_importances +from .classifier import precision_recall as plot_precision_recall +from .classifier import roc as plot_roc +from .clusterer import clusterer as plot_clusterer +from .clusterer import elbow_curve as plot_elbow_curve +from .clusterer import silhouette as plot_silhouette +from .regressor import outlier_candidates as plot_outlier_candidates +from .regressor import regressor as plot_regressor +from .regressor import residuals as plot_residuals +from .shared import learning_curve as plot_learning_curve +from .shared import summary_metrics as plot_summary_metrics + +__all__ = [ + "plot_classifier", + "plot_clusterer", + "plot_regressor", + "plot_summary_metrics", + "plot_learning_curve", + "plot_feature_importances", + "plot_class_proportions", + "plot_calibration_curve", + "plot_roc", + "plot_precision_recall", + "plot_confusion_matrix", + "plot_elbow_curve", + "plot_silhouette", + "plot_residuals", + "plot_outlier_candidates", +] diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/plot/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/sklearn/plot/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..83d59c99a6b8ef855c6afd277aa95ddb2db8ac9e Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/sklearn/plot/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/plot/__pycache__/classifier.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/sklearn/plot/__pycache__/classifier.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7e7d133da9435f2cab120286d5c55ec99f165329 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/sklearn/plot/__pycache__/classifier.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/plot/__pycache__/clusterer.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/sklearn/plot/__pycache__/clusterer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e41ff5cdd7bed58c9ba739a4fefc6d4183861e3d Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/sklearn/plot/__pycache__/clusterer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/plot/__pycache__/regressor.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/sklearn/plot/__pycache__/regressor.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..669ff276c34a0bbd06955d715d65584aa68ecdbf Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/sklearn/plot/__pycache__/regressor.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/plot/__pycache__/shared.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/sklearn/plot/__pycache__/shared.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..00e2d567a7689769c692291faa751204c399fe95 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/sklearn/plot/__pycache__/shared.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/plot/classifier.py b/lib/python3.12/site-packages/wandb/integration/sklearn/plot/classifier.py new file mode 100644 index 0000000000000000000000000000000000000000..e431c8090b565518315f7906c85314f7d4b0ad9b --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/sklearn/plot/classifier.py @@ -0,0 +1,329 @@ +"""Define plots for classification models built with scikit-learn.""" + +from warnings import simplefilter + +import numpy as np +from sklearn import naive_bayes + +import wandb +import wandb.plot +from wandb.integration.sklearn import calculate, utils + +from . import shared + +# ignore all future warnings +simplefilter(action="ignore", category=FutureWarning) + + +def classifier( + model, + X_train, # noqa: N803 + X_test, # noqa: N803 + y_train, + y_test, + y_pred, + y_probas, + labels, + is_binary=False, + model_name="Classifier", + feature_names=None, + log_learning_curve=False, +): + """Generate all sklearn classifier plots supported by W&B. + + The following plots are generated: + feature importances, confusion matrix, summary metrics, + class proportions, calibration curve, roc curve, precision-recall curve. + + Should only be called with a fitted classifier (otherwise an error is thrown). + + Args: + model: (classifier) Takes in a fitted classifier. + X_train: (arr) Training set features. + y_train: (arr) Training set labels. + X_test: (arr) Test set features. + y_test: (arr) Test set labels. + y_pred: (arr) Test set predictions by the model passed. + y_probas: (arr) Test set predicted probabilities by the model passed. + labels: (list) Named labels for target variable (y). Makes plots easier to + read by replacing target values with corresponding index. + For example if `labels=['dog', 'cat', 'owl']` all 0s are + replaced by dog, 1s by cat. + is_binary: (bool) Is the model passed a binary classifier? Defaults to False + model_name: (str) Model name. Defaults to 'Classifier' + feature_names: (list) Names for features. Makes plots easier to read by + replacing feature indexes with corresponding names. + log_learning_curve: (bool) Whether or not to log the learning curve. + Defaults to False. + + Returns: + None: To see plots, go to your W&B run page then expand the 'media' tab + under 'auto visualizations'. + + Example: + ```python + wandb.sklearn.plot_classifier( + model, + X_train, + X_test, + y_train, + y_test, + y_pred, + y_probas, + ["cat", "dog"], + False, + "RandomForest", + ["barks", "drools", "plays_fetch", "breed"], + ) + ``` + """ + wandb.termlog(f"\nPlotting {model_name}.") + + if not isinstance(model, naive_bayes.MultinomialNB): + feature_importances(model, feature_names) + wandb.termlog("Logged feature importances.") + + if log_learning_curve: + shared.learning_curve(model, X_train, y_train) + wandb.termlog("Logged learning curve.") + + confusion_matrix(y_test, y_pred, labels) + wandb.termlog("Logged confusion matrix.") + + shared.summary_metrics(model, X=X_train, y=y_train, X_test=X_test, y_test=y_test) + wandb.termlog("Logged summary metrics.") + + class_proportions(y_train, y_test, labels) + wandb.termlog("Logged class proportions.") + + if not isinstance(model, naive_bayes.MultinomialNB): + calibration_curve(model, X_train, y_train, model_name) + wandb.termlog("Logged calibration curve.") + + roc(y_test, y_probas, labels) + wandb.termlog("Logged roc curve.") + + precision_recall(y_test, y_probas, labels) + wandb.termlog("Logged precision-recall curve.") + + +def roc( + y_true=None, + y_probas=None, + labels=None, + plot_micro=True, + plot_macro=True, + classes_to_plot=None, +): + """Log the receiver-operating characteristic curve. + + Args: + y_true: (arr) Test set labels. + y_probas: (arr) Test set predicted probabilities. + labels: (list) Named labels for target variable (y). Makes plots easier to + read by replacing target values with corresponding index. + For example if `labels=['dog', 'cat', 'owl']` all 0s are + replaced by dog, 1s by cat. + + Returns: + None: To see plots, go to your W&B run page then expand the 'media' tab + under 'auto visualizations'. + + Example: + ```python + wandb.sklearn.plot_roc(y_true, y_probas, labels) + ``` + """ + roc_chart = wandb.plot.roc_curve(y_true, y_probas, labels, classes_to_plot) + wandb.log({"roc": roc_chart}) + + +def confusion_matrix( + y_true=None, + y_pred=None, + labels=None, + true_labels=None, + pred_labels=None, + normalize=False, +): + """Log a confusion matrix to W&B. + + Confusion matrices depict the pattern of misclassifications by a model. + + Args: + y_true: (arr) Test set labels. + y_probas: (arr) Test set predicted probabilities. + labels: (list) Named labels for target variable (y). Makes plots easier to + read by replacing target values with corresponding index. + For example if `labels=['dog', 'cat', 'owl']` all 0s are + replaced by dog, 1s by cat. + + Returns: + None: To see plots, go to your W&B run page then expand the 'media' tab + under 'auto visualizations'. + + Example: + ```python + wandb.sklearn.plot_confusion_matrix(y_true, y_probas, labels) + ``` + """ + y_true = np.asarray(y_true) + y_pred = np.asarray(y_pred) + + not_missing = utils.test_missing(y_true=y_true, y_pred=y_pred) + correct_types = utils.test_types(y_true=y_true, y_pred=y_pred) + + if not_missing and correct_types: + confusion_matrix_chart = calculate.confusion_matrix( + y_true, + y_pred, + labels, + true_labels, + pred_labels, + normalize, + ) + + wandb.log({"confusion_matrix": confusion_matrix_chart}) + + +def precision_recall( + y_true=None, y_probas=None, labels=None, plot_micro=True, classes_to_plot=None +): + """Log a precision-recall curve to W&B. + + Precision-recall curves depict the tradeoff between positive predictive value (precision) + and true positive rate (recall) as the threshold of a classifier is shifted. + + Args: + y_true: (arr) Test set labels. + y_probas: (arr) Test set predicted probabilities. + labels: (list) Named labels for target variable (y). Makes plots easier to + read by replacing target values with corresponding index. + For example if `labels=['dog', 'cat', 'owl']` all 0s are + replaced by dog, 1s by cat. + + Returns: + None: To see plots, go to your W&B run page then expand the 'media' tab + under 'auto visualizations'. + + Example: + ```python + wandb.sklearn.plot_precision_recall(y_true, y_probas, labels) + ``` + """ + precision_recall_chart = wandb.plot.pr_curve( + y_true, y_probas, labels, classes_to_plot + ) + + wandb.log({"precision_recall": precision_recall_chart}) + + +def feature_importances( + model=None, feature_names=None, title="Feature Importance", max_num_features=50 +): + """Log a plot depicting the relative importance of each feature for a classifier's decisions. + + Should only be called with a fitted classifier (otherwise an error is thrown). + Only works with classifiers that have a feature_importances_ attribute, like trees. + + Args: + model: (clf) Takes in a fitted classifier. + feature_names: (list) Names for features. Makes plots easier to read by + replacing feature indexes with corresponding names. + + Returns: + None: To see plots, go to your W&B run page then expand the 'media' tab + under 'auto visualizations'. + + Example: + ```python + wandb.sklearn.plot_feature_importances(model, ["width", "height", "length"]) + ``` + """ + not_missing = utils.test_missing(model=model) + correct_types = utils.test_types(model=model) + model_fitted = utils.test_fitted(model) + + if not_missing and correct_types and model_fitted: + feature_importance_chart = calculate.feature_importances(model, feature_names) + wandb.log({"feature_importances": feature_importance_chart}) + + +def class_proportions(y_train=None, y_test=None, labels=None): + """Plot the distribution of target classes in training and test sets. + + Useful for detecting imbalanced classes. + + Args: + y_train: (arr) Training set labels. + y_test: (arr) Test set labels. + labels: (list) Named labels for target variable (y). Makes plots easier to + read by replacing target values with corresponding index. + For example if `labels=['dog', 'cat', 'owl']` all 0s are + replaced by dog, 1s by cat. + + Returns: + None: To see plots, go to your W&B run page then expand the 'media' tab + under 'auto visualizations'. + + Example: + ```python + wandb.sklearn.plot_class_proportions(y_train, y_test, ["dog", "cat", "owl"]) + ``` + """ + not_missing = utils.test_missing(y_train=y_train, y_test=y_test) + correct_types = utils.test_types(y_train=y_train, y_test=y_test) + if not_missing and correct_types: + y_train, y_test = np.array(y_train), np.array(y_test) + class_proportions_chart = calculate.class_proportions(y_train, y_test, labels) + + wandb.log({"class_proportions": class_proportions_chart}) + + +def calibration_curve(clf=None, X=None, y=None, clf_name="Classifier"): # noqa: N803 + """Log a plot depicting how well-calibrated the predicted probabilities of a classifier are. + + Also suggests how to calibrate an uncalibrated classifier. Compares estimated predicted + probabilities by a baseline logistic regression model, the model passed as + an argument, and by both its isotonic calibration and sigmoid calibrations. + The closer the calibration curves are to a diagonal the better. + A sine wave like curve represents an overfitted classifier, while a cosine + wave like curve represents an underfitted classifier. + By training isotonic and sigmoid calibrations of the model and comparing + their curves we can figure out whether the model is over or underfitting and + if so which calibration (sigmoid or isotonic) might help fix this. + For more details, see https://scikit-learn.org/stable/auto_examples/calibration/plot_calibration_curve.html. + + Should only be called with a fitted classifier (otherwise an error is thrown). + + Please note this function fits variations of the model on the training set when called. + + Args: + clf: (clf) Takes in a fitted classifier. + X: (arr) Training set features. + y: (arr) Training set labels. + model_name: (str) Model name. Defaults to 'Classifier' + + Returns: + None: To see plots, go to your W&B run page then expand the 'media' tab + under 'auto visualizations'. + + Example: + ```python + wandb.sklearn.plot_calibration_curve(clf, X, y, "RandomForestClassifier") + ``` + """ + not_missing = utils.test_missing(clf=clf, X=X, y=y) + correct_types = utils.test_types(clf=clf, X=X, y=y) + is_fitted = utils.test_fitted(clf) + if not_missing and correct_types and is_fitted: + y = np.asarray(y) + if y.dtype.char == "U" or not ((y == 0) | (y == 1)).all(): + wandb.termwarn( + "This function only supports binary classification at the moment and therefore expects labels to be binary. Skipping calibration curve." + ) + return + + calibration_curve_chart = calculate.calibration_curves(clf, X, y, clf_name) + + wandb.log({"calibration_curve": calibration_curve_chart}) diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/plot/clusterer.py b/lib/python3.12/site-packages/wandb/integration/sklearn/plot/clusterer.py new file mode 100644 index 0000000000000000000000000000000000000000..bced65ae105780961e434b78e6ef316297b5182e --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/sklearn/plot/clusterer.py @@ -0,0 +1,146 @@ +"""Define plots for clustering models built with scikit-learn.""" + +from warnings import simplefilter + +import pandas as pd +import sklearn + +import wandb +from wandb.integration.sklearn import calculate, utils + +# ignore all future warnings +simplefilter(action="ignore", category=FutureWarning) + + +def clusterer(model, X_train, cluster_labels, labels=None, model_name="Clusterer"): # noqa: N803 + """Generates all sklearn clusterer plots supported by W&B. + + The following plots are generated: + elbow curve, silhouette plot. + + Should only be called with a fitted clusterer (otherwise an error is thrown). + + Args: + model: (clusterer) Takes in a fitted clusterer. + X_train: (arr) Training set features. + cluster_labels: (list) Names for cluster labels. Makes plots easier to read + by replacing cluster indexes with corresponding names. + labels: (list) Named labels for target variable (y). Makes plots easier to + read by replacing target values with corresponding index. + For example if `labels=['dog', 'cat', 'owl']` all 0s are + replaced by dog, 1s by cat. + model_name: (str) Model name. Defaults to 'Clusterer' + + Returns: + None: To see plots, go to your W&B run page then expand the 'media' tab + under 'auto visualizations'. + + Example: + ```python + wandb.sklearn.plot_clusterer(kmeans, X, cluster_labels, labels, "KMeans") + ``` + """ + wandb.termlog(f"\nPlotting {model_name}.") + if isinstance(model, sklearn.cluster.KMeans): + elbow_curve(model, X_train) + wandb.termlog("Logged elbow curve.") + + silhouette(model, X_train, cluster_labels, labels=labels, kmeans=True) + + else: + silhouette(model, X_train, cluster_labels, kmeans=False) + + wandb.termlog("Logged silhouette plot.") + + +def elbow_curve( + clusterer=None, + X=None, # noqa: N803 + cluster_ranges=None, + n_jobs=1, + show_cluster_time=True, +): + """Measures and plots variance explained as a function of the number of clusters. + + Useful in picking the optimal number of clusters. + + Should only be called with a fitted clusterer (otherwise an error is thrown). + + Please note this function fits the model on the training set when called. + + Args: + model: (clusterer) Takes in a fitted clusterer. + X: (arr) Training set features. + + Returns: + None: To see plots, go to your W&B run page then expand the 'media' tab + under 'auto visualizations'. + + Example: + ```python + wandb.sklearn.plot_elbow_curve(model, X_train) + ``` + """ + if not hasattr(clusterer, "n_clusters"): + wandb.termlog( + "n_clusters attribute not in classifier. Cannot plot elbow method." + ) + return + + not_missing = utils.test_missing(clusterer=clusterer) + correct_types = utils.test_types + is_fitted = utils.test_fitted(clusterer) + + if not_missing and correct_types and is_fitted: + elbow_curve_chart = calculate.elbow_curve( + clusterer, X, cluster_ranges, n_jobs, show_cluster_time + ) + + wandb.log({"elbow_curve": elbow_curve_chart}) + + +def silhouette( + clusterer=None, + X=None, # noqa: N803 + cluster_labels=None, + labels=None, + metric="euclidean", + kmeans=True, +): + """Measures & plots silhouette coefficients. + + Silhouette coefficients near +1 indicate that the sample is far away from + the neighboring clusters. A value near 0 indicates that the sample is on or + very close to the decision boundary between two neighboring clusters and + negative values indicate that the samples might have been assigned to the wrong cluster. + + Should only be called with a fitted clusterer (otherwise an error is thrown). + + Please note this function fits the model on the training set when called. + + Args: + model: (clusterer) Takes in a fitted clusterer. + X: (arr) Training set features. + cluster_labels: (list) Names for cluster labels. Makes plots easier to read + by replacing cluster indexes with corresponding names. + + Returns: + None: To see plots, go to your W&B run page then expand the 'media' tab + under 'auto visualizations'. + + Example: + ```python + wandb.sklearn.plot_silhouette(model, X_train, ["spam", "not spam"]) + ``` + """ + not_missing = utils.test_missing(clusterer=clusterer) + correct_types = utils.test_types(clusterer=clusterer) + is_fitted = utils.test_fitted(clusterer) + + if not_missing and correct_types and is_fitted: + if isinstance(X, (pd.DataFrame)): + X = X.values # noqa: N806 + silhouette_chart = calculate.silhouette( + clusterer, X, cluster_labels, labels, metric, kmeans + ) + wandb.log({"silhouette_plot": silhouette_chart}) diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/plot/regressor.py b/lib/python3.12/site-packages/wandb/integration/sklearn/plot/regressor.py new file mode 100644 index 0000000000000000000000000000000000000000..a2840a06dabfd27cbe0d980af2988f6cc1ef0604 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/sklearn/plot/regressor.py @@ -0,0 +1,121 @@ +"""Define plots for regression models built with scikit-learn.""" + +from warnings import simplefilter + +import numpy as np + +import wandb +from wandb.integration.sklearn import calculate, utils + +from . import shared + +# ignore all future warnings +simplefilter(action="ignore", category=FutureWarning) + + +def regressor(model, X_train, X_test, y_train, y_test, model_name="Regressor"): # noqa: N803 + """Generates all sklearn regressor plots supported by W&B. + + The following plots are generated: + learning curve, summary metrics, residuals plot, outlier candidates. + + Should only be called with a fitted regressor (otherwise an error is thrown). + + Args: + model: (regressor) Takes in a fitted regressor. + X_train: (arr) Training set features. + y_train: (arr) Training set labels. + X_test: (arr) Test set features. + y_test: (arr) Test set labels. + model_name: (str) Model name. Defaults to 'Regressor' + + Returns: + None: To see plots, go to your W&B run page then expand the 'media' tab + under 'auto visualizations'. + + Example: + ```python + wandb.sklearn.plot_regressor(reg, X_train, X_test, y_train, y_test, "Ridge") + ``` + """ + wandb.termlog(f"\nPlotting {model_name}.") + + shared.summary_metrics(model, X_train, y_train, X_test, y_test) + wandb.termlog("Logged summary metrics.") + + shared.learning_curve(model, X_train, y_train) + wandb.termlog("Logged learning curve.") + + outlier_candidates(model, X_train, y_train) + wandb.termlog("Logged outlier candidates.") + + residuals(model, X_train, y_train) + wandb.termlog("Logged residuals.") + + +def outlier_candidates(regressor=None, X=None, y=None): # noqa: N803 + """Measures a datapoint's influence on regression model via cook's distance. + + Instances with high influences could potentially be outliers. + + Should only be called with a fitted regressor (otherwise an error is thrown). + + Please note this function fits the model on the training set when called. + + Args: + model: (regressor) Takes in a fitted regressor. + X: (arr) Training set features. + y: (arr) Training set labels. + + Returns: + None: To see plots, go to your W&B run page then expand the 'media' tab + under 'auto visualizations'. + + Example: + ```python + wandb.sklearn.plot_outlier_candidates(model, X, y) + ``` + """ + is_missing = utils.test_missing(regressor=regressor, X=X, y=y) + correct_types = utils.test_types(regressor=regressor, X=X, y=y) + is_fitted = utils.test_fitted(regressor) + + if is_missing and correct_types and is_fitted: + y = np.asarray(y) + + outliers_chart = calculate.outlier_candidates(regressor, X, y) + wandb.log({"outlier_candidates": outliers_chart}) + + +def residuals(regressor=None, X=None, y=None): # noqa: N803 + """Measures and plots the regressor's predicted value against the residual. + + The marginal distribution of residuals is also calculated and plotted. + + Should only be called with a fitted regressor (otherwise an error is thrown). + + Please note this function fits variations of the model on the training set when called. + + Args: + regressor: (regressor) Takes in a fitted regressor. + X: (arr) Training set features. + y: (arr) Training set labels. + + Returns: + None: To see plots, go to your W&B run page then expand the 'media' tab + under 'auto visualizations'. + + Example: + ```python + wandb.sklearn.plot_residuals(model, X, y) + ``` + """ + not_missing = utils.test_missing(regressor=regressor, X=X, y=y) + correct_types = utils.test_types(regressor=regressor, X=X, y=y) + is_fitted = utils.test_fitted(regressor) + + if not_missing and correct_types and is_fitted: + y = np.asarray(y) + + residuals_chart = calculate.residuals(regressor, X, y) + wandb.log({"residuals": residuals_chart}) diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/plot/shared.py b/lib/python3.12/site-packages/wandb/integration/sklearn/plot/shared.py new file mode 100644 index 0000000000000000000000000000000000000000..871dbd742c960d681803f240fbccddff02000733 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/sklearn/plot/shared.py @@ -0,0 +1,91 @@ +"""Define plots used by multiple sklearn model classes.""" + +from warnings import simplefilter + +import numpy as np + +import wandb +from wandb.integration.sklearn import calculate, utils + +# ignore all future warnings +simplefilter(action="ignore", category=FutureWarning) + + +def summary_metrics(model=None, X=None, y=None, X_test=None, y_test=None): # noqa: N803 + """Logs a chart depicting summary metrics for a model. + + Should only be called with a fitted model (otherwise an error is thrown). + + Args: + model: (clf or reg) Takes in a fitted regressor or classifier. + X: (arr) Training set features. + y: (arr) Training set labels. + X_test: (arr) Test set features. + y_test: (arr) Test set labels. + + Returns: + None: To see plots, go to your W&B run page then expand the 'media' tab + under 'auto visualizations'. + + Example: + ```python + wandb.sklearn.plot_summary_metrics(model, X_train, y_train, X_test, y_test) + ``` + """ + not_missing = utils.test_missing( + model=model, X=X, y=y, X_test=X_test, y_test=y_test + ) + correct_types = utils.test_types( + model=model, X=X, y=y, X_test=X_test, y_test=y_test + ) + model_fitted = utils.test_fitted(model) + + if not_missing and correct_types and model_fitted: + metrics_chart = calculate.summary_metrics(model, X, y, X_test, y_test) + wandb.log({"summary_metrics": metrics_chart}) + + +def learning_curve( + model=None, + X=None, # noqa: N803 + y=None, + cv=None, + shuffle=False, + random_state=None, + train_sizes=None, + n_jobs=1, + scoring=None, +): + """Logs a plot depicting model performance against dataset size. + + Please note this function fits the model to datasets of varying sizes when called. + + Args: + model: (clf or reg) Takes in a fitted regressor or classifier. + X: (arr) Dataset features. + y: (arr) Dataset labels. + + For details on the other keyword arguments, see the documentation for + `sklearn.model_selection.learning_curve`. + + Returns: + None: To see plots, go to your W&B run page then expand the 'media' tab + under 'auto visualizations'. + + Example: + ```python + wandb.sklearn.plot_learning_curve(model, X, y) + ``` + """ + not_missing = utils.test_missing(model=model, X=X, y=y) + correct_types = utils.test_types(model=model, X=X, y=y) + if not_missing and correct_types: + if train_sizes is None: + train_sizes = np.linspace(0.1, 1.0, 5) + y = np.asarray(y) + + learning_curve_chart = calculate.learning_curve( + model, X, y, cv, shuffle, random_state, train_sizes, n_jobs, scoring + ) + + wandb.log({"learning_curve": learning_curve_chart}) diff --git a/lib/python3.12/site-packages/wandb/integration/sklearn/utils.py b/lib/python3.12/site-packages/wandb/integration/sklearn/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0b2ef628bbb4165e15d33cd26feb144bd2e40d11 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/sklearn/utils.py @@ -0,0 +1,184 @@ +"""Shared utilities for the modules in wandb.sklearn.""" + +from collections.abc import Iterable, Sequence + +import numpy as np +import pandas as pd +import scipy +import sklearn + +import wandb + +chart_limit = 1000 + + +def check_against_limit(count, chart, limit=None): + if limit is None: + limit = chart_limit + if count > limit: + warn_chart_limit(limit, chart) + return True + else: + return False + + +def warn_chart_limit(limit, chart): + warning = f"using only the first {limit} datapoints to create chart {chart}" + wandb.termwarn(warning) + + +def encode_labels(df): + le = sklearn.preprocessing.LabelEncoder() + # apply le on categorical feature columns + categorical_cols = df.select_dtypes( + exclude=["int", "float", "float64", "float32", "int32", "int64"] + ).columns + df[categorical_cols] = df[categorical_cols].apply(lambda col: le.fit_transform(col)) + + +def test_types(**kwargs): + test_passed = True + for k, v in kwargs.items(): + # check for incorrect types + if ( + (k == "X") + or (k == "X_test") + or (k == "y") + or (k == "y_test") + or (k == "y_true") + or (k == "y_probas") + ): + # FIXME: do this individually + if not isinstance( + v, + ( + Sequence, + Iterable, + np.ndarray, + np.generic, + pd.DataFrame, + pd.Series, + list, + ), + ): + wandb.termerror(f"{k} is not an array. Please try again.") + test_passed = False + # check for classifier types + if k == "model": + if (not sklearn.base.is_classifier(v)) and ( + not sklearn.base.is_regressor(v) + ): + wandb.termerror( + f"{k} is not a classifier or regressor. Please try again." + ) + test_passed = False + elif k == "clf" or k == "binary_clf": + if not (sklearn.base.is_classifier(v)): + wandb.termerror(f"{k} is not a classifier. Please try again.") + test_passed = False + elif k == "regressor": + if not sklearn.base.is_regressor(v): + wandb.termerror(f"{k} is not a regressor. Please try again.") + test_passed = False + elif k == "clusterer": + if not (getattr(v, "_estimator_type", None) == "clusterer"): + wandb.termerror(f"{k} is not a clusterer. Please try again.") + test_passed = False + return test_passed + + +def test_fitted(model): + try: + model.predict(np.zeros((7, 3))) + except sklearn.exceptions.NotFittedError: + wandb.termerror("Please fit the model before passing it in.") + return False + except AttributeError: + # Some clustering models (LDA, PCA, Agglomerative) don't implement ``predict`` + try: + sklearn.utils.validation.check_is_fitted( + model, + [ + "coef_", + "estimator_", + "labels_", + "n_clusters_", + "children_", + "components_", + "n_components_", + "n_iter_", + "n_batch_iter_", + "explained_variance_", + "singular_values_", + "mean_", + ], + all_or_any=any, + ) + except sklearn.exceptions.NotFittedError: + wandb.termerror("Please fit the model before passing it in.") + return False + else: + return True + except Exception: + # Assume it's fitted, since ``NotFittedError`` wasn't raised + return True + + +# Test Asummptions for plotting parameters and datasets +def test_missing(**kwargs): + test_passed = True + for k, v in kwargs.items(): + # Missing/empty params/datapoint arrays + if v is None: + wandb.termerror(f"{k} is None. Please try again.") + test_passed = False + if (k == "X") or (k == "X_test"): + if isinstance(v, scipy.sparse.csr.csr_matrix): + v = v.toarray() + elif isinstance(v, (pd.DataFrame, pd.Series)): + v = v.to_numpy() + elif isinstance(v, list): + v = np.asarray(v) + + # Warn the user about missing values + missing = 0 + missing = np.count_nonzero(pd.isnull(v)) + if missing > 0: + wandb.termwarn(f"{k} contains {missing} missing values. ") + test_passed = False + # Ensure the dataset contains only integers + non_nums = 0 + if v.ndim == 1: + non_nums = sum( + 1 + for val in v + if ( + not isinstance(val, (int, float, complex)) + and not isinstance(val, np.number) + ) + ) + else: + non_nums = sum( + 1 + for sl in v + for val in sl + if ( + not isinstance(val, (int, float, complex)) + and not isinstance(val, np.number) + ) + ) + if non_nums > 0: + wandb.termerror( + f"{k} contains values that are not numbers. Please vectorize, label encode or one hot encode {k} " + "and call the plotting function again." + ) + test_passed = False + return test_passed + + +def round_3(n): + return round(n, 3) + + +def round_2(n): + return round(n, 2) diff --git a/lib/python3.12/site-packages/wandb/integration/tensorboard/__init__.py b/lib/python3.12/site-packages/wandb/integration/tensorboard/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..80f60ef258a10fb2e7ad2f81594084ccd6098f6c --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/tensorboard/__init__.py @@ -0,0 +1,10 @@ +"""wandb integration tensorboard module.""" + +from .log import _log, log, reset_state, tf_summary_to_dict +from .monkeypatch import patch, unpatch + +__all__ = [ + "patch", + "unpatch", + "log", +] diff --git a/lib/python3.12/site-packages/wandb/integration/tensorboard/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/tensorboard/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2d41e05b7a02c8e42291ca9b1043fdd8d33f2085 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/tensorboard/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/tensorboard/__pycache__/log.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/tensorboard/__pycache__/log.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0ec14e183ae3a0fa8108991abff0f4c2d4825f2a Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/tensorboard/__pycache__/log.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/tensorboard/__pycache__/monkeypatch.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/tensorboard/__pycache__/monkeypatch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9b1fc395996436bf1d0e99f3919b9607defb0d98 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/tensorboard/__pycache__/monkeypatch.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/tensorboard/log.py b/lib/python3.12/site-packages/wandb/integration/tensorboard/log.py new file mode 100644 index 0000000000000000000000000000000000000000..5062a9b0d14109aae1901e751e7864d701d2b7da --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/tensorboard/log.py @@ -0,0 +1,351 @@ +import io +import re +import time +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union + +import wandb +import wandb.util +from wandb.sdk.lib import telemetry + +if TYPE_CHECKING: + import numpy as np + + from wandb.sdk.internal.tb_watcher import TBHistory + +# We have at least the default namestep and a global step to track +# TODO: reset this structure on wandb.finish +STEPS: Dict[str, Dict[str, Any]] = { + "": {"step": 0}, + "global": {"step": 0, "last_log": None}, +} +# TODO(cling): Set these when tensorboard behavior is configured. +# We support rate limited logging by setting this to number of seconds, +# can be a floating point. +RATE_LIMIT_SECONDS: Optional[Union[float, int]] = None +IGNORE_KINDS = ["graphs"] +tensor_util = wandb.util.get_module("tensorboard.util.tensor_util") + + +# prefer tensorboard, fallback to protobuf in tensorflow when tboard isn't available +pb = wandb.util.get_module( + "tensorboard.compat.proto.summary_pb2" +) or wandb.util.get_module("tensorflow.core.framework.summary_pb2") + +Summary = pb.Summary if pb else None + + +def make_ndarray(tensor: Any) -> Optional["np.ndarray"]: + if tensor_util: + res = tensor_util.make_ndarray(tensor) + # Tensorboard can log generic objects, and we don't want to save them + if res.dtype == "object": + return None + else: + return res # type: ignore + else: + wandb.termwarn( + "Can't convert tensor summary, upgrade tensorboard with `pip" + " install tensorboard --upgrade`" + ) + return None + + +def namespaced_tag(tag: str, namespace: str = "") -> str: + if not namespace: + return tag + else: + return namespace + "/" + tag + + +def history_image_key(key: str, namespace: str = "") -> str: + """Convert invalid filesystem characters to _ for use in History keys. + + Unfortunately this means currently certain image keys will collide silently. We + implement this mapping up here in the TensorFlow stuff rather than in the History + stuff so that we don't have to store a mapping anywhere from the original keys to + the safe ones. + """ + return namespaced_tag(re.sub(r"[/\\]", "_", key), namespace) + + +def tf_summary_to_dict( # noqa: C901 + tf_summary_str_or_pb: Any, namespace: str = "" +) -> Optional[Dict[str, Any]]: + """Convert a Tensorboard Summary to a dictionary. + + Accepts a tensorflow.summary.Summary, one encoded as a string, + or a list of such encoded as strings. + """ + values = {} + if hasattr(tf_summary_str_or_pb, "summary"): + summary_pb = tf_summary_str_or_pb.summary + values[namespaced_tag("global_step", namespace)] = tf_summary_str_or_pb.step + values["_timestamp"] = tf_summary_str_or_pb.wall_time + elif isinstance(tf_summary_str_or_pb, (str, bytes, bytearray)): + summary_pb = Summary() + summary_pb.ParseFromString(tf_summary_str_or_pb) + elif hasattr(tf_summary_str_or_pb, "__iter__"): + summary_pb = [Summary() for _ in range(len(tf_summary_str_or_pb))] + for i, summary in enumerate(tf_summary_str_or_pb): + summary_pb[i].ParseFromString(summary) + if i > 0: + summary_pb[0].MergeFrom(summary_pb[i]) + summary_pb = summary_pb[0] + else: + summary_pb = tf_summary_str_or_pb + + if not hasattr(summary_pb, "value") or len(summary_pb.value) == 0: + # Ignore these, caller is responsible for handling None + return None + + def encode_images(_img_strs: List[bytes], _value: Any) -> None: + try: + from PIL import Image + except ImportError: + wandb.termwarn( + "Install pillow if you are logging images with Tensorboard. " + "To install, run `pip install pillow`.", + repeat=False, + ) + return None + + if len(_img_strs) == 0: + return None + + images: List[Union[wandb.Video, wandb.Image]] = [] + for _img_str in _img_strs: + # Supports gifs from TensorboardX + if _img_str.startswith(b"GIF"): + images.append(wandb.Video(io.BytesIO(_img_str), format="gif")) + else: + images.append(wandb.Image(Image.open(io.BytesIO(_img_str)))) + tag_idx = _value.tag.rsplit("/", 1) + if len(tag_idx) > 1 and tag_idx[1].isdigit(): + tag, idx = tag_idx + values.setdefault(history_image_key(tag, namespace), []).extend(images) + else: + values[history_image_key(_value.tag, namespace)] = images + + return None + + for value in summary_pb.value: + kind = value.WhichOneof("value") + if kind in IGNORE_KINDS: + continue + if kind == "simple_value": + values[namespaced_tag(value.tag, namespace)] = value.simple_value + elif kind == "tensor": + plugin_name = value.metadata.plugin_data.plugin_name + if plugin_name == "scalars" or plugin_name == "": + values[namespaced_tag(value.tag, namespace)] = make_ndarray( + value.tensor + ) + elif plugin_name == "images": + img_strs = value.tensor.string_val[2:] # First two items are dims. + encode_images(img_strs, value) + elif plugin_name == "histograms": + # https://github.com/tensorflow/tensorboard/blob/master/tensorboard/plugins/histogram/summary_v2.py#L15-L26 + ndarray = make_ndarray(value.tensor) + if ndarray is None: + continue + shape = ndarray.shape + counts = [] + bins = [] + if shape[0] > 1: + bins.append(ndarray[0][0]) # Add the left most edge + for v in ndarray: + counts.append(v[2]) + bins.append(v[1]) # Add the right most edges + elif shape[0] == 1: + counts = [ndarray[0][2]] + bins = ndarray[0][:2] + if len(counts) > 0: + try: + # TODO: we should just re-bin if there are too many buckets + values[namespaced_tag(value.tag, namespace)] = wandb.Histogram( + np_histogram=(counts, bins) # type: ignore + ) + except ValueError: + wandb.termwarn( + f'Not logging key "{namespaced_tag(value.tag, namespace)}". ' + f"Histograms must have fewer than {wandb.Histogram.MAX_LENGTH} bins", + repeat=False, + ) + elif plugin_name == "pr_curves": + pr_curve_data = make_ndarray(value.tensor) + if pr_curve_data is None: + continue + precision = pr_curve_data[-2, :].tolist() + recall = pr_curve_data[-1, :].tolist() + # TODO: (kdg) implement spec for showing additional info in tool tips + # true_pos = pr_curve_data[1,:] + # false_pos = pr_curve_data[2,:] + # true_neg = pr_curve_data[1,:] + # false_neg = pr_curve_data[1,:] + # threshold = [1.0 / n for n in range(len(true_pos), 0, -1)] + # min of each in case tensorboard ever changes their pr_curve + # to allow for different length outputs + data = [] + for i in range(min(len(precision), len(recall))): + # drop additional threshold values if they exist + if precision[i] != 0 or recall[i] != 0: + data.append((recall[i], precision[i])) + # sort data so custom chart looks the same as tb generated pr curve + # ascending recall, descending precision for the same recall values + data = sorted(data, key=lambda x: (x[0], -x[1])) + data_table = wandb.Table(data=data, columns=["recall", "precision"]) + name = namespaced_tag(value.tag, namespace) + + values[name] = wandb.plot_table( + "wandb/line/v0", + data_table, + {"x": "recall", "y": "precision"}, + {"title": f"{name} Precision v. Recall"}, + ) + elif kind == "image": + img_str = value.image.encoded_image_string + encode_images([img_str], value) + # Coming soon... + # elif kind == "audio": + # audio = wandb.Audio( + # six.BytesIO(value.audio.encoded_audio_string), + # sample_rate=value.audio.sample_rate, + # content_type=value.audio.content_type, + # ) + elif kind == "histo": + tag = namespaced_tag(value.tag, namespace) + if len(value.histo.bucket_limit) >= 3: + first = ( + value.histo.bucket_limit[0] + + value.histo.bucket_limit[0] + - value.histo.bucket_limit[1] + ) + last = ( + value.histo.bucket_limit[-2] + + value.histo.bucket_limit[-2] + - value.histo.bucket_limit[-3] + ) + np_histogram = ( + list(value.histo.bucket), + [first] + value.histo.bucket_limit[:-1] + [last], + ) + try: + # TODO: we should just re-bin if there are too many buckets + values[tag] = wandb.Histogram(np_histogram=np_histogram) # type: ignore + except ValueError: + wandb.termwarn( + f"Not logging key {tag!r}. " + f"Histograms must have fewer than {wandb.Histogram.MAX_LENGTH} bins", + repeat=False, + ) + else: + # TODO: is there a case where we can render this? + wandb.termwarn( + f"Not logging key {tag!r}. Found a histogram with only 2 bins.", + repeat=False, + ) + # TODO(jhr): figure out how to share this between userspace and internal process or dont + # elif value.tag == "_hparams_/session_start_info": + # if wandb.util.get_module("tensorboard.plugins.hparams"): + # from tensorboard.plugins.hparams import plugin_data_pb2 + # + # plugin_data = plugin_data_pb2.HParamsPluginData() # + # plugin_data.ParseFromString(value.metadata.plugin_data.content) + # for key, param in six.iteritems(plugin_data.session_start_info.hparams): + # if not wandb.run.config.get(key): + # wandb.run.config[key] = ( + # param.number_value or param.string_value or param.bool_value + # ) + # else: + # wandb.termerror( + # "Received hparams tf.summary, but could not import " + # "the hparams plugin from tensorboard" + # ) + return values + + +def reset_state() -> None: + """Internal method for resetting state, called by wandb.finish().""" + global STEPS + STEPS = {"": {"step": 0}, "global": {"step": 0, "last_log": None}} + + +def _log( + tf_summary_str_or_pb: Any, + history: Optional["TBHistory"] = None, + step: int = 0, + namespace: str = "", + **kwargs: Any, +) -> None: + """Logs a tfsummary to wandb. + + Can accept a tf summary string or parsed event. Will use wandb.run.history unless a + history object is passed. Can optionally namespace events. Results are committed + when step increases for this namespace. + + NOTE: This assumes that events being passed in are in chronological order + """ + global STEPS + global RATE_LIMIT_SECONDS + # To handle multiple global_steps, we keep track of them here instead + # of the global log + last_step = STEPS.get(namespace, {"step": 0}) + + # Commit our existing data if this namespace increased its step + commit = False + if last_step["step"] < step: + commit = True + + log_dict = tf_summary_to_dict(tf_summary_str_or_pb, namespace) + if log_dict is None: + # not an event, just return + return + + # Pass timestamp to history for loading historic data + timestamp = log_dict.get("_timestamp", time.time()) + # Store our initial timestamp + if STEPS["global"]["last_log"] is None: + STEPS["global"]["last_log"] = timestamp + # Rollup events that share the same step across namespaces + if commit and step == STEPS["global"]["step"]: + commit = False + # Always add the biggest global_step key for non-default namespaces + if step > STEPS["global"]["step"]: + STEPS["global"]["step"] = step + if namespace != "": + log_dict["global_step"] = STEPS["global"]["step"] + + # Keep internal step counter + STEPS[namespace] = {"step": step} + + if commit: + # Only commit our data if we're below the rate limit or don't have one + if ( + RATE_LIMIT_SECONDS is None + or timestamp - STEPS["global"]["last_log"] >= RATE_LIMIT_SECONDS + ): + if history is None: + if wandb.run is not None: + wandb.run._log({}) + else: + history.add({}) + + STEPS["global"]["last_log"] = timestamp + + if history is None: + if wandb.run is not None: + wandb.run._log(log_dict, commit=False) + else: + history._row_update(log_dict) + + +def log(tf_summary_str_or_pb: Any, step: int = 0, namespace: str = "") -> None: + if wandb.run is None: + raise wandb.Error( + "You must call `wandb.init()` before calling `wandb.tensorflow.log`" + ) + + with telemetry.context() as tel: + tel.feature.tensorboard_log = True + + _log(tf_summary_str_or_pb, namespace=namespace, step=step) diff --git a/lib/python3.12/site-packages/wandb/integration/tensorboard/monkeypatch.py b/lib/python3.12/site-packages/wandb/integration/tensorboard/monkeypatch.py new file mode 100644 index 0000000000000000000000000000000000000000..a0db49aaf79bad018642eea9e1d555f65827ba6e --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/tensorboard/monkeypatch.py @@ -0,0 +1,186 @@ +"""monkeypatch: patch code to add tensorboard hooks.""" + +import os +import re +import socket +from typing import Any, Optional + +import wandb +import wandb.util + +TENSORBOARD_C_MODULE = "tensorflow.python.ops.gen_summary_ops" +TENSORBOARD_X_MODULE = "tensorboardX.writer" +TENSORFLOW_PY_MODULE = "tensorflow.python.summary.writer.writer" +TENSORBOARD_WRITER_MODULE = "tensorboard.summary.writer.event_file_writer" +TENSORBOARD_PYTORCH_MODULE = "torch.utils.tensorboard.writer" + + +def unpatch() -> None: + for module, method in wandb.patched["tensorboard"]: + writer = wandb.util.get_module(module, lazy=False) + setattr(writer, method, getattr(writer, f"orig_{method}")) + wandb.patched["tensorboard"] = [] + + +def patch( + save: bool = True, + tensorboard_x: Optional[bool] = None, + pytorch: Optional[bool] = None, + root_logdir: str = "", +) -> None: + if len(wandb.patched["tensorboard"]) > 0: + raise ValueError( + "Tensorboard already patched. Call `wandb.tensorboard.unpatch()` first; " + "remove `sync_tensorboard=True` from `wandb.init`; " + "or only call `wandb.tensorboard.patch` once." + ) + + # TODO: Some older versions of tensorflow don't require tensorboard to be present. + # we may want to lift this requirement, but it's safer to have it for now + wandb.util.get_module( + "tensorboard", required="Please install tensorboard package", lazy=False + ) + c_writer = wandb.util.get_module(TENSORBOARD_C_MODULE, lazy=False) + py_writer = wandb.util.get_module(TENSORFLOW_PY_MODULE, lazy=False) + tb_writer = wandb.util.get_module(TENSORBOARD_WRITER_MODULE, lazy=False) + pt_writer = wandb.util.get_module(TENSORBOARD_PYTORCH_MODULE, lazy=False) + tbx_writer = wandb.util.get_module(TENSORBOARD_X_MODULE, lazy=False) + + if not pytorch and not tensorboard_x and c_writer: + _patch_tensorflow2( + writer=c_writer, + module=TENSORBOARD_C_MODULE, + save=save, + root_logdir=root_logdir, + ) + # This is for tensorflow <= 1.15 (tf.compat.v1.summary.FileWriter) + if py_writer: + _patch_file_writer( + writer=py_writer, + module=TENSORFLOW_PY_MODULE, + save=save, + root_logdir=root_logdir, + ) + if tb_writer: + _patch_file_writer( + writer=tb_writer, + module=TENSORBOARD_WRITER_MODULE, + save=save, + root_logdir=root_logdir, + ) + if pt_writer: + _patch_file_writer( + writer=pt_writer, + module=TENSORBOARD_PYTORCH_MODULE, + save=save, + root_logdir=root_logdir, + ) + if tbx_writer: + _patch_file_writer( + writer=tbx_writer, + module=TENSORBOARD_X_MODULE, + save=save, + root_logdir=root_logdir, + ) + if not c_writer and not tb_writer and not tb_writer: + wandb.termerror("Unsupported tensorboard configuration") + + +def _patch_tensorflow2( + writer: Any, + module: Any, + save: bool = True, + root_logdir: str = "", +) -> None: + # This configures TensorFlow 2 style Tensorboard logging + old_csfw_func = writer.create_summary_file_writer + logdir_hist = [] + + def new_csfw_func(*args: Any, **kwargs: Any) -> Any: + logdir = ( + kwargs["logdir"].numpy().decode("utf8") + if hasattr(kwargs["logdir"], "numpy") + else kwargs["logdir"] + ) + logdir_hist.append(logdir) + root_logdir_arg = root_logdir + + if len(set(logdir_hist)) > 1 and root_logdir == "": + wandb.termwarn( + "When using several event log directories, " + 'please call `wandb.tensorboard.patch(root_logdir="...")` before `wandb.init`' + ) + # if the logdir contains the hostname, the writer was not given a logdir. + # In this case, the generated logdir + # is generated and ends with the hostname, update the root_logdir to match. + hostname = socket.gethostname() + search = re.search(rf"-\d+_{hostname}", logdir) + if search: + root_logdir_arg = logdir[: search.span()[1]] + elif root_logdir is not None and not os.path.abspath(logdir).startswith( + os.path.abspath(root_logdir) + ): + wandb.termwarn( + "Found log directory outside of given root_logdir, " + f"dropping given root_logdir for event file in {logdir}" + ) + root_logdir_arg = "" + + _notify_tensorboard_logdir(logdir, save=save, root_logdir=root_logdir_arg) + return old_csfw_func(*args, **kwargs) + + writer.orig_create_summary_file_writer = old_csfw_func + writer.create_summary_file_writer = new_csfw_func + wandb.patched["tensorboard"].append([module, "create_summary_file_writer"]) + + +def _patch_file_writer( + writer: Any, + module: Any, + save: bool = True, + root_logdir: str = "", +) -> None: + # This configures non-TensorFlow Tensorboard logging, or tensorflow <= 1.15 + logdir_hist = [] + + class TBXEventFileWriter(writer.EventFileWriter): + def __init__(self, logdir: str, *args: Any, **kwargs: Any) -> None: + logdir_hist.append(logdir) + root_logdir_arg = root_logdir + if len(set(logdir_hist)) > 1 and root_logdir == "": + wandb.termwarn( + "When using several event log directories, " + 'please call `wandb.tensorboard.patch(root_logdir="...")` before `wandb.init`' + ) + + # if the logdir contains the hostname, the writer was not given a logdir. + # In this case, the logdir is generated and ends with the hostname, + # update the root_logdir to match. + hostname = socket.gethostname() + search = re.search(rf"-\d+_{hostname}", logdir) + if search: + root_logdir_arg = logdir[: search.span()[1]] + + elif root_logdir is not None and not os.path.abspath(logdir).startswith( + os.path.abspath(root_logdir) + ): + wandb.termwarn( + "Found log directory outside of given root_logdir, " + f"dropping given root_logdir for event file in {logdir}" + ) + root_logdir_arg = "" + + _notify_tensorboard_logdir(logdir, save=save, root_logdir=root_logdir_arg) + + super().__init__(logdir, *args, **kwargs) + + writer.orig_EventFileWriter = writer.EventFileWriter + writer.EventFileWriter = TBXEventFileWriter + wandb.patched["tensorboard"].append([module, "EventFileWriter"]) + + +def _notify_tensorboard_logdir( + logdir: str, save: bool = True, root_logdir: str = "" +) -> None: + if wandb.run is not None: + wandb.run._tensorboard_callback(logdir, save=save, root_logdir=root_logdir) diff --git a/lib/python3.12/site-packages/wandb/integration/tensorflow/__init__.py b/lib/python3.12/site-packages/wandb/integration/tensorflow/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b5a5838c7d3eb5e5ea9e50b420bf42a605084dc6 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/tensorflow/__init__.py @@ -0,0 +1,5 @@ +"""api.""" + +from wandb.integration.tensorboard import log + +from .estimator_hook import WandbHook diff --git a/lib/python3.12/site-packages/wandb/integration/tensorflow/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/tensorflow/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..857a7623f0e1d7e38e20c2743ed9fb5a55e2e692 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/tensorflow/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/tensorflow/__pycache__/estimator_hook.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/tensorflow/__pycache__/estimator_hook.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6f2c9af578b37b38ccb951c9c5c6a74ac7618273 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/tensorflow/__pycache__/estimator_hook.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/tensorflow/estimator_hook.py b/lib/python3.12/site-packages/wandb/integration/tensorflow/estimator_hook.py new file mode 100644 index 0000000000000000000000000000000000000000..88d58abe3c6ee81759ce01a8e7254e1e324034de --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/tensorflow/estimator_hook.py @@ -0,0 +1,54 @@ +import tensorflow as tf + +import wandb +from wandb.sdk.lib import telemetry + +if hasattr(tf.estimator, "SessionRunHook"): + # In tf 1.14 and beyond, SessionRunHook is in the estimator package. + SessionRunHook = tf.estimator.SessionRunHook + SessionRunArgs = tf.estimator.SessionRunArgs +else: + # In older versions it's in train. + SessionRunHook = tf.train.SessionRunHook + SessionRunArgs = tf.train.SessionRunArgs + +if hasattr(tf.train, "get_global_step"): + get_global_step = tf.train.get_global_step +else: + get_global_step = tf.compat.v1.train.get_global_step + +if hasattr(tf.summary, "merge_all"): + merge_all_summaries = tf.summary.merge_all +else: + merge_all_summaries = tf.compat.v1.summary.merge_all + + +class WandbHook(SessionRunHook): + def __init__(self, summary_op=None, steps_per_log=1000, history=None): + self._summary_op = summary_op + self._steps_per_log = steps_per_log + self._history = history + + with telemetry.context() as tel: + tel.feature.estimator_hook = True + + def begin(self): + if wandb.run is None: + raise wandb.Error("You must call `wandb.init()` before calling `WandbHook`") + if self._summary_op is None: + self._summary_op = merge_all_summaries() + self._step = -1 + + def before_run(self, run_context): + return SessionRunArgs( + {"summary": self._summary_op, "global_step": get_global_step()} + ) + + def after_run(self, run_context, run_values): + step = run_values.results["global_step"] + if step % self._steps_per_log == 0: + wandb.tensorboard._log( + run_values.results["summary"], + history=self._history, + step=step, + ) diff --git a/lib/python3.12/site-packages/wandb/integration/torch/__init__.py b/lib/python3.12/site-packages/wandb/integration/torch/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/wandb/integration/torch/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/torch/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2e597e621b94cb1e4ba10e34901ca477ac7d0e38 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/torch/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/torch/__pycache__/wandb_torch.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/torch/__pycache__/wandb_torch.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..50eeb3e26949a1c5ca534e12925343976baca14a Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/torch/__pycache__/wandb_torch.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/torch/wandb_torch.py b/lib/python3.12/site-packages/wandb/integration/torch/wandb_torch.py new file mode 100644 index 0000000000000000000000000000000000000000..28d5971f8fd640ef6315aa22cb795163bc9d0724 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/torch/wandb_torch.py @@ -0,0 +1,554 @@ +"""PyTorch-specific functionality.""" + +import itertools +from functools import reduce +from operator import mul +from typing import TYPE_CHECKING, List + +import wandb +from wandb import util +from wandb.data_types import Node + +torch = None + +if TYPE_CHECKING: + from torch import Tensor + from torch.nn import Module + + +def nested_shape(array_or_tuple, seen=None): + """Figure out the shape of tensors possibly embedded in tuples. + + for example: + - [0,0] returns (2) + - ([0,0], [0,0]) returns (2,2) + - (([0,0], [0,0]),[0,0]) returns ((2,2),2). + """ + if seen is None: + seen = set() + if hasattr(array_or_tuple, "size"): + # pytorch tensors use V.size() to get size of tensor + return list(array_or_tuple.size()) + elif hasattr(array_or_tuple, "get_shape"): + # tensorflow uses V.get_shape() to get size of tensor + return array_or_tuple.get_shape().as_list() + elif hasattr(array_or_tuple, "shape"): + return array_or_tuple.shape + + seen.add(id(array_or_tuple)) + try: + # treat object as iterable + return [ + nested_shape(item, seen) if id(item) not in seen else 0 + for item in list(array_or_tuple) + ] + except TypeError: + # object is not actually iterable + # LB: Maybe we should throw an error? + return [] + + +LOG_TRACK_COUNT, LOG_TRACK_THRESHOLD = range(2) + + +def log_track_init(log_freq: int) -> List[int]: + """Create tracking structure used by log_track_update.""" + log_track = [0, 0] + log_track[LOG_TRACK_THRESHOLD] = log_freq + return log_track + + +def log_track_update(log_track: int) -> bool: + """Count (log_track[0]) up to threshold (log_track[1]), reset count (log_track[0]) and return true when reached.""" + log_track[LOG_TRACK_COUNT] += 1 + if log_track[LOG_TRACK_COUNT] < log_track[LOG_TRACK_THRESHOLD]: + return False + log_track[LOG_TRACK_COUNT] = 0 + return True + + +class TorchHistory: + """History methods specific to PyTorch.""" + + def __init__(self): + global torch + torch = wandb.util.get_module("torch", "Could not import torch") + self._hook_handles = {} + self._num_bins = 64 + self._is_cuda_histc_supported = None + self.hook_torch = TorchGraph.hook_torch + + def add_log_parameters_hook( + self, + module: "Module", + name: str = "", + prefix: str = "", + log_freq: int = 0, + ) -> None: + """This instruments hooks into the pytorch module. + + log parameters after a forward pass + log_freq - log gradients/parameters every N batches. + """ + # if name is not None: + prefix = prefix + name + + if not hasattr(module, "_wandb_hook_names"): + module._wandb_hook_names = [] + + def parameter_log_hook(module, input_, output, log_track): + if not log_track_update(log_track): + return + for name, parameter in module.named_parameters(): + # for pytorch 0.3 Variables + if isinstance(parameter, torch.autograd.Variable): + data = parameter.data + else: + data = parameter + self.log_tensor_stats(data.cpu(), "parameters/" + prefix + name) + + log_track_params = log_track_init(log_freq) + try: + hook = module.register_forward_hook( + lambda mod, inp, outp: parameter_log_hook( + mod, inp, outp, log_track_params + ) + ) + self._hook_handles["parameters/" + prefix] = hook + module._wandb_hook_names.append("parameters/" + prefix) + except RuntimeError as e: + wandb.termwarn( + f"Trying to register forward_hook failed ({e}) - skipping parameter tracking." + ) + + def add_log_gradients_hook( + self, + module: "Module", + name: str = "", + prefix: str = "", + log_freq: int = 0, + ) -> None: + """This instruments hooks into the PyTorch module slog gradients after a backward pass. + + Args: + module: torch.nn.Module - the module to instrument + name: str - the name of the module + prefix: str - the prefix to add to the name + log_freq: log gradients/parameters every N batches + """ + # if name is not None: + prefix = prefix + name + + if not hasattr(module, "_wandb_hook_names"): + module._wandb_hook_names = [] + + for name, parameter in module.named_parameters(): + if parameter.requires_grad: + log_track_grad = log_track_init(log_freq) + module._wandb_hook_names.append("gradients/" + prefix + name) + self._hook_variable_gradient_stats( + parameter, "gradients/" + prefix + name, log_track_grad + ) + + def log_tensor_stats(self, tensor, name): # noqa: C901 + """Add distribution statistics on a tensor's elements to the current History entry.""" + # TODO Handle the case of duplicate names. + if isinstance(tensor, (tuple, list)): + while isinstance(tensor, (tuple, list)) and isinstance( + tensor[0], (tuple, list) + ): + tensor = [item for sublist in tensor for item in sublist] + tensor = torch.cat([t.detach().clone().reshape(-1) for t in tensor]) + + tensor = tensor.detach().clone() + # checking for inheritance from _TensorBase didn't work for some reason + if not hasattr(tensor, "shape"): + cls = type(tensor) + raise TypeError(f"Expected Tensor, not {cls.__module__}.{cls.__name__}") + + # Sparse tensors have a bunch of implicit zeros. In order to histo them correctly, + # we have to count them up and add them to the histo ourselves. + sparse_zeros = None + if tensor.is_sparse: + # Have to call this on a sparse tensor before most other ops. + tensor = tensor.cpu().coalesce() + + backing_values = tensor._values() + sparse_zeros = tensor.numel() - backing_values.numel() + tensor = backing_values + + flat = tensor.reshape(-1) + + if flat.is_cuda: + if self._is_cuda_histc_supported is None: + try: + flat.histc(bins=self._num_bins) + except RuntimeError: + self._is_cuda_histc_supported = False + else: + self._is_cuda_histc_supported = True + + # As of torch 1.0.1.post2+nightly, float16 cuda summary ops are not supported (convert to float32) + if not self._is_cuda_histc_supported: + flat = flat.cpu() + elif not isinstance( + flat, (torch.cuda.FloatTensor, torch.cuda.DoubleTensor) + ): + flat = flat.type(torch.cuda.FloatTensor) + + # Since we use histc, we need to make sure that torch supports the operation on CPU, + # otherwise we'll get a runtime error. Hence, we need to upcast to float32. + if not flat.is_cuda and not isinstance( + flat, (torch.FloatTensor, torch.DoubleTensor) + ): + flat = flat.type(torch.FloatTensor) + + # Skip logging if all values are nan or inf or the tensor is empty. + if self._no_finite_values(flat): + return + + # Remove nans and infs if present. There's no good way to represent that in histograms. + flat = self._remove_infs_nans(flat) + + tmin = flat.min().item() + tmax = flat.max().item() + if sparse_zeros: + # If we've got zeros to add in, make sure zero is in the hist range. + tmin = 0 if tmin > 0 else tmin + tmax = 0 if tmax < 0 else tmax + # Anecdotally, this can somehow happen sometimes. Maybe a precision error + # in min()/max() above. Swap here to prevent a runtime error. + # If all values are equal, just return a single bin. + if tmin > tmax: + tmin, tmax = tmax, tmin + if tmin == tmax: + tensor = torch.Tensor([flat.numel()]) + tensor = tensor.cpu().clone().detach() + bins = torch.Tensor([tmin, tmax]) + else: + tensor = flat.histc(bins=self._num_bins, min=tmin, max=tmax) + tensor = tensor.cpu().detach().clone() + bins = torch.linspace(tmin, tmax, steps=self._num_bins + 1) + + # Add back zeroes from a sparse tensor. + if sparse_zeros: + bins_np = bins.numpy() + tensor_np = tensor.numpy() + bin_idx = 0 + num_buckets = len(bins_np) - 1 + for i in range(num_buckets): + start = bins_np[i] + end = bins_np[i + 1] + # There are 3 cases to consider here, all of which mean we've found the right bucket + # 1. The bucket range contains zero. + # 2. The bucket range lower bound *is* zero. + # 3. This is the last bucket and the bucket range upper bound is zero. + if (start <= 0 and end > 0) or (i == num_buckets - 1 and end == 0): + bin_idx = i + break + + tensor_np[bin_idx] += sparse_zeros + tensor = torch.Tensor(tensor_np) + bins = torch.Tensor(bins_np) + + wandb.run._log( + {name: wandb.Histogram(np_histogram=(tensor.tolist(), bins.tolist()))}, + commit=False, + ) + + def _hook_variable_gradient_stats(self, var, name, log_track): + """Logs a Variable's gradient's distribution statistics next time backward() is called on it.""" + if not isinstance(var, torch.autograd.Variable): + cls = type(var) + raise TypeError( + f"Expected torch.Variable, not {cls.__module__}.{cls.__name__}" + ) + + handle = self._hook_handles.get(name) + if handle is not None and self._torch_hook_handle_is_valid(handle): + raise ValueError(f'A hook has already been set under name "{name}"') + + def _callback(grad, log_track): + if not log_track_update(log_track): + return + self.log_tensor_stats(grad.data, name) + + handle = var.register_hook(lambda grad: _callback(grad, log_track)) + self._hook_handles[name] = handle + return handle + + def unhook_all(self): + for handle in self._hook_handles.values(): + handle.remove() + self._hook_handles = {} + + def unhook(self, name): + handle = self._hook_handles.pop(name) + handle.remove() + + def _torch_hook_handle_is_valid(self, handle): + d = handle.hooks_dict_ref() + if d is None: + return False + else: + return handle.id in d + + def _no_finite_values(self, tensor: "Tensor") -> bool: + return tensor.shape == torch.Size([0]) or (~torch.isfinite(tensor)).all().item() + + def _remove_infs_nans(self, tensor: "Tensor") -> "Tensor": + if not torch.isfinite(tensor).all(): + tensor = tensor[torch.isfinite(tensor)] + + return tensor + + +class TorchGraph(wandb.data_types.Graph): + def __init__(self): + super().__init__("torch") + self._graph_hooks = set() + + @classmethod + def hook_torch(cls, model, criterion=None, graph_idx=0): + wandb.termlog("logging graph, to disable use `wandb.watch(log_graph=False)`") + graph = TorchGraph() + graph.hook_torch_modules(model, criterion, graph_idx=graph_idx) + return graph + + def create_forward_hook(self, name, graph_idx): + graph = self + + def after_forward_hook(module, input, output): + if id(module) not in self._graph_hooks: + # hook already processed -> noop + return + if not isinstance(output, tuple): + output = (output,) + parameters = [ + (pname, list(param.size())) + for pname, param in module.named_parameters() + ] + + node = Node( + id=id(module), + name=name, + class_name=str(module), + output_shape=nested_shape(output), + parameters=parameters, + num_parameters=[reduce(mul, size, 1) for (pname, size) in parameters], + ) + graph.nodes_by_id[id(module)] = node + for param in module.parameters(): + graph.nodes_by_id[id(param)] = node + graph.add_node(node) + if not graph.criterion_passed: + if hasattr(output[0], "grad_fn"): + graph.criterion = output[0].grad_fn + elif ( + isinstance(output[0], list) + and output[0] + and hasattr(output[0][0], "grad_fn") + ): + graph.criterion = output[0][0].grad_fn + + # hook has been processed + self._graph_hooks -= {id(module)} + + if not self._graph_hooks: + # we went through the entire graph + wandb.run.summary[f"graph_{graph_idx}"] = self + + return after_forward_hook + + def hook_torch_modules( + self, module, criterion=None, prefix=None, graph_idx=0, parent=None + ): + torch = util.get_module("torch", "Could not import torch") + layers = 0 + graph = self + if hasattr(module, "_wandb_watch_called") and module._wandb_watch_called: + raise ValueError( + "You can only call `wandb.watch` once per model. Pass a new instance of the model if you need to call wandb.watch again in your code." + ) + module._wandb_watch_called = True + if criterion: + graph.criterion = criterion + graph.criterion_passed = True + + for name, sub_module in module.named_children(): + name = name or str(layers) + if prefix: + name = prefix + "." + name + layers += 1 + if not isinstance(sub_module, torch.nn.Module): + # TODO: Why does this happen? + break + + # Trying to support torch >0.3 making this code complicated + # We want a list of types that we should recurse into + # Torch 0.3 uses containers + # 0.4 has ModuleList + # 0.4.1 has ModuleDict + module_types = [ + getattr(torch.nn, module_classname) + for module_classname in ( + "Container", + "Sequential", + "ModuleList", + "ModuleDict", + ) + if hasattr(torch.nn, module_classname) + ] + if parent is None: + parent = module + + if isinstance(sub_module, tuple(module_types)): + self.hook_torch_modules(sub_module, prefix=name, parent=parent) + else: + self._graph_hooks |= {id(sub_module)} + try: + graph_hook = sub_module.register_forward_hook( + self.create_forward_hook(name, graph_idx) + ) + wandb.run._torch._hook_handles[ + "topology/" + str(id(graph_hook)) + ] = graph_hook + if not hasattr(parent, "_wandb_hook_names"): + # should never happen but let's be extra safe + parent._wandb_hook_names = [] + parent._wandb_hook_names.append("topology/" + str(id(graph_hook))) + except RuntimeError as e: + wandb.termwarn( + f"Trying to register forward_hook failed ({e}) - skipping graph tracking.", + repeat=False, + ) + + @classmethod + def from_torch_layers(cls, module_graph, variable): + """Recover something like neural net layers from PyTorch Module's and the compute graph from a Variable. + + Example output for a multi-layer RNN. We confusingly assign shared embedding values + to the encoder, but ordered next to the decoder. + + rnns.0.linear.module.weight_raw rnns.0 + rnns.0.linear.module.bias rnns.0 + rnns.1.linear.module.weight_raw rnns.1 + rnns.1.linear.module.bias rnns.1 + rnns.2.linear.module.weight_raw rnns.2 + rnns.2.linear.module.bias rnns.2 + rnns.3.linear.module.weight_raw rnns.3 + rnns.3.linear.module.bias rnns.3 + decoder.weight encoder + decoder.bias decoder + """ + # TODO: We're currently not using this, but I left it here in case we want to resurrect! - CVP + torch = util.get_module("torch", "Could not import torch") + + module_nodes_by_hash = {id(n): n for n in module_graph.nodes} + module_parameter_nodes = [ + n for n in module_graph.nodes if isinstance(n.obj, torch.nn.Parameter) + ] + + names_by_pid = {id(n.obj): n.name for n in module_parameter_nodes} + + reachable_param_nodes = module_graph[0].reachable_descendents() + reachable_params = {} + module_reachable_params = {} + names = {} + for pid, reachable_nodes in reachable_param_nodes.items(): + node = module_nodes_by_hash[pid] + if not isinstance(node.obj, torch.nn.Module): + continue + module = node.obj + reachable_params = {} # by object id + module_reachable_params[id(module)] = reachable_params + names[node.name] = set() + for reachable_hash in reachable_nodes: + reachable = module_nodes_by_hash[reachable_hash] + if isinstance(reachable.obj, torch.nn.Parameter): + param = reachable.obj + reachable_params[id(param)] = param + names[node.name].add(names_by_pid[id(param)]) + + # we look for correspondences between sets of parameters used in subtrees of the + # computation graph and sets of parameters contained in subtrees of the module + # graph + node_depths = {id(n): d for n, d in module_graph[0].descendent_bfs()} + parameter_module_names = {} + parameter_modules = {} + for param_node in ( + n for n in module_graph.nodes if isinstance(n.obj, torch.nn.Parameter) + ): + pid = id(param_node.obj) + best_node = None + best_depth = None + best_reachable_params = None + for node in module_graph.nodes: + if not isinstance(node.obj, torch.nn.Module): + continue + module = node.obj + reachable_params = module_reachable_params[id(module)] + if pid in reachable_params: + depth = node_depths[id(node)] + if best_node is None or (len(reachable_params), depth) <= ( + len(best_reachable_params), + best_depth, + ): + best_node = node + best_depth = depth + best_reachable_params = reachable_params + + parameter_modules[pid] = best_node + parameter_module_names[param_node.name] = best_node.name + + # contains all parameters but only a minimal set of modules necessary + # to contain them (and which ideally correspond to conceptual layers) + reduced_module_graph = cls() + rmg_ids = itertools.count() + rmg_root = Node(id=next(rmg_ids), node=module_graph[0]) + reduced_module_graph.add_node(rmg_root) + reduced_module_graph.root = rmg_root + rmg_nodes_by_pid = {} + + module_nodes_by_pid = {id(n.obj): n for n in module_graph.nodes} + + compute_graph, compute_node_vars = cls.from_torch_compute_graph(variable) + for node, _ in reversed(list(compute_graph[0].ancestor_bfs())): + param = compute_node_vars.get(node.id) + pid = id(param) + if not isinstance(param, torch.nn.Parameter): + continue + if pid not in module_nodes_by_pid: + # not all Parameters that occur in the compute graph come from the Module graph + continue + + # add the nodes in the order we want to display them on the frontend + mid = id(parameter_modules[pid].obj) + if mid in rmg_nodes_by_pid: + rmg_module = rmg_nodes_by_pid[mid] + else: + rmg_module = rmg_nodes_by_pid[mid] = Node( + id=next(rmg_ids), node=module_nodes_by_pid[mid] + ) + reduced_module_graph.add_node(rmg_module) + reduced_module_graph.add_edge(rmg_root, rmg_module) + + rmg_param = Node(id=next(rmg_ids), node=module_nodes_by_pid[pid]) + rmg_nodes_by_pid[pid] = rmg_param + reduced_module_graph.add_node(rmg_param) + + reduced_module_graph.add_edge(rmg_module, rmg_param) + return reduced_module_graph + + @classmethod + def node_from_module(cls, nid, module): + numpy = util.get_module("numpy", "Could not import numpy") + + node = wandb.Node() + node.id = nid + node.child_parameters = 0 + for parameter in module.parameters(): + node.child_parameters += numpy.prod(parameter.size()) + node.class_name = type(module).__name__ + + return node diff --git a/lib/python3.12/site-packages/wandb/integration/ultralytics/__init__.py b/lib/python3.12/site-packages/wandb/integration/ultralytics/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c25ae320c8c946a25150b870770710f58a8bf713 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/ultralytics/__init__.py @@ -0,0 +1,11 @@ +"""Tools for integrating with [`ultralytics`](https://docs.ultralytics.com/). + +Ultralytics is a computer vision framework for training and deploying YOLOv8 models. +""" + +from wandb.integration.ultralytics.callback import ( + WandBUltralyticsCallback, + add_wandb_callback, +) + +__all__ = ("WandBUltralyticsCallback", "add_wandb_callback") diff --git a/lib/python3.12/site-packages/wandb/integration/ultralytics/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/ultralytics/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c223b9d06879bb85b5f38da54951d10912f31e03 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/ultralytics/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/ultralytics/__pycache__/bbox_utils.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/ultralytics/__pycache__/bbox_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2f712a59a9ce76d49b704330ba9ad7d38242801d Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/ultralytics/__pycache__/bbox_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/ultralytics/__pycache__/callback.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/ultralytics/__pycache__/callback.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6b4bdb2ea0758085e3ee0ac7ef9929668ddbaa45 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/ultralytics/__pycache__/callback.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/ultralytics/__pycache__/classification_utils.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/ultralytics/__pycache__/classification_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5fee1f6d28a4ead53162240ba58a0edd341226ef Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/ultralytics/__pycache__/classification_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/ultralytics/__pycache__/mask_utils.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/ultralytics/__pycache__/mask_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..be9eae122f8de0b683c84fd7ece0032dfdf47f93 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/ultralytics/__pycache__/mask_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/ultralytics/__pycache__/pose_utils.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/ultralytics/__pycache__/pose_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3fa957be0935d734e8a71b216beb4e2f0f92c8a3 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/ultralytics/__pycache__/pose_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/ultralytics/bbox_utils.py b/lib/python3.12/site-packages/wandb/integration/ultralytics/bbox_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..4f87bd167bfb45cb4f9a996df301985fa46a0dc9 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/ultralytics/bbox_utils.py @@ -0,0 +1,215 @@ +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch +from tqdm.auto import tqdm +from ultralytics.engine.results import Results +from ultralytics.models.yolo.detect import DetectionPredictor +from ultralytics.utils import ops + +import wandb + + +def scale_bounding_box_to_original_image_shape( + box: torch.Tensor, + resized_image_shape: Tuple, + original_image_shape: Tuple, + ratio_pad: bool, +) -> List[int]: + """YOLOv8 resizes images during training and the label values are normalized based on this resized shape. + + This function rescales the bounding box labels to the original + image shape. + + Reference: https://github.com/ultralytics/ultralytics/blob/main/ultralytics/yolo/utils/callbacks/comet.py#L105 + """ + resized_image_height, resized_image_width = resized_image_shape + # Convert normalized xywh format predictions to xyxy in resized scale format + box = ops.xywhn2xyxy(box, h=resized_image_height, w=resized_image_width) + # Scale box predictions from resized image scale back to original image scale + box = ops.scale_boxes(resized_image_shape, box, original_image_shape, ratio_pad) + # # Convert bounding box format from xyxy to xywh for Comet logging + box = ops.xyxy2xywh(box) + return box.tolist() + + +def get_ground_truth_bbox_annotations( + img_idx: int, image_path: str, batch: Dict, class_name_map: Dict = None +) -> List[Dict[str, Any]]: + """Get ground truth bounding box annotation data in the form required for `wandb.Image` overlay system.""" + indices = batch["batch_idx"] == img_idx + bboxes = batch["bboxes"][indices] + if len(batch["cls"][indices]): + cls_labels = batch["cls"][indices].squeeze(1).tolist() + else: + cls_labels = [] + + class_name_map_reverse = {v: k for k, v in class_name_map.items()} + + if len(bboxes) == 0: + wandb.termwarn( + f"Image: {image_path} has no bounding boxes labels", repeat=False + ) + return None + + if len(batch["cls"][indices]): + cls_labels = batch["cls"][indices].squeeze(1).tolist() + else: + cls_labels = [] + + if class_name_map: + cls_labels = [str(class_name_map[label]) for label in cls_labels] + + original_image_shape = batch["ori_shape"][img_idx] + resized_image_shape = batch["resized_shape"][img_idx] + ratio_pad = batch["ratio_pad"][img_idx] + + data = [] + for box, label in zip(bboxes, cls_labels): + box = scale_bounding_box_to_original_image_shape( + box, resized_image_shape, original_image_shape, ratio_pad + ) + data.append( + { + "position": { + "middle": [int(box[0]), int(box[1])], + "width": int(box[2]), + "height": int(box[3]), + }, + "domain": "pixel", + "class_id": class_name_map_reverse[label], + "box_caption": label, + } + ) + + return data + + +def get_mean_confidence_map( + classes: List, confidence: List, class_id_to_label: Dict +) -> Dict[str, float]: + """Get Mean-confidence map from the predictions to be logged into a `wandb.Table`.""" + confidence_map = {v: [] for _, v in class_id_to_label.items()} + for class_idx, confidence_value in zip(classes, confidence): + confidence_map[class_id_to_label[class_idx]].append(confidence_value) + updated_confidence_map = {} + for label, confidence_list in confidence_map.items(): + if len(confidence_list) > 0: + updated_confidence_map[label] = sum(confidence_list) / len(confidence_list) + else: + updated_confidence_map[label] = 0 + return updated_confidence_map + + +def get_boxes(result: Results) -> Tuple[Dict, Dict]: + """Convert an ultralytics prediction result into metadata for the `wandb.Image` overlay system.""" + boxes = result.boxes.xywh.long().numpy() + classes = result.boxes.cls.long().numpy() + confidence = result.boxes.conf.numpy() + class_id_to_label = {int(k): str(v) for k, v in result.names.items()} + mean_confidence_map = get_mean_confidence_map( + classes, confidence, class_id_to_label + ) + box_data = [] + for idx in range(len(boxes)): + box_data.append( + { + "position": { + "middle": [int(boxes[idx][0]), int(boxes[idx][1])], + "width": int(boxes[idx][2]), + "height": int(boxes[idx][3]), + }, + "domain": "pixel", + "class_id": int(classes[idx]), + "box_caption": class_id_to_label[int(classes[idx])], + "scores": {"confidence": float(confidence[idx])}, + } + ) + boxes = { + "predictions": { + "box_data": box_data, + "class_labels": class_id_to_label, + }, + } + return boxes, mean_confidence_map + + +def plot_bbox_predictions( + result: Results, model_name: str, table: Optional[wandb.Table] = None +) -> Union[wandb.Table, Tuple[wandb.Image, Dict, Dict]]: + """Plot the images with the W&B overlay system. + + The `wandb.Image` is either added to a `wandb.Table` or returned. + """ + result = result.to("cpu") + boxes, mean_confidence_map = get_boxes(result) + image = wandb.Image(result.orig_img[:, :, ::-1], boxes=boxes) + if table is not None: + table.add_data( + model_name, + image, + len(boxes["predictions"]["box_data"]), + mean_confidence_map, + result.speed, + ) + return table + return image, boxes["predictions"], mean_confidence_map + + +def plot_detection_validation_results( + dataloader: Any, + class_label_map: Dict, + model_name: str, + predictor: DetectionPredictor, + table: wandb.Table, + max_validation_batches: int, + epoch: Optional[int] = None, +) -> wandb.Table: + """Plot validation results in a table.""" + data_idx = 0 + num_dataloader_batches = len(dataloader.dataset) // dataloader.batch_size + max_validation_batches = min(max_validation_batches, num_dataloader_batches) + for batch_idx, batch in enumerate(dataloader): + prediction_results = predictor(batch["im_file"]) + progress_bar_result_iterable = tqdm( + enumerate(prediction_results), + total=len(prediction_results), + desc=f"Generating Visualizations for batch-{batch_idx + 1}/{max_validation_batches}", + ) + for img_idx, prediction_result in progress_bar_result_iterable: + prediction_result = prediction_result.to("cpu") + _, prediction_box_data, mean_confidence_map = plot_bbox_predictions( + prediction_result, model_name + ) + try: + ground_truth_data = get_ground_truth_bbox_annotations( + img_idx, batch["im_file"][img_idx], batch, class_label_map + ) + wandb_image = wandb.Image( + batch["im_file"][img_idx], + boxes={ + "ground-truth": { + "box_data": ground_truth_data, + "class_labels": class_label_map, + }, + "predictions": { + "box_data": prediction_box_data["box_data"], + "class_labels": class_label_map, + }, + }, + ) + table_rows = [ + data_idx, + batch_idx, + wandb_image, + mean_confidence_map, + prediction_result.speed, + ] + table_rows = [epoch] + table_rows if epoch is not None else table_rows + table_rows = [model_name] + table_rows + table.add_data(*table_rows) + data_idx += 1 + except TypeError: + pass + if batch_idx + 1 == max_validation_batches: + break + return table diff --git a/lib/python3.12/site-packages/wandb/integration/ultralytics/callback.py b/lib/python3.12/site-packages/wandb/integration/ultralytics/callback.py new file mode 100644 index 0000000000000000000000000000000000000000..c47605a9579dafece3ef364bd3916f4cfd3c96f9 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/ultralytics/callback.py @@ -0,0 +1,528 @@ +import copy +from datetime import datetime +from typing import Callable, Dict, Optional, Union + +from packaging import version + +try: + import dill as pickle +except ImportError: + import pickle + +import wandb +from wandb.sdk.lib import telemetry + +try: + import torch + import ultralytics + from tqdm.auto import tqdm + + if version.parse(ultralytics.__version__) > version.parse("8.0.238"): + wandb.termwarn( + """This integration is tested and supported for ultralytics v8.0.238 and below. + Please report any issues to https://github.com/wandb/wandb/issues with the tag `yolov8`.""", + repeat=False, + ) + + from ultralytics.models import YOLO + from ultralytics.models.sam.predict import Predictor as SAMPredictor + from ultralytics.models.yolo.classify import ( + ClassificationPredictor, + ClassificationTrainer, + ClassificationValidator, + ) + from ultralytics.models.yolo.detect import ( + DetectionPredictor, + DetectionTrainer, + DetectionValidator, + ) + from ultralytics.models.yolo.pose import PosePredictor, PoseTrainer, PoseValidator + from ultralytics.models.yolo.segment import ( + SegmentationPredictor, + SegmentationTrainer, + SegmentationValidator, + ) + from ultralytics.utils.torch_utils import de_parallel + + try: + from ultralytics.yolo.utils import RANK, __version__ + except ModuleNotFoundError: + from ultralytics.utils import RANK, __version__ + + from wandb.integration.ultralytics.bbox_utils import ( + plot_bbox_predictions, + plot_detection_validation_results, + ) + from wandb.integration.ultralytics.classification_utils import ( + plot_classification_predictions, + plot_classification_validation_results, + ) + from wandb.integration.ultralytics.mask_utils import ( + plot_mask_predictions, + plot_sam_predictions, + plot_segmentation_validation_results, + ) + from wandb.integration.ultralytics.pose_utils import ( + plot_pose_predictions, + plot_pose_validation_results, + ) +except Exception as e: + wandb.Error(e) + + +TRAINER_TYPE = Union[ + ClassificationTrainer, DetectionTrainer, SegmentationTrainer, PoseTrainer +] +VALIDATOR_TYPE = Union[ + ClassificationValidator, DetectionValidator, SegmentationValidator, PoseValidator +] +PREDICTOR_TYPE = Union[ + ClassificationPredictor, + DetectionPredictor, + SegmentationPredictor, + PosePredictor, + SAMPredictor, +] + + +class WandBUltralyticsCallback: + """Stateful callback for logging to W&B. + + In particular, it will log model checkpoints, predictions, and + ground-truth annotations with interactive overlays for bounding boxes + to Weights & Biases Tables during training, validation and prediction + for a `ultratytics` workflow. + + Example: + ```python + from ultralytics.yolo.engine.model import YOLO + from wandb.yolov8 import add_wandb_callback + + # initialize YOLO model + model = YOLO("yolov8n.pt") + + # add wandb callback + add_wandb_callback( + model, max_validation_batches=2, enable_model_checkpointing=True + ) + + # train + model.train(data="coco128.yaml", epochs=5, imgsz=640) + + # validate + model.val() + + # perform inference + model(["img1.jpeg", "img2.jpeg"]) + ``` + + Args: + model: (ultralytics.yolo.engine.model.YOLO) YOLO Model of type + `ultralytics.yolo.engine.model.YOLO`. + epoch_logging_interval: (int) interval to log the prediction visualizations + during training. + max_validation_batches: (int) maximum number of validation batches to log to + a table per epoch. + enable_model_checkpointing: (bool) enable logging model checkpoints as + artifacts at the end of eveny epoch if set to `True`. + visualize_skeleton: (bool) visualize pose skeleton by drawing lines connecting + keypoints for human pose. + """ + + def __init__( + self, + model: YOLO, + epoch_logging_interval: int = 1, + max_validation_batches: int = 1, + enable_model_checkpointing: bool = False, + visualize_skeleton: bool = False, + ) -> None: + self.epoch_logging_interval = epoch_logging_interval + self.max_validation_batches = max_validation_batches + self.enable_model_checkpointing = enable_model_checkpointing + self.visualize_skeleton = visualize_skeleton + self.task = model.task + self.task_map = model.task_map + self.model_name = ( + model.overrides["model"].split(".")[0] + if "model" in model.overrides + else None + ) + self._make_tables() + self._make_predictor(model) + self.supported_tasks = ["detect", "segment", "pose", "classify"] + self.prompts = None + self.run_id = None + self.train_epoch = None + + def _make_tables(self): + if self.task in ["detect", "segment"]: + validation_columns = [ + "Data-Index", + "Batch-Index", + "Image", + "Mean-Confidence", + "Speed", + ] + train_columns = ["Epoch"] + validation_columns + self.train_validation_table = wandb.Table( + columns=["Model-Name"] + train_columns + ) + self.validation_table = wandb.Table( + columns=["Model-Name"] + validation_columns + ) + self.prediction_table = wandb.Table( + columns=[ + "Model-Name", + "Image", + "Num-Objects", + "Mean-Confidence", + "Speed", + ] + ) + elif self.task == "classify": + classification_columns = [ + "Image", + "Predicted-Category", + "Prediction-Confidence", + "Top-5-Prediction-Categories", + "Top-5-Prediction-Confindence", + "Probabilities", + "Speed", + ] + validation_columns = ["Data-Index", "Batch-Index"] + classification_columns + validation_columns.insert(3, "Ground-Truth-Category") + self.train_validation_table = wandb.Table( + columns=["Model-Name", "Epoch"] + validation_columns + ) + self.validation_table = wandb.Table( + columns=["Model-Name"] + validation_columns + ) + self.prediction_table = wandb.Table( + columns=["Model-Name"] + classification_columns + ) + elif self.task == "pose": + validation_columns = [ + "Data-Index", + "Batch-Index", + "Image-Ground-Truth", + "Image-Prediction", + "Num-Instances", + "Mean-Confidence", + "Speed", + ] + train_columns = ["Epoch"] + validation_columns + self.train_validation_table = wandb.Table( + columns=["Model-Name"] + train_columns + ) + self.validation_table = wandb.Table( + columns=["Model-Name"] + validation_columns + ) + self.prediction_table = wandb.Table( + columns=[ + "Model-Name", + "Image-Prediction", + "Num-Instances", + "Mean-Confidence", + "Speed", + ] + ) + + def _make_predictor(self, model: YOLO): + overrides = copy.deepcopy(model.overrides) + overrides["conf"] = 0.1 + self.predictor = self.task_map[self.task]["predictor"](overrides=overrides) + self.predictor.callbacks = {} + self.predictor.args.save = False + self.predictor.args.save_txt = False + self.predictor.args.save_crop = False + self.predictor.args.verbose = None + + def _save_model(self, trainer: TRAINER_TYPE): + model_checkpoint_artifact = wandb.Artifact(f"run_{wandb.run.id}_model", "model") + checkpoint_dict = { + "epoch": trainer.epoch, + "best_fitness": trainer.best_fitness, + "model": copy.deepcopy(de_parallel(self.model)).half(), + "ema": copy.deepcopy(trainer.ema.ema).half(), + "updates": trainer.ema.updates, + "optimizer": trainer.optimizer.state_dict(), + "train_args": vars(trainer.args), + "date": datetime.now().isoformat(), + "version": __version__, + } + checkpoint_path = trainer.wdir / f"epoch{trainer.epoch}.pt" + torch.save(checkpoint_dict, checkpoint_path, pickle_module=pickle) + model_checkpoint_artifact.add_file(checkpoint_path) + wandb.log_artifact( + model_checkpoint_artifact, aliases=[f"epoch_{trainer.epoch}"] + ) + + def on_train_start(self, trainer: TRAINER_TYPE): + with telemetry.context(run=wandb.run) as tel: + tel.feature.ultralytics_yolov8 = True + wandb.config.train = vars(trainer.args) + self.run_id = wandb.run.id + + @torch.no_grad() + def on_fit_epoch_end(self, trainer: DetectionTrainer): + if self.task in self.supported_tasks and self.train_epoch != trainer.epoch: + self.train_epoch = trainer.epoch + if (self.train_epoch + 1) % self.epoch_logging_interval == 0: + validator = trainer.validator + dataloader = validator.dataloader + class_label_map = validator.names + self.device = next(trainer.model.parameters()).device + if isinstance(trainer.model, torch.nn.parallel.DistributedDataParallel): + model = trainer.model.module + else: + model = trainer.model + self.model = copy.deepcopy(model).eval().to(self.device) + self.predictor.setup_model(model=self.model, verbose=False) + if self.task == "pose": + self.train_validation_table = plot_pose_validation_results( + dataloader=dataloader, + class_label_map=class_label_map, + model_name=self.model_name, + predictor=self.predictor, + visualize_skeleton=self.visualize_skeleton, + table=self.train_validation_table, + max_validation_batches=self.max_validation_batches, + epoch=trainer.epoch, + ) + elif self.task == "segment": + self.train_validation_table = plot_segmentation_validation_results( + dataloader=dataloader, + class_label_map=class_label_map, + model_name=self.model_name, + predictor=self.predictor, + table=self.train_validation_table, + max_validation_batches=self.max_validation_batches, + epoch=trainer.epoch, + ) + elif self.task == "detect": + self.train_validation_table = plot_detection_validation_results( + dataloader=dataloader, + class_label_map=class_label_map, + model_name=self.model_name, + predictor=self.predictor, + table=self.train_validation_table, + max_validation_batches=self.max_validation_batches, + epoch=trainer.epoch, + ) + elif self.task == "classify": + self.train_validation_table = ( + plot_classification_validation_results( + dataloader=dataloader, + model_name=self.model_name, + predictor=self.predictor, + table=self.train_validation_table, + max_validation_batches=self.max_validation_batches, + epoch=trainer.epoch, + ) + ) + if self.enable_model_checkpointing: + self._save_model(trainer) + trainer.model.to(self.device) + + def on_train_end(self, trainer: TRAINER_TYPE): + if self.task in self.supported_tasks: + wandb.log({"Train-Table": self.train_validation_table}, commit=False) + + def on_val_start(self, validator: VALIDATOR_TYPE): + wandb.run or wandb.init( + project=validator.args.project or "YOLOv8", + job_type="validation_" + validator.args.task, + ) + + @torch.no_grad() + def on_val_end(self, trainer: VALIDATOR_TYPE): + if self.task in self.supported_tasks: + validator = trainer + dataloader = validator.dataloader + class_label_map = validator.names + if self.task == "pose": + self.validation_table = plot_pose_validation_results( + dataloader=dataloader, + class_label_map=class_label_map, + model_name=self.model_name, + predictor=self.predictor, + visualize_skeleton=self.visualize_skeleton, + table=self.validation_table, + max_validation_batches=self.max_validation_batches, + ) + elif self.task == "segment": + self.validation_table = plot_segmentation_validation_results( + dataloader=dataloader, + class_label_map=class_label_map, + model_name=self.model_name, + predictor=self.predictor, + table=self.validation_table, + max_validation_batches=self.max_validation_batches, + ) + elif self.task == "detect": + self.validation_table = plot_detection_validation_results( + dataloader=dataloader, + class_label_map=class_label_map, + model_name=self.model_name, + predictor=self.predictor, + table=self.validation_table, + max_validation_batches=self.max_validation_batches, + ) + elif self.task == "classify": + self.validation_table = plot_classification_validation_results( + dataloader=dataloader, + model_name=self.model_name, + predictor=self.predictor, + table=self.validation_table, + max_validation_batches=self.max_validation_batches, + ) + wandb.log({"Validation-Table": self.validation_table}, commit=False) + + def on_predict_start(self, predictor: PREDICTOR_TYPE): + wandb.run or wandb.init( + project=predictor.args.project or "YOLOv8", + config=vars(predictor.args), + job_type="prediction_" + predictor.args.task, + ) + if isinstance(predictor, SAMPredictor): + self.prompts = copy.deepcopy(predictor.prompts) + self.prediction_table = wandb.Table(columns=["Image"]) + + def on_predict_end(self, predictor: PREDICTOR_TYPE): + wandb.config.prediction_configs = vars(predictor.args) + if self.task in self.supported_tasks: + for result in tqdm(predictor.results): + if self.task == "pose": + self.prediction_table = plot_pose_predictions( + result, + self.model_name, + self.visualize_skeleton, + self.prediction_table, + ) + elif self.task == "segment": + if isinstance(predictor, SegmentationPredictor): + self.prediction_table = plot_mask_predictions( + result, self.model_name, self.prediction_table + ) + elif isinstance(predictor, SAMPredictor): + self.prediction_table = plot_sam_predictions( + result, self.prompts, self.prediction_table + ) + elif self.task == "detect": + self.prediction_table = plot_bbox_predictions( + result, self.model_name, self.prediction_table + ) + elif self.task == "classify": + self.prediction_table = plot_classification_predictions( + result, self.model_name, self.prediction_table + ) + + wandb.log({"Prediction-Table": self.prediction_table}, commit=False) + + @property + def callbacks(self) -> Dict[str, Callable]: + """Property contains all the relevant callbacks to add to the YOLO model for the Weights & Biases logging.""" + return { + "on_train_start": self.on_train_start, + "on_fit_epoch_end": self.on_fit_epoch_end, + "on_train_end": self.on_train_end, + "on_val_start": self.on_val_start, + "on_val_end": self.on_val_end, + "on_predict_start": self.on_predict_start, + "on_predict_end": self.on_predict_end, + } + + +# TODO: Add epoch interval +def add_wandb_callback( + model: YOLO, + epoch_logging_interval: int = 1, + enable_model_checkpointing: bool = False, + enable_train_validation_logging: bool = True, + enable_validation_logging: bool = True, + enable_prediction_logging: bool = True, + max_validation_batches: Optional[int] = 1, + visualize_skeleton: Optional[bool] = True, +): + """Function to add the `WandBUltralyticsCallback` callback to the `YOLO` model. + + Example: + ```python + from ultralytics.yolo.engine.model import YOLO + from wandb.yolov8 import add_wandb_callback + + # initialize YOLO model + model = YOLO("yolov8n.pt") + + # add wandb callback + add_wandb_callback( + model, max_validation_batches=2, enable_model_checkpointing=True + ) + + # train + model.train(data="coco128.yaml", epochs=5, imgsz=640) + + # validate + model.val() + + # perform inference + model(["img1.jpeg", "img2.jpeg"]) + ``` + + Args: + model: (ultralytics.yolo.engine.model.YOLO) YOLO Model of type + `ultralytics.yolo.engine.model.YOLO`. + epoch_logging_interval: (int) interval to log the prediction visualizations + during training. + enable_model_checkpointing: (bool) enable logging model checkpoints as + artifacts at the end of eveny epoch if set to `True`. + enable_train_validation_logging: (bool) enable logging the predictions and + ground-truths as interactive image overlays on the images from + the validation dataloader to a `wandb.Table` along with + mean-confidence of the predictions per-class at the end of each + training epoch. + enable_validation_logging: (bool) enable logging the predictions and + ground-truths as interactive image overlays on the images from the + validation dataloader to a `wandb.Table` along with + mean-confidence of the predictions per-class at the end of + validation. + enable_prediction_logging: (bool) enable logging the predictions and + ground-truths as interactive image overlays on the images from the + validation dataloader to a `wandb.Table` along with mean-confidence + of the predictions per-class at the end of each prediction. + max_validation_batches: (Optional[int]) maximum number of validation batches to log to + a table per epoch. + visualize_skeleton: (Optional[bool]) visualize pose skeleton by drawing lines connecting + keypoints for human pose. + + Returns: + An instance of `ultralytics.yolo.engine.model.YOLO` with the `WandBUltralyticsCallback`. + """ + if RANK in [-1, 0]: + wandb_callback = WandBUltralyticsCallback( + copy.deepcopy(model), + epoch_logging_interval, + max_validation_batches, + enable_model_checkpointing, + visualize_skeleton, + ) + callbacks = wandb_callback.callbacks + if not enable_train_validation_logging: + _ = callbacks.pop("on_fit_epoch_end") + _ = callbacks.pop("on_train_end") + if not enable_validation_logging: + _ = callbacks.pop("on_val_start") + _ = callbacks.pop("on_val_end") + if not enable_prediction_logging: + _ = callbacks.pop("on_predict_start") + _ = callbacks.pop("on_predict_end") + for event, callback_fn in callbacks.items(): + model.add_callback(event, callback_fn) + else: + wandb.termerror( + "The RANK of the process to add the callbacks was neither 0 or " + "-1. No Weights & Biases callbacks were added to this instance " + "of the YOLO model." + ) + return model diff --git a/lib/python3.12/site-packages/wandb/integration/ultralytics/classification_utils.py b/lib/python3.12/site-packages/wandb/integration/ultralytics/classification_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9db6db72768dd2d16b7b44ee1ad0514a6322752f --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/ultralytics/classification_utils.py @@ -0,0 +1,83 @@ +from typing import Any, Optional + +import numpy as np +from tqdm.auto import tqdm +from ultralytics.engine.results import Results +from ultralytics.models.yolo.classify import ClassificationPredictor + +import wandb + + +def plot_classification_predictions( + result: Results, + model_name: str, + table: Optional[wandb.Table] = None, + original_image: Optional[np.array] = None, +): + """Plot classification prediction results to a `wandb.Table` if the table is passed otherwise return the data.""" + result = result.to("cpu") + probabilities = result.probs + probabilities_list = probabilities.data.numpy().tolist() + class_id_to_label = {int(k): str(v) for k, v in result.names.items()} + original_image = ( + wandb.Image(original_image) + if original_image is not None + else wandb.Image(result.orig_img) + ) + table_row = [ + model_name, + original_image, + class_id_to_label[int(probabilities.top1)], + probabilities.top1conf, + [class_id_to_label[int(class_idx)] for class_idx in list(probabilities.top5)], + [probabilities_list[int(class_idx)] for class_idx in list(probabilities.top5)], + { + class_id_to_label[int(class_idx)]: probability + for class_idx, probability in enumerate(probabilities_list) + }, + result.speed, + ] + if table is not None: + table.add_data(*table_row) + return table + return class_id_to_label, table_row + + +def plot_classification_validation_results( + dataloader: Any, + model_name: str, + predictor: ClassificationPredictor, + table: wandb.Table, + max_validation_batches: int, + epoch: Optional[int] = None, +) -> wandb.Table: + """Plot classification results to a `wandb.Table`.""" + data_idx = 0 + num_dataloader_batches = len(dataloader.dataset) // dataloader.batch_size + max_validation_batches = min(max_validation_batches, num_dataloader_batches) + for batch_idx, batch in enumerate(dataloader): + image_batch = [ + image for image in np.transpose(batch["img"].numpy(), (0, 2, 3, 1)) + ] + ground_truth = batch["cls"].numpy().tolist() + progress_bar_result_iterable = tqdm( + range(max_validation_batches), + desc=f"Generating Visualizations for batch-{batch_idx + 1}/{max_validation_batches}", + ) + for img_idx in progress_bar_result_iterable: + try: + prediction_result = predictor(image_batch[img_idx])[0] + class_id_to_label, table_row = plot_classification_predictions( + prediction_result, model_name, original_image=image_batch[img_idx] + ) + table_row = [data_idx, batch_idx] + table_row[1:] + table_row.insert(3, class_id_to_label[ground_truth[img_idx]]) + table_row = [epoch] + table_row if epoch is not None else table_row + table_row = [model_name] + table_row + table.add_data(*table_row) + data_idx += 1 + except Exception: + pass + if batch_idx + 1 == max_validation_batches: + break + return table diff --git a/lib/python3.12/site-packages/wandb/integration/ultralytics/mask_utils.py b/lib/python3.12/site-packages/wandb/integration/ultralytics/mask_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..5392d964b30ae1330349ce0978ef58eb3e5ee5dd --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/ultralytics/mask_utils.py @@ -0,0 +1,202 @@ +from typing import Dict, Optional, Tuple + +import cv2 +import numpy as np +from tqdm.auto import tqdm +from ultralytics.engine.results import Results +from ultralytics.models.yolo.segment import SegmentationPredictor +from ultralytics.utils.ops import scale_image + +import wandb +from wandb.integration.ultralytics.bbox_utils import ( + get_ground_truth_bbox_annotations, + get_mean_confidence_map, +) + + +def instance_mask_to_semantic_mask(instance_mask, class_indices): + height, width, num_instances = instance_mask.shape + semantic_mask = np.zeros((height, width), dtype=np.uint8) + for i in range(num_instances): + instance_map = instance_mask[:, :, i] + class_index = class_indices[i] + semantic_mask[instance_map == 1] = class_index + return semantic_mask + + +def get_boxes_and_masks(result: Results) -> Tuple[Dict, Dict, Dict]: + boxes = result.boxes.xywh.long().numpy() + classes = result.boxes.cls.long().numpy() + confidence = result.boxes.conf.numpy() + class_id_to_label = {int(k): str(v) for k, v in result.names.items()} + class_id_to_label.update({len(result.names.items()): "background"}) + mean_confidence_map = get_mean_confidence_map( + classes, confidence, class_id_to_label + ) + masks = None + if result.masks is not None: + scaled_instance_mask = scale_image( + np.transpose(result.masks.data.numpy(), (1, 2, 0)), + result.orig_img[:, :, ::-1].shape, + ) + scaled_semantic_mask = instance_mask_to_semantic_mask( + scaled_instance_mask, classes.tolist() + ) + scaled_semantic_mask[scaled_semantic_mask == 0] = len(result.names.items()) + masks = { + "predictions": { + "mask_data": scaled_semantic_mask, + "class_labels": class_id_to_label, + } + } + box_data, total_confidence = [], 0.0 + for idx in range(len(boxes)): + box_data.append( + { + "position": { + "middle": [int(boxes[idx][0]), int(boxes[idx][1])], + "width": int(boxes[idx][2]), + "height": int(boxes[idx][3]), + }, + "domain": "pixel", + "class_id": int(classes[idx]), + "box_caption": class_id_to_label[int(classes[idx])], + "scores": {"confidence": float(confidence[idx])}, + } + ) + total_confidence += float(confidence[idx]) + + boxes = { + "predictions": { + "box_data": box_data, + "class_labels": class_id_to_label, + }, + } + return boxes, masks, mean_confidence_map + + +def plot_mask_predictions( + result: Results, model_name: str, table: Optional[wandb.Table] = None +) -> Tuple[wandb.Image, Dict, Dict, Dict]: + result = result.to("cpu") + boxes, masks, mean_confidence_map = get_boxes_and_masks(result) + image = wandb.Image(result.orig_img[:, :, ::-1], boxes=boxes, masks=masks) + if table is not None: + table.add_data( + model_name, + image, + len(boxes["predictions"]["box_data"]), + mean_confidence_map, + result.speed, + ) + return table + return image, masks, boxes["predictions"], mean_confidence_map + + +def structure_prompts_and_image(image: np.array, prompt: Dict) -> Dict: + wb_box_data = [] + if prompt["bboxes"] is not None: + wb_box_data.append( + { + "position": { + "middle": [prompt["bboxes"][0], prompt["bboxes"][1]], + "width": prompt["bboxes"][2], + "height": prompt["bboxes"][3], + }, + "domain": "pixel", + "class_id": 1, + "box_caption": "Prompt-Box", + } + ) + if prompt["points"] is not None: + image = image.copy().astype(np.uint8) + image = cv2.circle( + image, tuple(prompt["points"]), 5, (0, 255, 0), -1, lineType=cv2.LINE_AA + ) + wb_box_data = { + "prompts": { + "box_data": wb_box_data, + "class_labels": {1: "Prompt-Box"}, + } + } + return image, wb_box_data + + +def plot_sam_predictions( + result: Results, prompt: Dict, table: wandb.Table +) -> wandb.Table: + result = result.to("cpu") + image = result.orig_img[:, :, ::-1] + image, wb_box_data = structure_prompts_and_image(image, prompt) + image = wandb.Image( + image, + boxes=wb_box_data, + masks={ + "predictions": { + "mask_data": np.squeeze(result.masks.data.cpu().numpy().astype(int)), + "class_labels": {0: "Background", 1: "Prediction"}, + } + }, + ) + table.add_data(image) + return table + + +def plot_segmentation_validation_results( + dataloader, + class_label_map, + model_name: str, + predictor: SegmentationPredictor, + table: wandb.Table, + max_validation_batches: int, + epoch: Optional[int] = None, +): + data_idx = 0 + num_dataloader_batches = len(dataloader.dataset) // dataloader.batch_size + max_validation_batches = min(max_validation_batches, num_dataloader_batches) + for batch_idx, batch in enumerate(dataloader): + prediction_results = predictor(batch["im_file"]) + progress_bar_result_iterable = tqdm( + enumerate(prediction_results), + total=len(prediction_results), + desc=f"Generating Visualizations for batch-{batch_idx + 1}/{max_validation_batches}", + ) + for img_idx, prediction_result in progress_bar_result_iterable: + prediction_result = prediction_result.to("cpu") + ( + _, + prediction_mask_data, + prediction_box_data, + mean_confidence_map, + ) = plot_mask_predictions(prediction_result, model_name) + try: + ground_truth_data = get_ground_truth_bbox_annotations( + img_idx, batch["im_file"][img_idx], batch, class_label_map + ) + wandb_image = wandb.Image( + batch["im_file"][img_idx], + boxes={ + "ground-truth": { + "box_data": ground_truth_data, + "class_labels": class_label_map, + }, + "predictions": prediction_box_data, + }, + masks=prediction_mask_data, + ) + table_rows = [ + data_idx, + batch_idx, + wandb_image, + mean_confidence_map, + prediction_result.speed, + ] + table_rows = [epoch] + table_rows if epoch is not None else table_rows + table_rows = [model_name] + table_rows + table.add_data(*table_rows) + data_idx += 1 + except TypeError: + pass + if batch_idx + 1 == max_validation_batches: + break + return table diff --git a/lib/python3.12/site-packages/wandb/integration/ultralytics/pose_utils.py b/lib/python3.12/site-packages/wandb/integration/ultralytics/pose_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..4cef5a90353d063e281ae54836d76675029aaa4e --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/ultralytics/pose_utils.py @@ -0,0 +1,103 @@ +from typing import Any, Optional + +import numpy as np +from PIL import Image +from tqdm.auto import tqdm +from ultralytics.engine.results import Results +from ultralytics.models.yolo.pose import PosePredictor +from ultralytics.utils.plotting import Annotator + +import wandb +from wandb.integration.ultralytics.bbox_utils import ( + get_boxes, + get_ground_truth_bbox_annotations, +) + + +def annotate_keypoint_results(result: Results, visualize_skeleton: bool): + annotator = Annotator(np.ascontiguousarray(result.orig_img[:, :, ::-1])) + key_points = result.keypoints.data.numpy() + for idx in range(key_points.shape[0]): + annotator.kpts(key_points[idx], kpt_line=visualize_skeleton) + return annotator.im + + +def annotate_keypoint_batch(image_path: str, keypoints: Any, visualize_skeleton: bool): + with Image.open(image_path) as original_image: + original_image = np.ascontiguousarray(original_image) + annotator = Annotator(original_image) + annotator.kpts(keypoints.numpy(), kpt_line=visualize_skeleton) + return annotator.im + + +def plot_pose_predictions( + result: Results, + model_name: str, + visualize_skeleton: bool, + table: Optional[wandb.Table] = None, +): + result = result.to("cpu") + boxes, mean_confidence_map = get_boxes(result) + annotated_image = annotate_keypoint_results(result, visualize_skeleton) + prediction_image = wandb.Image(annotated_image, boxes=boxes) + table_row = [ + model_name, + prediction_image, + len(boxes["predictions"]["box_data"]), + mean_confidence_map, + result.speed, + ] + if table is not None: + table.add_data(*table_row) + return table + return table_row + + +def plot_pose_validation_results( + dataloader, + class_label_map, + model_name: str, + predictor: PosePredictor, + visualize_skeleton: bool, + table: wandb.Table, + max_validation_batches: int, + epoch: Optional[int] = None, +) -> wandb.Table: + data_idx = 0 + num_dataloader_batches = len(dataloader.dataset) // dataloader.batch_size + max_validation_batches = min(max_validation_batches, num_dataloader_batches) + for batch_idx, batch in enumerate(dataloader): + prediction_results = predictor(batch["im_file"]) + progress_bar_result_iterable = tqdm( + enumerate(prediction_results), + total=len(prediction_results), + desc=f"Generating Visualizations for batch-{batch_idx + 1}/{max_validation_batches}", + ) + for img_idx, prediction_result in progress_bar_result_iterable: + prediction_result = prediction_result.to("cpu") + table_row = plot_pose_predictions( + prediction_result, model_name, visualize_skeleton + ) + ground_truth_image = wandb.Image( + annotate_keypoint_batch( + batch["im_file"][img_idx], + batch["keypoints"][img_idx], + visualize_skeleton, + ), + boxes={ + "ground-truth": { + "box_data": get_ground_truth_bbox_annotations( + img_idx, batch["im_file"][img_idx], batch, class_label_map + ), + "class_labels": class_label_map, + }, + }, + ) + table_row = [data_idx, batch_idx, ground_truth_image] + table_row[1:] + table_row = [epoch] + table_row if epoch is not None else table_row + table_row = [model_name] + table_row + table.add_data(*table_row) + data_idx += 1 + if batch_idx + 1 == max_validation_batches: + break + return table diff --git a/lib/python3.12/site-packages/wandb/integration/xgboost/__init__.py b/lib/python3.12/site-packages/wandb/integration/xgboost/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..052083e0cf7034a469a82582422ff7fee62965a5 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/xgboost/__init__.py @@ -0,0 +1,11 @@ +"""W&B callback for xgboost. + +Simple callback to get logging for each tree + +Use the `wandb_callback` to add `wandb` logging to any `XGboost` model. However, it will +be deprecated in favor of WandbCallback. Use it instead for more features. +""" + +from .xgboost import WandbCallback, wandb_callback + +__all__ = ["wandb_callback", "WandbCallback"] diff --git a/lib/python3.12/site-packages/wandb/integration/xgboost/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/xgboost/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ad31c44bb8af3f55c7d00ecad55f2ed6596505d9 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/xgboost/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/xgboost/__pycache__/xgboost.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/xgboost/__pycache__/xgboost.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a473ea36dca288edc09a58e79b7d3808d90c2164 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/xgboost/__pycache__/xgboost.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/xgboost/xgboost.py b/lib/python3.12/site-packages/wandb/integration/xgboost/xgboost.py new file mode 100644 index 0000000000000000000000000000000000000000..2a8b63dfa2f39db9d40989dc8562cc732ba11d98 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/xgboost/xgboost.py @@ -0,0 +1,189 @@ +"""xgboost init!""" + +import json +import warnings +from pathlib import Path +from typing import TYPE_CHECKING, cast + +import xgboost as xgb # type: ignore +from xgboost import Booster + +import wandb +from wandb.sdk.lib import telemetry as wb_telemetry + +MINIMIZE_METRICS = [ + "rmse", + "rmsle", + "mae", + "mape", + "mphe", + "logloss", + "error", + "error@t", + "merror", +] + +MAXIMIZE_METRICS = ["auc", "aucpr", "ndcg", "map", "ndcg@n", "map@n"] + + +if TYPE_CHECKING: + from typing import Callable, List, NamedTuple + + class CallbackEnv(NamedTuple): + evaluation_result_list: List + + +def wandb_callback() -> "Callable": + """Old style callback that will be deprecated in favor of WandbCallback. Please try the new logger for more features.""" + warnings.warn( + "wandb_callback will be deprecated in favor of WandbCallback. Please use WandbCallback for more features.", + UserWarning, + stacklevel=2, + ) + + with wb_telemetry.context() as tel: + tel.feature.xgboost_old_wandb_callback = True + + def callback(env: "CallbackEnv") -> None: + for k, v in env.evaluation_result_list: + wandb.log({k: v}, commit=False) + wandb.log({}) + + return callback + + +class WandbCallback(xgb.callback.TrainingCallback): + """`WandbCallback` automatically integrates XGBoost with wandb. + + Args: + log_model: (boolean) if True save and upload the model to Weights & Biases Artifacts + log_feature_importance: (boolean) if True log a feature importance bar plot + importance_type: (str) one of {weight, gain, cover, total_gain, total_cover} for tree model. weight for linear model. + define_metric: (boolean) if True (default) capture model performance at the best step, instead of the last step, of training in your `wandb.summary`. + + Passing `WandbCallback` to XGBoost will: + + - log the booster model configuration to Weights & Biases + - log evaluation metrics collected by XGBoost, such as rmse, accuracy etc. to Weights & Biases + - log training metric collected by XGBoost (if you provide training data to eval_set) + - log the best score and the best iteration + - save and upload your trained model to Weights & Biases Artifacts (when `log_model = True`) + - log feature importance plot when `log_feature_importance=True` (default). + - Capture the best eval metric in `wandb.summary` when `define_metric=True` (default). + + Example: + ```python + bst_params = dict( + objective="reg:squarederror", + colsample_bytree=0.3, + learning_rate=0.1, + max_depth=5, + alpha=10, + n_estimators=10, + tree_method="hist", + callbacks=[WandbCallback()], + ) + + xg_reg = xgb.XGBRegressor(**bst_params) + xg_reg.fit( + X_train, + y_train, + eval_set=[(X_test, y_test)], + ) + ``` + """ + + def __init__( + self, + log_model: bool = False, + log_feature_importance: bool = True, + importance_type: str = "gain", + define_metric: bool = True, + ): + self.log_model: bool = log_model + self.log_feature_importance: bool = log_feature_importance + self.importance_type: str = importance_type + self.define_metric: bool = define_metric + + if wandb.run is None: + raise wandb.Error("You must call wandb.init() before WandbCallback()") + + with wb_telemetry.context() as tel: + tel.feature.xgboost_wandb_callback = True + + def before_training(self, model: Booster) -> Booster: + """Run before training is finished.""" + # Update W&B config + config = model.save_config() + wandb.config.update(json.loads(config)) + + return model + + def after_training(self, model: Booster) -> Booster: + """Run after training is finished.""" + # Log the booster model as artifacts + if self.log_model: + self._log_model_as_artifact(model) + + # Plot feature importance + if self.log_feature_importance: + self._log_feature_importance(model) + + # Log the best score and best iteration + if model.attr("best_score") is not None: + wandb.log( + { + "best_score": float(cast(str, model.attr("best_score"))), + "best_iteration": int(cast(str, model.attr("best_iteration"))), + } + ) + + return model + + def after_iteration(self, model: Booster, epoch: int, evals_log: dict) -> bool: + """Run after each iteration. Return True when training should stop.""" + # Log metrics + for data, metric in evals_log.items(): + for metric_name, log in metric.items(): + if self.define_metric: + self._define_metric(data, metric_name) + wandb.log({f"{data}-{metric_name}": log[-1]}, commit=False) + else: + wandb.log({f"{data}-{metric_name}": log[-1]}, commit=False) + + wandb.log({"epoch": epoch}) + + self.define_metric = False + + return False + + def _log_model_as_artifact(self, model: Booster) -> None: + model_name = f"{wandb.run.id}_model.json" # type: ignore + model_path = Path(wandb.run.dir) / model_name # type: ignore + model.save_model(str(model_path)) + + model_artifact = wandb.Artifact(name=model_name, type="model") + model_artifact.add_file(str(model_path)) + wandb.log_artifact(model_artifact) + + def _log_feature_importance(self, model: Booster) -> None: + fi = model.get_score(importance_type=self.importance_type) + fi_data = [[k, fi[k]] for k in fi] + table = wandb.Table(data=fi_data, columns=["Feature", "Importance"]) + wandb.log( + { + "Feature Importance": wandb.plot.bar( + table, "Feature", "Importance", title="Feature Importance" + ) + } + ) + + def _define_metric(self, data: str, metric_name: str) -> None: + if "loss" in str.lower(metric_name): + wandb.define_metric(f"{data}-{metric_name}", summary="min") + elif str.lower(metric_name) in MINIMIZE_METRICS: + wandb.define_metric(f"{data}-{metric_name}", summary="min") + elif str.lower(metric_name) in MAXIMIZE_METRICS: + wandb.define_metric(f"{data}-{metric_name}", summary="max") + else: + pass diff --git a/lib/python3.12/site-packages/wandb/integration/yolov8/__init__.py b/lib/python3.12/site-packages/wandb/integration/yolov8/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/wandb/integration/yolov8/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/yolov8/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..905e9eeda0d84eda17bdddb315820c093c48f3fa Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/yolov8/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/yolov8/__pycache__/yolov8.cpython-312.pyc b/lib/python3.12/site-packages/wandb/integration/yolov8/__pycache__/yolov8.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c0408b319fd8be70e714e505b42e41d00832b407 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/integration/yolov8/__pycache__/yolov8.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/integration/yolov8/yolov8.py b/lib/python3.12/site-packages/wandb/integration/yolov8/yolov8.py new file mode 100644 index 0000000000000000000000000000000000000000..1c5dbd0fe8de39f387b6427b4696f28f9c029b44 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/integration/yolov8/yolov8.py @@ -0,0 +1,284 @@ +from typing import Any, Callable, Dict, List, Optional + +from ultralytics.yolo.engine.model import YOLO +from ultralytics.yolo.engine.trainer import BaseTrainer + +try: + from ultralytics.yolo.utils import RANK + from ultralytics.yolo.utils.torch_utils import get_flops, get_num_params +except ModuleNotFoundError: + from ultralytics.utils import RANK + from ultralytics.utils.torch_utils import get_flops, get_num_params +from ultralytics.yolo.v8.classify.train import ClassificationTrainer + +import wandb +from wandb.sdk.lib import telemetry + + +class WandbCallback: + """An internal YOLO model wrapper that tracks metrics, and logs models to Weights & Biases. + + Usage: + ```python + from wandb.integration.yolov8.yolov8 import WandbCallback + + model = YOLO("yolov8n.pt") + wandb_logger = WandbCallback( + model, + ) + for event, callback_fn in wandb_logger.callbacks.items(): + model.add_callback(event, callback_fn) + ``` + """ + + def __init__( + self, + yolo: YOLO, + run_name: Optional[str] = None, + project: Optional[str] = None, + tags: Optional[List[str]] = None, + resume: Optional[str] = None, + **kwargs: Optional[Any], + ) -> None: + """A utility class to manage wandb run and various callbacks for the ultralytics YOLOv8 framework. + + Args: + yolo: A YOLOv8 model that's inherited from `:class:ultralytics.yolo.engine.model.YOLO` + run_name, str: The name of the Weights & Biases run, defaults to an auto generated run_name if `trainer.args.name` is not defined. + project, str: The name of the Weights & Biases project, defaults to `"YOLOv8"` if `trainer.args.project` is not defined. + tags, List[str]: A list of tags to be added to the Weights & Biases run, defaults to `["YOLOv8"]`. + resume, str: Whether to resume a previous run on Weights & Biases, defaults to `None`. + **kwargs: Additional arguments to be passed to `wandb.init()`. + """ + self.yolo = yolo + self.run_name = run_name + self.project = project + self.tags = tags + self.resume = resume + self.kwargs = kwargs + + def on_pretrain_routine_start(self, trainer: BaseTrainer) -> None: + """Starts a new wandb run to track the training process and log to Weights & Biases. + + Args: + trainer: A task trainer that's inherited from `:class:ultralytics.yolo.engine.trainer.BaseTrainer` + that contains the model training and optimization routine. + """ + if wandb.run is None: + self.run = wandb.init( + name=self.run_name if self.run_name else trainer.args.name, + project=self.project + if self.project + else trainer.args.project or "YOLOv8", + tags=self.tags if self.tags else ["YOLOv8"], + config=vars(trainer.args), + resume=self.resume if self.resume else None, + **self.kwargs, + ) + else: + self.run = wandb.run + assert self.run is not None + self.run.define_metric("epoch", hidden=True) + self.run.define_metric( + "train/*", step_metric="epoch", step_sync=True, summary="min" + ) + + self.run.define_metric( + "val/*", step_metric="epoch", step_sync=True, summary="min" + ) + + self.run.define_metric( + "metrics/*", step_metric="epoch", step_sync=True, summary="max" + ) + self.run.define_metric( + "lr/*", step_metric="epoch", step_sync=True, summary="last" + ) + + with telemetry.context(run=wandb.run) as tel: + tel.feature.ultralytics_yolov8 = True + + def on_pretrain_routine_end(self, trainer: BaseTrainer) -> None: + assert self.run is not None + self.run.summary.update( + { + "model/parameters": get_num_params(trainer.model), + "model/GFLOPs": round(get_flops(trainer.model), 3), + } + ) + + def on_train_epoch_start(self, trainer: BaseTrainer) -> None: + """On train epoch start we only log epoch number to the Weights & Biases run.""" + # We log the epoch number here to commit the previous step, + assert self.run is not None + self.run.log({"epoch": trainer.epoch + 1}) + + def on_train_epoch_end(self, trainer: BaseTrainer) -> None: + """On train epoch end we log all the metrics to the Weights & Biases run.""" + assert self.run is not None + self.run.log( + { + **trainer.metrics, + **trainer.label_loss_items(trainer.tloss, prefix="train"), + **trainer.lr, + }, + ) + # Currently only the detection and segmentation trainers save images to the save_dir + if not isinstance(trainer, ClassificationTrainer): + self.run.log( + { + "train_batch_images": [ + wandb.Image(str(image_path), caption=image_path.stem) + for image_path in trainer.save_dir.glob("train_batch*.jpg") + ] + } + ) + + def on_fit_epoch_end(self, trainer: BaseTrainer) -> None: + """On fit epoch end we log all the best metrics and model detail to Weights & Biases run summary.""" + assert self.run is not None + if trainer.epoch == 0: + speeds = [ + trainer.validator.speed.get( + key, + ) + for key in (1, "inference") + ] + speed = speeds[0] if speeds[0] else speeds[1] + if speed: + self.run.summary.update( + { + "model/speed(ms/img)": round(speed, 3), + } + ) + if trainer.best_fitness == trainer.fitness: + self.run.summary.update( + { + "best/epoch": trainer.epoch + 1, + **{f"best/{key}": val for key, val in trainer.metrics.items()}, + } + ) + + def on_train_end(self, trainer: BaseTrainer) -> None: + """On train end we log all the media, including plots, images and best model artifact to Weights & Biases.""" + # Currently only the detection and segmentation trainers save images to the save_dir + assert self.run is not None + if not isinstance(trainer, ClassificationTrainer): + assert self.run is not None + self.run.log( + { + "plots": [ + wandb.Image(str(image_path), caption=image_path.stem) + for image_path in trainer.save_dir.glob("*.png") + ], + "val_images": [ + wandb.Image(str(image_path), caption=image_path.stem) + for image_path in trainer.validator.save_dir.glob("val*.jpg") + ], + }, + ) + + if trainer.best.exists(): + assert self.run is not None + self.run.log_artifact( + str(trainer.best), + type="model", + name=f"{self.run.name}_{trainer.args.task}.pt", + aliases=["best", f"epoch_{trainer.epoch + 1}"], + ) + + def on_model_save(self, trainer: BaseTrainer) -> None: + """On model save we log the model as an artifact to Weights & Biases.""" + assert self.run is not None + self.run.log_artifact( + str(trainer.last), + type="model", + name=f"{self.run.name}_{trainer.args.task}.pt", + aliases=["last", f"epoch_{trainer.epoch + 1}"], + ) + + def teardown(self, _trainer: BaseTrainer) -> None: + """On teardown, we finish the Weights & Biases run and set it to None.""" + assert self.run is not None + self.run.finish() + self.run = None + + @property + def callbacks( + self, + ) -> Dict[str, Callable]: + """Property contains all the relevant callbacks to add to the YOLO model for the Weights & Biases logging.""" + return { + "on_pretrain_routine_start": self.on_pretrain_routine_start, + "on_pretrain_routine_end": self.on_pretrain_routine_end, + "on_train_epoch_start": self.on_train_epoch_start, + "on_train_epoch_end": self.on_train_epoch_end, + "on_fit_epoch_end": self.on_fit_epoch_end, + "on_train_end": self.on_train_end, + "on_model_save": self.on_model_save, + "teardown": self.teardown, + } + + +def add_callbacks( + yolo: YOLO, + run_name: Optional[str] = None, + project: Optional[str] = None, + tags: Optional[List[str]] = None, + resume: Optional[str] = None, + **kwargs: Optional[Any], +) -> YOLO: + """A YOLO model wrapper that tracks metrics, and logs models to Weights & Biases. + + Args: + yolo: A YOLOv8 model that's inherited from `:class:ultralytics.yolo.engine.model.YOLO` + run_name, str: The name of the Weights & Biases run, defaults to an auto generated name if `trainer.args.name` is not defined. + project, str: The name of the Weights & Biases project, defaults to `"YOLOv8"` if `trainer.args.project` is not defined. + tags, List[str]: A list of tags to be added to the Weights & Biases run, defaults to `["YOLOv8"]`. + resume, str: Whether to resume a previous run on Weights & Biases, defaults to `None`. + **kwargs: Additional arguments to be passed to `wandb.init()`. + + Usage: + ```python + from wandb.integration.yolov8 import add_callbacks as add_wandb_callbacks + + model = YOLO("yolov8n.pt") + add_wandb_callbacks( + model, + ) + model.train( + data="coco128.yaml", + epochs=3, + imgsz=640, + ) + ``` + """ + wandb.termwarn( + """The wandb callback is currently in beta and is subject to change based on updates to `ultralytics yolov8`. + The callback is tested and supported for ultralytics v8.0.43 and above. + Please report any issues to https://github.com/wandb/wandb/issues with the tag `yolov8`. + """, + repeat=False, + ) + wandb.termwarn( + """This wandb callback is no longer functional and would be deprecated in the near future. + We recommend you to use the updated callback using `from wandb.integration.ultralytics import add_wandb_callback`. + The updated callback is tested and supported for ultralytics 8.0.167 and above. + You can refer to https://docs.wandb.ai/guides/integrations/ultralytics for the updated documentation. + Please report any issues to https://github.com/wandb/wandb/issues with the tag `yolov8`. + """, + repeat=False, + ) + + if RANK in [-1, 0]: + wandb_logger = WandbCallback( + yolo, run_name=run_name, project=project, tags=tags, resume=resume, **kwargs + ) + for event, callback_fn in wandb_logger.callbacks.items(): + yolo.add_callback(event, callback_fn) + return yolo + else: + wandb.termerror( + "The RANK of the process to add the callbacks was neither 0 or -1." + "No Weights & Biases callbacks were added to this instance of the YOLO model." + ) + return yolo diff --git a/lib/python3.12/site-packages/wandb/mpmain/__init__.py b/lib/python3.12/site-packages/wandb/mpmain/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/wandb/mpmain/__main__.py b/lib/python3.12/site-packages/wandb/mpmain/__main__.py new file mode 100644 index 0000000000000000000000000000000000000000..02e8a24d9fbce482e31cde126d8407586a61b3fd --- /dev/null +++ b/lib/python3.12/site-packages/wandb/mpmain/__main__.py @@ -0,0 +1 @@ +# This module is initialized after multiprocessing spawn diff --git a/lib/python3.12/site-packages/wandb/mpmain/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/mpmain/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1fb1bcf1bb7c49e869c6bc4eff17d08f7d5e3ed3 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/mpmain/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/mpmain/__pycache__/__main__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/mpmain/__pycache__/__main__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1c34b09b995a84a19ce15cae80616d42430e5db4 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/mpmain/__pycache__/__main__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/v5/__pycache__/wandb_base_pb2.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/v5/__pycache__/wandb_base_pb2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ed2314512459c912a6237d21bf140774b61371d8 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/v5/__pycache__/wandb_base_pb2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/v5/__pycache__/wandb_internal_pb2.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/v5/__pycache__/wandb_internal_pb2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c48ae7ffe68438086bff458249f858542c9f4d42 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/v5/__pycache__/wandb_internal_pb2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/v5/__pycache__/wandb_server_pb2.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/v5/__pycache__/wandb_server_pb2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..edf3b41efd527db4b79cb71ca0dbede4effc45c7 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/v5/__pycache__/wandb_server_pb2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/v5/__pycache__/wandb_settings_pb2.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/v5/__pycache__/wandb_settings_pb2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3c0959c8e394be1e2d71b887fcc66d16a61e7be6 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/v5/__pycache__/wandb_settings_pb2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/v5/__pycache__/wandb_telemetry_pb2.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/v5/__pycache__/wandb_telemetry_pb2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1c19d8831e6117015026af2d4dde45ce57e7ffde Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/v5/__pycache__/wandb_telemetry_pb2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/v5/wandb_server_pb2.py b/lib/python3.12/site-packages/wandb/proto/v5/wandb_server_pb2.py new file mode 100644 index 0000000000000000000000000000000000000000..d046eda055b092313965c9f8b1ef7ef94d8f70a6 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/proto/v5/wandb_server_pb2.py @@ -0,0 +1,68 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: wandb/proto/wandb_server.proto +# Protobuf Python Version: 5.26.1 +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import builder as _builder +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from wandb.proto import wandb_base_pb2 as wandb_dot_proto_dot_wandb__base__pb2 +from wandb.proto import wandb_internal_pb2 as wandb_dot_proto_dot_wandb__internal__pb2 +from wandb.proto import wandb_settings_pb2 as wandb_dot_proto_dot_wandb__settings__pb2 + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1ewandb/proto/wandb_server.proto\x12\x0ewandb_internal\x1a\x1cwandb/proto/wandb_base.proto\x1a wandb/proto/wandb_internal.proto\x1a wandb/proto/wandb_settings.proto\"k\n\x19ServerAuthenticateRequest\x12\x0f\n\x07\x61pi_key\x18\x01 \x01(\t\x12\x10\n\x08\x62\x61se_url\x18\x02 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"w\n\x1aServerAuthenticateResponse\x12\x16\n\x0e\x64\x65\x66\x61ult_entity\x18\x01 \x01(\t\x12\x14\n\x0c\x65rror_status\x18\x02 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"D\n\x15ServerShutdownRequest\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x18\n\x16ServerShutdownResponse\"B\n\x13ServerStatusRequest\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x16\n\x14ServerStatusResponse\"r\n\x17ServerInformInitRequest\x12*\n\x08settings\x18\x01 \x01(\x0b\x32\x18.wandb_internal.Settings\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x1a\n\x18ServerInformInitResponse\"s\n\x18ServerInformStartRequest\x12*\n\x08settings\x18\x01 \x01(\x0b\x32\x18.wandb_internal.Settings\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x1b\n\x19ServerInformStartResponse\"H\n\x19ServerInformFinishRequest\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x1c\n\x1aServerInformFinishResponse\"H\n\x19ServerInformAttachRequest\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"u\n\x1aServerInformAttachResponse\x12*\n\x08settings\x18\x01 \x01(\x0b\x32\x18.wandb_internal.Settings\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"H\n\x19ServerInformDetachRequest\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x1c\n\x1aServerInformDetachResponse\"]\n\x1bServerInformTeardownRequest\x12\x11\n\texit_code\x18\x01 \x01(\x05\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x1e\n\x1cServerInformTeardownResponse\"\xfb\x04\n\rServerRequest\x12\x12\n\nrequest_id\x18\n \x01(\t\x12\x30\n\x0erecord_publish\x18\x01 \x01(\x0b\x32\x16.wandb_internal.RecordH\x00\x12\x34\n\x12record_communicate\x18\x02 \x01(\x0b\x32\x16.wandb_internal.RecordH\x00\x12>\n\x0binform_init\x18\x03 \x01(\x0b\x32\'.wandb_internal.ServerInformInitRequestH\x00\x12\x42\n\rinform_finish\x18\x04 \x01(\x0b\x32).wandb_internal.ServerInformFinishRequestH\x00\x12\x42\n\rinform_attach\x18\x05 \x01(\x0b\x32).wandb_internal.ServerInformAttachRequestH\x00\x12\x42\n\rinform_detach\x18\x06 \x01(\x0b\x32).wandb_internal.ServerInformDetachRequestH\x00\x12\x46\n\x0finform_teardown\x18\x07 \x01(\x0b\x32+.wandb_internal.ServerInformTeardownRequestH\x00\x12@\n\x0cinform_start\x18\x08 \x01(\x0b\x32(.wandb_internal.ServerInformStartRequestH\x00\x12\x41\n\x0c\x61uthenticate\x18\t \x01(\x0b\x32).wandb_internal.ServerAuthenticateRequestH\x00\x42\x15\n\x13server_request_type\"\x91\x05\n\x0eServerResponse\x12\x12\n\nrequest_id\x18\n \x01(\t\x12\x34\n\x12result_communicate\x18\x02 \x01(\x0b\x32\x16.wandb_internal.ResultH\x00\x12H\n\x14inform_init_response\x18\x03 \x01(\x0b\x32(.wandb_internal.ServerInformInitResponseH\x00\x12L\n\x16inform_finish_response\x18\x04 \x01(\x0b\x32*.wandb_internal.ServerInformFinishResponseH\x00\x12L\n\x16inform_attach_response\x18\x05 \x01(\x0b\x32*.wandb_internal.ServerInformAttachResponseH\x00\x12L\n\x16inform_detach_response\x18\x06 \x01(\x0b\x32*.wandb_internal.ServerInformDetachResponseH\x00\x12P\n\x18inform_teardown_response\x18\x07 \x01(\x0b\x32,.wandb_internal.ServerInformTeardownResponseH\x00\x12J\n\x15inform_start_response\x18\x08 \x01(\x0b\x32).wandb_internal.ServerInformStartResponseH\x00\x12K\n\x15\x61uthenticate_response\x18\t \x01(\x0b\x32*.wandb_internal.ServerAuthenticateResponseH\x00\x42\x16\n\x14server_response_typeB\x1bZ\x19\x63ore/pkg/service_go_protob\x06proto3') + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'wandb.proto.wandb_server_pb2', _globals) +if not _descriptor._USE_C_DESCRIPTORS: + _globals['DESCRIPTOR']._loaded_options = None + _globals['DESCRIPTOR']._serialized_options = b'Z\031core/pkg/service_go_proto' + _globals['_SERVERAUTHENTICATEREQUEST']._serialized_start=148 + _globals['_SERVERAUTHENTICATEREQUEST']._serialized_end=255 + _globals['_SERVERAUTHENTICATERESPONSE']._serialized_start=257 + _globals['_SERVERAUTHENTICATERESPONSE']._serialized_end=376 + _globals['_SERVERSHUTDOWNREQUEST']._serialized_start=378 + _globals['_SERVERSHUTDOWNREQUEST']._serialized_end=446 + _globals['_SERVERSHUTDOWNRESPONSE']._serialized_start=448 + _globals['_SERVERSHUTDOWNRESPONSE']._serialized_end=472 + _globals['_SERVERSTATUSREQUEST']._serialized_start=474 + _globals['_SERVERSTATUSREQUEST']._serialized_end=540 + _globals['_SERVERSTATUSRESPONSE']._serialized_start=542 + _globals['_SERVERSTATUSRESPONSE']._serialized_end=564 + _globals['_SERVERINFORMINITREQUEST']._serialized_start=566 + _globals['_SERVERINFORMINITREQUEST']._serialized_end=680 + _globals['_SERVERINFORMINITRESPONSE']._serialized_start=682 + _globals['_SERVERINFORMINITRESPONSE']._serialized_end=708 + _globals['_SERVERINFORMSTARTREQUEST']._serialized_start=710 + _globals['_SERVERINFORMSTARTREQUEST']._serialized_end=825 + _globals['_SERVERINFORMSTARTRESPONSE']._serialized_start=827 + _globals['_SERVERINFORMSTARTRESPONSE']._serialized_end=854 + _globals['_SERVERINFORMFINISHREQUEST']._serialized_start=856 + _globals['_SERVERINFORMFINISHREQUEST']._serialized_end=928 + _globals['_SERVERINFORMFINISHRESPONSE']._serialized_start=930 + _globals['_SERVERINFORMFINISHRESPONSE']._serialized_end=958 + _globals['_SERVERINFORMATTACHREQUEST']._serialized_start=960 + _globals['_SERVERINFORMATTACHREQUEST']._serialized_end=1032 + _globals['_SERVERINFORMATTACHRESPONSE']._serialized_start=1034 + _globals['_SERVERINFORMATTACHRESPONSE']._serialized_end=1151 + _globals['_SERVERINFORMDETACHREQUEST']._serialized_start=1153 + _globals['_SERVERINFORMDETACHREQUEST']._serialized_end=1225 + _globals['_SERVERINFORMDETACHRESPONSE']._serialized_start=1227 + _globals['_SERVERINFORMDETACHRESPONSE']._serialized_end=1255 + _globals['_SERVERINFORMTEARDOWNREQUEST']._serialized_start=1257 + _globals['_SERVERINFORMTEARDOWNREQUEST']._serialized_end=1350 + _globals['_SERVERINFORMTEARDOWNRESPONSE']._serialized_start=1352 + _globals['_SERVERINFORMTEARDOWNRESPONSE']._serialized_end=1382 + _globals['_SERVERREQUEST']._serialized_start=1385 + _globals['_SERVERREQUEST']._serialized_end=2020 + _globals['_SERVERRESPONSE']._serialized_start=2023 + _globals['_SERVERRESPONSE']._serialized_end=2680 +# @@protoc_insertion_point(module_scope) diff --git a/lib/python3.12/site-packages/wandb/proto/v5/wandb_telemetry_pb2.py b/lib/python3.12/site-packages/wandb/proto/v5/wandb_telemetry_pb2.py new file mode 100644 index 0000000000000000000000000000000000000000..e626d85c05c9dc8cdab4f2f15b9c7b6298b889c9 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/proto/v5/wandb_telemetry_pb2.py @@ -0,0 +1,42 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: wandb/proto/wandb_telemetry.proto +# Protobuf Python Version: 5.26.1 +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import builder as _builder +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from wandb.proto import wandb_base_pb2 as wandb_dot_proto_dot_wandb__base__pb2 + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n!wandb/proto/wandb_telemetry.proto\x12\x0ewandb_internal\x1a\x1cwandb/proto/wandb_base.proto\"\xdb\x03\n\x0fTelemetryRecord\x12-\n\x0cimports_init\x18\x01 \x01(\x0b\x32\x17.wandb_internal.Imports\x12/\n\x0eimports_finish\x18\x02 \x01(\x0b\x32\x17.wandb_internal.Imports\x12(\n\x07\x66\x65\x61ture\x18\x03 \x01(\x0b\x32\x17.wandb_internal.Feature\x12\x16\n\x0epython_version\x18\x04 \x01(\t\x12\x13\n\x0b\x63li_version\x18\x05 \x01(\t\x12\x1b\n\x13huggingface_version\x18\x06 \x01(\t\x12 \n\x03\x65nv\x18\x08 \x01(\x0b\x32\x13.wandb_internal.Env\x12%\n\x05label\x18\t \x01(\x0b\x32\x16.wandb_internal.Labels\x12.\n\ndeprecated\x18\n \x01(\x0b\x32\x1a.wandb_internal.Deprecated\x12&\n\x06issues\x18\x0b \x01(\x0b\x32\x16.wandb_internal.Issues\x12\x14\n\x0c\x63ore_version\x18\x0c \x01(\t\x12\x10\n\x08platform\x18\r \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x11\n\x0fTelemetryResult\"\x9a\x0e\n\x07Imports\x12\r\n\x05torch\x18\x01 \x01(\x08\x12\r\n\x05keras\x18\x02 \x01(\x08\x12\x12\n\ntensorflow\x18\x03 \x01(\x08\x12\x0e\n\x06\x66\x61stai\x18\x04 \x01(\x08\x12\x0f\n\x07sklearn\x18\x05 \x01(\x08\x12\x0f\n\x07xgboost\x18\x06 \x01(\x08\x12\x10\n\x08\x63\x61tboost\x18\x07 \x01(\x08\x12\x10\n\x08lightgbm\x18\x08 \x01(\x08\x12\x19\n\x11pytorch_lightning\x18\t \x01(\x08\x12\x0e\n\x06ignite\x18\n \x01(\x08\x12\x14\n\x0ctransformers\x18\x0b \x01(\x08\x12\x0b\n\x03jax\x18\x0c \x01(\x08\x12\x10\n\x08metaflow\x18\r \x01(\x08\x12\x10\n\x08\x61llennlp\x18\x0e \x01(\x08\x12\x11\n\tautogluon\x18\x0f \x01(\x08\x12\x11\n\tautokeras\x18\x10 \x01(\x08\x12\x10\n\x08\x63\x61talyst\x18\x12 \x01(\x08\x12\x10\n\x08\x64\x65\x65pchem\x18\x15 \x01(\x08\x12\x0f\n\x07\x64\x65\x65pctr\x18\x16 \x01(\x08\x12\x0f\n\x07pycaret\x18\x1c \x01(\x08\x12\x14\n\x0cpytorchvideo\x18\x1d \x01(\x08\x12\x0b\n\x03ray\x18\x1e \x01(\x08\x12\x1a\n\x12simpletransformers\x18\x1f \x01(\x08\x12\x0e\n\x06skorch\x18 \x01(\x08\x12\r\n\x05spacy\x18! \x01(\x08\x12\r\n\x05\x66lash\x18\" \x01(\x08\x12\x0e\n\x06optuna\x18# \x01(\x08\x12\x0f\n\x07recbole\x18$ \x01(\x08\x12\x0c\n\x04mmcv\x18% \x01(\x08\x12\r\n\x05mmdet\x18& \x01(\x08\x12\x11\n\ttorchdrug\x18\' \x01(\x08\x12\x11\n\ttorchtext\x18( \x01(\x08\x12\x13\n\x0btorchvision\x18) \x01(\x08\x12\r\n\x05\x65legy\x18* \x01(\x08\x12\x12\n\ndetectron2\x18+ \x01(\x08\x12\r\n\x05\x66lair\x18, \x01(\x08\x12\x0c\n\x04\x66lax\x18- \x01(\x08\x12\x0c\n\x04syft\x18. \x01(\x08\x12\x0b\n\x03TTS\x18/ \x01(\x08\x12\r\n\x05monai\x18\x30 \x01(\x08\x12\x17\n\x0fhuggingface_hub\x18\x31 \x01(\x08\x12\r\n\x05hydra\x18\x32 \x01(\x08\x12\x10\n\x08\x64\x61tasets\x18\x33 \x01(\x08\x12\x0e\n\x06sacred\x18\x34 \x01(\x08\x12\x0e\n\x06joblib\x18\x35 \x01(\x08\x12\x0c\n\x04\x64\x61sk\x18\x36 \x01(\x08\x12\x11\n\tpaddleocr\x18\x38 \x01(\x08\x12\r\n\x05ppdet\x18\x39 \x01(\x08\x12\x11\n\tpaddleseg\x18: \x01(\x08\x12\x11\n\tpaddlenlp\x18; \x01(\x08\x12\r\n\x05mmseg\x18< \x01(\x08\x12\r\n\x05mmocr\x18= \x01(\x08\x12\r\n\x05mmcls\x18> \x01(\x08\x12\x0c\n\x04timm\x18? \x01(\x08\x12\x0f\n\x07\x66\x61irseq\x18@ \x01(\x08\x12\x12\n\ndeepchecks\x18\x41 \x01(\x08\x12\x10\n\x08\x63omposer\x18\x42 \x01(\x08\x12\x10\n\x08sparseml\x18\x43 \x01(\x08\x12\x10\n\x08\x61nomalib\x18\x44 \x01(\x08\x12\r\n\x05zenml\x18\x45 \x01(\x08\x12\x12\n\ncolossalai\x18\x46 \x01(\x08\x12\x12\n\naccelerate\x18G \x01(\x08\x12\x0e\n\x06merlin\x18H \x01(\x08\x12\x0f\n\x07nanodet\x18I \x01(\x08\x12#\n\x1bsegmentation_models_pytorch\x18J \x01(\x08\x12\x1d\n\x15sentence_transformers\x18K \x01(\x08\x12\x0b\n\x03\x64gl\x18L \x01(\x08\x12\x17\n\x0ftorch_geometric\x18M \x01(\x08\x12\x0c\n\x04jina\x18N \x01(\x08\x12\x0e\n\x06kornia\x18O \x01(\x08\x12\x16\n\x0e\x61lbumentations\x18P \x01(\x08\x12\x10\n\x08keras_cv\x18Q \x01(\x08\x12\x10\n\x08mmengine\x18R \x01(\x08\x12\x11\n\tdiffusers\x18S \x01(\x08\x12\x0b\n\x03trl\x18T \x01(\x08\x12\x0c\n\x04trlx\x18U \x01(\x08\x12\x11\n\tlangchain\x18V \x01(\x08\x12\x13\n\x0bllama_index\x18W \x01(\x08\x12\x15\n\rstability_sdk\x18X \x01(\x08\x12\x0f\n\x07prefect\x18Y \x01(\x08\x12\x13\n\x0bprefect_ray\x18Z \x01(\x08\x12\x10\n\x08pinecone\x18[ \x01(\x08\x12\x10\n\x08\x63hromadb\x18\\ \x01(\x08\x12\x10\n\x08weaviate\x18] \x01(\x08\x12\x13\n\x0bpromptlayer\x18^ \x01(\x08\x12\x0e\n\x06openai\x18_ \x01(\x08\x12\x0e\n\x06\x63ohere\x18` \x01(\x08\x12\x11\n\tanthropic\x18\x61 \x01(\x08\x12\x0c\n\x04peft\x18\x62 \x01(\x08\x12\x0f\n\x07optimum\x18\x63 \x01(\x08\x12\x10\n\x08\x65valuate\x18\x64 \x01(\x08\x12\x10\n\x08langflow\x18\x65 \x01(\x08\x12\x12\n\nkeras_core\x18\x66 \x01(\x08\x12\x18\n\x10lightning_fabric\x18g \x01(\x08\x12\x1c\n\x14\x63urated_transformers\x18h \x01(\x08\x12\x0e\n\x06orjson\x18i \x01(\x08\x12\x11\n\tlightning\x18j \x01(\x08J\x04\x08\x11\x10\x12J\x04\x08\x13\x10\x14J\x04\x08\x14\x10\x15J\x04\x08\x17\x10\x18J\x04\x08\x18\x10\x19J\x04\x08\x19\x10\x1aJ\x04\x08\x1a\x10\x1bJ\x04\x08\x1b\x10\x1cJ\x04\x08\x37\x10\x38\"\x95\r\n\x07\x46\x65\x61ture\x12\r\n\x05watch\x18\x01 \x01(\x08\x12\x0e\n\x06\x66inish\x18\x02 \x01(\x08\x12\x0c\n\x04save\x18\x03 \x01(\x08\x12\x0f\n\x07offline\x18\x04 \x01(\x08\x12\x0f\n\x07resumed\x18\x05 \x01(\x08\x12\x0c\n\x04grpc\x18\x06 \x01(\x08\x12\x0e\n\x06metric\x18\x07 \x01(\x08\x12\r\n\x05keras\x18\x08 \x01(\x08\x12\x11\n\tsagemaker\x18\t \x01(\x08\x12\x1c\n\x14\x61rtifact_incremental\x18\n \x01(\x08\x12\x10\n\x08metaflow\x18\x0b \x01(\x08\x12\x0f\n\x07prodigy\x18\x0c \x01(\x08\x12\x15\n\rset_init_name\x18\r \x01(\x08\x12\x13\n\x0bset_init_id\x18\x0e \x01(\x08\x12\x15\n\rset_init_tags\x18\x0f \x01(\x08\x12\x17\n\x0fset_init_config\x18\x10 \x01(\x08\x12\x14\n\x0cset_run_name\x18\x11 \x01(\x08\x12\x14\n\x0cset_run_tags\x18\x12 \x01(\x08\x12\x17\n\x0fset_config_item\x18\x13 \x01(\x08\x12\x0e\n\x06launch\x18\x14 \x01(\x08\x12\x1c\n\x14torch_profiler_trace\x18\x15 \x01(\x08\x12\x0b\n\x03sb3\x18\x16 \x01(\x08\x12\x17\n\x0finit_return_run\x18\x18 \x01(\x08\x12\x1f\n\x17lightgbm_wandb_callback\x18\x19 \x01(\x08\x12\x1c\n\x14lightgbm_log_summary\x18\x1a \x01(\x08\x12\x1f\n\x17\x63\x61tboost_wandb_callback\x18\x1b \x01(\x08\x12\x1c\n\x14\x63\x61tboost_log_summary\x18\x1c \x01(\x08\x12\x17\n\x0ftensorboard_log\x18\x1d \x01(\x08\x12\x16\n\x0e\x65stimator_hook\x18\x1e \x01(\x08\x12\x1e\n\x16xgboost_wandb_callback\x18\x1f \x01(\x08\x12\"\n\x1axgboost_old_wandb_callback\x18 \x01(\x08\x12\x0e\n\x06\x61ttach\x18! \x01(\x08\x12\x19\n\x11tensorboard_patch\x18\" \x01(\x08\x12\x18\n\x10tensorboard_sync\x18# \x01(\x08\x12\x15\n\rkfp_wandb_log\x18$ \x01(\x08\x12\x1b\n\x13maybe_run_overwrite\x18% \x01(\x08\x12\x1c\n\x14keras_metrics_logger\x18& \x01(\x08\x12\x1e\n\x16keras_model_checkpoint\x18\' \x01(\x08\x12!\n\x19keras_wandb_eval_callback\x18( \x01(\x08\x12\x1d\n\x15\x66low_control_overflow\x18) \x01(\x08\x12\x0c\n\x04sync\x18* \x01(\x08\x12\x1d\n\x15\x66low_control_disabled\x18+ \x01(\x08\x12\x1b\n\x13\x66low_control_custom\x18, \x01(\x08\x12\x14\n\x0copen_metrics\x18. \x01(\x08\x12\x1a\n\x12ultralytics_yolov8\x18/ \x01(\x08\x12\x17\n\x0fimporter_mlflow\x18\x30 \x01(\x08\x12\x15\n\rsync_tfevents\x18\x31 \x01(\x08\x12\x15\n\rasync_uploads\x18\x32 \x01(\x08\x12\x16\n\x0eopenai_autolog\x18\x33 \x01(\x08\x12\x18\n\x10langchain_tracer\x18\x34 \x01(\x08\x12\x16\n\x0e\x63ohere_autolog\x18\x35 \x01(\x08\x12\x1b\n\x13hf_pipeline_autolog\x18\x36 \x01(\x08\x12\r\n\x05lib_c\x18\x38 \x01(\x08\x12\x0f\n\x07lib_cpp\x18\x39 \x01(\x08\x12\x19\n\x11openai_finetuning\x18: \x01(\x08\x12\x19\n\x11\x64iffusers_autolog\x18; \x01(\x08\x12\x1f\n\x17lightning_fabric_logger\x18< \x01(\x08\x12\x14\n\x0cset_step_log\x18= \x01(\x08\x12\x13\n\x0bset_summary\x18> \x01(\x08\x12\x16\n\x0emetric_summary\x18? \x01(\x08\x12\x13\n\x0bmetric_goal\x18@ \x01(\x08\x12\x15\n\rmetric_hidden\x18\x41 \x01(\x08\x12\x18\n\x10metric_step_sync\x18\x42 \x01(\x08\x12\x13\n\x0bshared_mode\x18\x43 \x01(\x08\x12#\n\x1bserver_side_derived_summary\x18\x44 \x01(\x08\x12\x1b\n\x13user_provided_label\x18\x45 \x01(\x08\x12\x1e\n\x16\x64\x63gm_profiling_enabled\x18\x46 \x01(\x08\x12\x11\n\tfork_mode\x18G \x01(\x08\x12\x13\n\x0brewind_mode\x18H \x01(\x08J\x04\x08\x17\x10\x18J\x04\x08-\x10.J\x04\x08\x37\x10\x38\"\xc9\x01\n\x03\x45nv\x12\x0f\n\x07jupyter\x18\x01 \x01(\x08\x12\x0e\n\x06kaggle\x18\x02 \x01(\x08\x12\x0f\n\x07windows\x18\x03 \x01(\x08\x12\x0e\n\x06m1_gpu\x18\x04 \x01(\x08\x12\x10\n\x08trainium\x18\n \x01(\x08\x12\x0b\n\x03pex\x18\x0b \x01(\x08\x12\r\n\x05\x63olab\x18\x0c \x01(\x08\x12\x0f\n\x07ipython\x18\r \x01(\x08\x12\x12\n\naws_lambda\x18\x0e \x01(\x08\x12\x0f\n\x07\x61md_gpu\x18\x0f \x01(\x08J\x04\x08\x05\x10\x06J\x04\x08\x06\x10\x07J\x04\x08\x07\x10\x08J\x04\x08\x08\x10\tJ\x04\x08\t\x10\n\"H\n\x06Labels\x12\x13\n\x0b\x63ode_string\x18\x01 \x01(\t\x12\x13\n\x0brepo_string\x18\x02 \x01(\t\x12\x14\n\x0c\x63ode_version\x18\x03 \x01(\t\"\xd5\x05\n\nDeprecated\x12!\n\x19keras_callback__data_type\x18\x01 \x01(\x08\x12\r\n\x05plots\x18\x05 \x01(\x08\x12!\n\x19init__config_include_keys\x18\x07 \x01(\x08\x12!\n\x19init__config_exclude_keys\x18\x08 \x01(\x08\x12\"\n\x1akeras_callback__save_model\x18\t \x01(\x08\x12\x18\n\x10langchain_tracer\x18\n \x01(\x08\x12\x1a\n\x12\x61rtifact__get_path\x18\x0b \x01(\x08\x12#\n\x1b\x61rtifactmanifestentry__name\x18\x0c \x01(\x08\x12\x1e\n\x16\x61pi__artifact_versions\x18\r \x01(\x08\x12(\n artifact_collection__change_type\x18\x0e \x01(\x08\x12\x1f\n\x17run__define_metric_copy\x18\x0f \x01(\x08\x12\x14\n\x0crun_disabled\x18\x10 \x01(\x08\x12\x16\n\x0ekeras_callback\x18\x11 \x01(\x08\x12$\n\x1crun__define_metric_best_goal\x18\x12 \x01(\x08\x12\x19\n\x11run__finish_quiet\x18\x13 \x01(\x08\x12\x18\n\x10run__reinit_bool\x18\x14 \x01(\x08\x12\x14\n\x0crun__get_url\x18\x15 \x01(\x08\x12\x19\n\x11run__project_name\x18\x16 \x01(\x08\x12\x1c\n\x14run__get_project_url\x18\x17 \x01(\x08\x12\x1a\n\x12run__get_sweep_url\x18\x18 \x01(\x08\x12 \n\x18run__use_artifact_use_as\x18\x19 \x01(\x08\x12\x18\n\x10\x61rtifact__use_as\x18\x1a \x01(\x08\x12\x1d\n\x15\x61rtifact__init_use_as\x18\x1b \x01(\x08J\x04\x08\x02\x10\x03J\x04\x08\x03\x10\x04J\x04\x08\x04\x10\x05J\x04\x08\x06\x10\x07\"|\n\x06Issues\x12%\n\x1dsettings__validation_warnings\x18\x01 \x01(\x08\x12!\n\x19settings__unexpected_args\x18\x02 \x01(\x08\x12(\n settings__preprocessing_warnings\x18\x03 \x01(\x08\x42\x1bZ\x19\x63ore/pkg/service_go_protob\x06proto3') + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'wandb.proto.wandb_telemetry_pb2', _globals) +if not _descriptor._USE_C_DESCRIPTORS: + _globals['DESCRIPTOR']._loaded_options = None + _globals['DESCRIPTOR']._serialized_options = b'Z\031core/pkg/service_go_proto' + _globals['_TELEMETRYRECORD']._serialized_start=84 + _globals['_TELEMETRYRECORD']._serialized_end=559 + _globals['_TELEMETRYRESULT']._serialized_start=561 + _globals['_TELEMETRYRESULT']._serialized_end=578 + _globals['_IMPORTS']._serialized_start=581 + _globals['_IMPORTS']._serialized_end=2399 + _globals['_FEATURE']._serialized_start=2402 + _globals['_FEATURE']._serialized_end=4087 + _globals['_ENV']._serialized_start=4090 + _globals['_ENV']._serialized_end=4291 + _globals['_LABELS']._serialized_start=4293 + _globals['_LABELS']._serialized_end=4365 + _globals['_DEPRECATED']._serialized_start=4368 + _globals['_DEPRECATED']._serialized_end=5093 + _globals['_ISSUES']._serialized_start=5095 + _globals['_ISSUES']._serialized_end=5219 +# @@protoc_insertion_point(module_scope) diff --git a/lib/python3.12/site-packages/wandb/sdk/__init__.py b/lib/python3.12/site-packages/wandb/sdk/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..94e25e9daca3c691b28af748aa8b21ea84be5707 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/__init__.py @@ -0,0 +1,37 @@ +"""W&B SDK module.""" + +__all__ = ( + "Config", + "Settings", + "Summary", + "Artifact", + "AlertLevel", + "init", + "setup", + "_attach", + "_sync", + "login", + "require", + "finish", + "teardown", + "_watch", + "_unwatch", + "sweep", + "controller", + "helper", +) + +from . import wandb_helper as helper +from .artifacts.artifact import Artifact +from .wandb_alerts import AlertLevel +from .wandb_config import Config +from .wandb_init import _attach, init +from .wandb_login import login +from .wandb_require import require +from .wandb_run import finish +from .wandb_settings import Settings +from .wandb_setup import setup, teardown +from .wandb_summary import Summary +from .wandb_sweep import controller, sweep +from .wandb_sync import _sync +from .wandb_watch import _unwatch, _watch diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/__init__.py b/lib/python3.12/site-packages/wandb/sdk/internal/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..88612aeebcbc177fa6faad1dd83cfd2ca1e88b76 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/context.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/context.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c4e14c51d46e2c114f2840641fdb5d876b8cb82a Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/context.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/datastore.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/datastore.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..04a503db28b29f724154366c107c9aa61592acb0 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/datastore.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/file_pusher.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/file_pusher.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d5015c6f4fd4d9056a923e9475dcfc7192fea1f3 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/file_pusher.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/file_stream.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/file_stream.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..64aff6880ab3d1496d1557e44ccf792c685c91ae Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/file_stream.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/handler.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/handler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7af82f9c08e499c11b5b0f729594bb7406da79a1 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/handler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/incremental_table_util.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/incremental_table_util.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..becbe505747a862960c313ad84e779e423a7f199 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/incremental_table_util.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/job_builder.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/job_builder.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..56e5c4c5641caa203f901b91c0fb8e20798b93d2 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/job_builder.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/profiler.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/profiler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5f29edb7500456b97bed6fac85d64fa20b43f321 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/profiler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/progress.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/progress.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a609256ef03d57dc87c68f8e936c6fb5b24c400c Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/progress.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/run.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/run.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2c2876e343b53549d54bfeafb5c2445ed1387185 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/run.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/sample.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/sample.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..57f1d67beed88aac92afe021d8a49eb25bf18985 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/sample.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/sender.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/sender.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..283913125b6a51a03cd42ed184ec82f3533b9925 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/sender.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/sender_config.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/sender_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cff77958a009083296f26aa892863d4a06d245d6 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/sender_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/settings_static.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/settings_static.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..012683fd37c9b27abaf4f872505214d4770eaf09 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/settings_static.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/tb_watcher.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/tb_watcher.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ebd6bf41d69f798213c8be9f8248a232ca101c19 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/tb_watcher.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/thread_local_settings.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/thread_local_settings.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bfc74085d58bfecb239b9217d29005ed3e739f0c Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/internal/__pycache__/thread_local_settings.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/_generated/__init__.py b/lib/python3.12/site-packages/wandb/sdk/internal/_generated/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..59f141690e45f602762598fa942d3860ddba0616 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/internal/_generated/__init__.py @@ -0,0 +1,15 @@ +# Generated by ariadne-codegen + +from .operations import SERVER_FEATURES_QUERY_GQL +from .server_features_query import ( + ServerFeaturesQuery, + ServerFeaturesQueryServerInfo, + ServerFeaturesQueryServerInfoFeatures, +) + +__all__ = [ + "SERVER_FEATURES_QUERY_GQL", + "ServerFeaturesQuery", + "ServerFeaturesQueryServerInfo", + "ServerFeaturesQueryServerInfoFeatures", +] diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/_generated/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/internal/_generated/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..85f1c3920d23658ddd4d292270c346215d0dcf71 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/internal/_generated/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/_generated/__pycache__/enums.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/internal/_generated/__pycache__/enums.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..482de40c7782e9dcb155c2f17a20b81b85e8e2b4 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/internal/_generated/__pycache__/enums.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/_generated/__pycache__/input_types.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/internal/_generated/__pycache__/input_types.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..26eaae9732928c8ccb1004c62804b5d27eb5b808 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/internal/_generated/__pycache__/input_types.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/_generated/__pycache__/operations.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/internal/_generated/__pycache__/operations.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..73df3c0f4db4d18a1d0aabffeed871fde1abd473 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/internal/_generated/__pycache__/operations.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/_generated/__pycache__/server_features_query.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sdk/internal/_generated/__pycache__/server_features_query.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ace03b6008006c407f14f3d1f5b7d91d2b24524a Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sdk/internal/_generated/__pycache__/server_features_query.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/_generated/enums.py b/lib/python3.12/site-packages/wandb/sdk/internal/_generated/enums.py new file mode 100644 index 0000000000000000000000000000000000000000..cc7d61d95669bab416e88c3acc378913704e64e8 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/internal/_generated/enums.py @@ -0,0 +1,4 @@ +# Generated by ariadne-codegen +# Source: core/api/graphql/schemas/schema-latest.graphql + +from __future__ import annotations diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/_generated/input_types.py b/lib/python3.12/site-packages/wandb/sdk/internal/_generated/input_types.py new file mode 100644 index 0000000000000000000000000000000000000000..cc7d61d95669bab416e88c3acc378913704e64e8 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/internal/_generated/input_types.py @@ -0,0 +1,4 @@ +# Generated by ariadne-codegen +# Source: core/api/graphql/schemas/schema-latest.graphql + +from __future__ import annotations diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/_generated/operations.py b/lib/python3.12/site-packages/wandb/sdk/internal/_generated/operations.py new file mode 100644 index 0000000000000000000000000000000000000000..6d204236d0154d142ccc4ac22ade6befe7e6e012 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/internal/_generated/operations.py @@ -0,0 +1,15 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/utils/ + +__all__ = ["SERVER_FEATURES_QUERY_GQL"] + +SERVER_FEATURES_QUERY_GQL = """ +query ServerFeaturesQuery { + serverInfo { + features { + name + isEnabled + } + } +} +""" diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/_generated/server_features_query.py b/lib/python3.12/site-packages/wandb/sdk/internal/_generated/server_features_query.py new file mode 100644 index 0000000000000000000000000000000000000000..e9bd8543434ed3cbfdef6f2636ba472d8bfa7f00 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/internal/_generated/server_features_query.py @@ -0,0 +1,27 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/utils/ + +from __future__ import annotations + +from typing import List, Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase + + +class ServerFeaturesQuery(GQLBase): + server_info: Optional[ServerFeaturesQueryServerInfo] = Field(alias="serverInfo") + + +class ServerFeaturesQueryServerInfo(GQLBase): + features: List[Optional[ServerFeaturesQueryServerInfoFeatures]] + + +class ServerFeaturesQueryServerInfoFeatures(GQLBase): + name: str + is_enabled: bool = Field(alias="isEnabled") + + +ServerFeaturesQuery.model_rebuild() +ServerFeaturesQueryServerInfo.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/context.py b/lib/python3.12/site-packages/wandb/sdk/internal/context.py new file mode 100644 index 0000000000000000000000000000000000000000..1ad49ee15a1e15b76c1cfacdbc2158fe8cfa9ad7 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/internal/context.py @@ -0,0 +1,89 @@ +"""Context Keeper.""" + +import logging +import threading +from typing import Dict, Optional + +from wandb.proto.wandb_internal_pb2 import Record, Result + +logger = logging.getLogger(__name__) + + +class Context: + _cancel_event: threading.Event + # TODO(debug_context) add debug setting to enable this + # _debug_record: Optional[Record] + + def __init__(self) -> None: + self._cancel_event = threading.Event() + # TODO(debug_context) see above + # self._debug_record = None + + def cancel(self) -> None: + self._cancel_event.set() + + @property + def cancel_event(self) -> threading.Event: + return self._cancel_event + + +def context_id_from_record(record: Record) -> str: + context_id = record.control.mailbox_slot + return context_id + + +def context_id_from_result(result: Result) -> str: + context_id = result.control.mailbox_slot + return context_id + + +class ContextKeeper: + _active_items: Dict[str, Context] + + def __init__(self) -> None: + self._active_items = {} + + def add_from_record(self, record: Record) -> Optional[Context]: + context_id = context_id_from_record(record) + if not context_id: + return None + context_obj = self.add(context_id) + + # TODO(debug_context) see above + # context_obj._debug_record = record + + return context_obj + + def add(self, context_id: str) -> Context: + assert context_id + context_obj = Context() + self._active_items[context_id] = context_obj + return context_obj + + def get(self, context_id: str) -> Optional[Context]: + item = self._active_items.get(context_id) + return item + + def release(self, context_id: str) -> None: + if not context_id: + return + _ = self._active_items.pop(context_id, None) + + def cancel(self, context_id: str) -> bool: + item = self.get(context_id) + if item: + item.cancel() + return True + return False + + # TODO(debug_context) see above + # def _debug_print_orphans(self, print_to_stdout: bool) -> None: + # for context_id, context in self._active_items.items(): + # record = context._debug_record + # record_type = record.WhichOneof("record_type") if record else "unknown" + # message = ( + # f"Context: {context_id} {context.cancel_event.is_set()} {record_type}" + # ) + # logger.warning(message) + # if print_to_stdout: + # print(message) diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/datastore.py b/lib/python3.12/site-packages/wandb/sdk/internal/datastore.py new file mode 100644 index 0000000000000000000000000000000000000000..76376f9b1535473d698d56db85a3af316f76ec16 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/internal/datastore.py @@ -0,0 +1,293 @@ +"""leveldb log datastore. + +Format is described at: + https://github.com/google/leveldb/blob/master/doc/log_format.md + +block := record* trailer? +record := + checksum: uint32 // crc32c of type and data[] ; little-endian + length: uint16 // little-endian + type: uint8 // One of FULL, FIRST, MIDDLE, LAST + data: uint8[length] + +header := + ident: char[4] + magic: uint16 + version: uint8 +""" + +# TODO: possibly restructure code by porting the C++ or go implementation + +import logging +import os +import struct +import zlib +from typing import TYPE_CHECKING, Optional, Tuple + +import wandb + +if TYPE_CHECKING: + from typing import IO, Any + + from wandb.proto.wandb_internal_pb2 import Record + +logger = logging.getLogger(__name__) + +LEVELDBLOG_HEADER_LEN = 7 +LEVELDBLOG_BLOCK_LEN = 32768 +LEVELDBLOG_DATA_LEN = LEVELDBLOG_BLOCK_LEN - LEVELDBLOG_HEADER_LEN + +LEVELDBLOG_FULL = 1 +LEVELDBLOG_FIRST = 2 +LEVELDBLOG_MIDDLE = 3 +LEVELDBLOG_LAST = 4 + +LEVELDBLOG_HEADER_IDENT = ":W&B" +LEVELDBLOG_HEADER_MAGIC = ( + 0xBEE1 # zlib.crc32(bytes("Weights & Biases", 'iso8859-1')) & 0xffff +) +LEVELDBLOG_HEADER_VERSION = 0 + +try: + bytes("", "ascii") + + def strtobytes(x): + """Strtobytes.""" + return bytes(x, "iso8859-1") + + # def bytestostr(x): + # return str(x, 'iso8859-1') + +except Exception: + strtobytes = str + # bytestostr = str + + +class DataStore: + _index: int + _flush_offset: int + + def __init__(self) -> None: + self._opened_for_scan = False + self._fp: Optional[IO[Any]] = None + self._index = 0 + self._flush_offset = 0 + self._size_bytes = 0 + + self._crc = [0] * (LEVELDBLOG_LAST + 1) + for x in range(1, LEVELDBLOG_LAST + 1): + self._crc[x] = zlib.crc32(strtobytes(chr(x))) & 0xFFFFFFFF + + assert ( + wandb._assert_is_internal_process # type: ignore + ), "DataStore can only be used in the internal process" + + def open_for_write(self, fname: str) -> None: + self._fname = fname + logger.info("open: %s", fname) + open_flags = "xb" + self._fp = open(fname, open_flags) + self._write_header() + + def open_for_append(self, fname): + # TODO: implement + self._fname = fname + logger.info("open: %s", fname) + self._fp = open(fname, "wb") + # do something with _index + + def open_for_scan(self, fname): + self._fname = fname + logger.info("open for scan: %s", fname) + self._fp = open(fname, "r+b") + self._index = 0 + self._size_bytes = os.stat(fname).st_size + self._opened_for_scan = True + self._read_header() + + def seek(self, offset: int) -> None: + self._fp.seek(offset) # type: ignore + self._index = offset + + def get_offset(self) -> int: + offset = self._fp.tell() # type: ignore + return offset + + def in_last_block(self): + """Determine if we're in the last block to handle in-progress writes.""" + return self._index > self._size_bytes - LEVELDBLOG_DATA_LEN + + def scan_record(self): + assert self._opened_for_scan, "file not open for scanning" + # TODO(jhr): handle some assertions as file corruption issues + # assume we have enough room to read header, checked by caller? + header = self._fp.read(LEVELDBLOG_HEADER_LEN) + if len(header) == 0: + return None + assert len(header) == LEVELDBLOG_HEADER_LEN, ( + f"record header is {len(header)} bytes instead of the expected {LEVELDBLOG_HEADER_LEN}" + ) + fields = struct.unpack(" LEVELDBLOG_DATA_LEN: + self._write_record( + s[data_used : data_used + LEVELDBLOG_DATA_LEN], + LEVELDBLOG_MIDDLE, + ) + data_used += LEVELDBLOG_DATA_LEN + data_left -= LEVELDBLOG_DATA_LEN + + # write last and flush the entire block to disk + self._write_record(s[data_used:], LEVELDBLOG_LAST) + self._fp.flush() + os.fsync(self._fp.fileno()) + self._flush_offset = self._index + + return start_offset, self._index, self._flush_offset + + def ensure_flushed(self, off: int) -> None: + self._fp.flush() # type: ignore + + def write(self, obj: "Record") -> Tuple[int, int, int]: + """Write a protocol buffer. + + Args: + obj: Protocol buffer to write. + + Returns: + (start_offset, end_offset, flush_offset) if successful, + None otherwise + + """ + raw_size = obj.ByteSize() + s = obj.SerializeToString() + assert len(s) == raw_size, "invalid serialization" + ret = self._write_data(s) + return ret + + def close(self) -> None: + if self._fp is not None: + logger.info("close: %s", self._fname) + self._fp.close() diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/file_pusher.py b/lib/python3.12/site-packages/wandb/sdk/internal/file_pusher.py new file mode 100644 index 0000000000000000000000000000000000000000..72b5b2d8ca41af9d44ccb284c78d600289570f6c --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/internal/file_pusher.py @@ -0,0 +1,177 @@ +import concurrent.futures +import logging +import os +import queue +import tempfile +import threading +import time +from typing import TYPE_CHECKING, Optional, Tuple + +import wandb +import wandb.util +from wandb.filesync import stats, step_checksum, step_upload +from wandb.sdk.lib.paths import LogicalPath + +if TYPE_CHECKING: + from wandb.sdk.artifacts.artifact_manifest import ArtifactManifest + from wandb.sdk.artifacts.artifact_saver import SaveFn + from wandb.sdk.internal import file_stream, internal_api + from wandb.sdk.internal.settings_static import SettingsStatic + + +logger = logging.getLogger(__name__) + + +class FilePusher: + """Parallel file upload class. + + This manages uploading multiple files in parallel. It will restart a given file's + upload job if it receives a notification that that file has been modified. The + finish() method will block until all events have been processed and all uploads are + complete. + """ + + MAX_UPLOAD_JOBS = 64 + + def __init__( + self, + api: "internal_api.Api", + file_stream: "file_stream.FileStreamApi", + settings: Optional["SettingsStatic"] = None, + ) -> None: + self._api = api + + # Temporary directory for copies we make of some file types to + # reduce the probability that the file gets changed while we're + # uploading it. + self._tempdir = tempfile.TemporaryDirectory("wandb") + + self._stats = stats.Stats() + + self._incoming_queue: queue.Queue[step_checksum.Event] = queue.Queue() + self._event_queue: queue.Queue[step_upload.Event] = queue.Queue() + + self._step_checksum = step_checksum.StepChecksum( + self._api, + self._tempdir, + self._incoming_queue, + self._event_queue, + self._stats, + ) + self._step_checksum.start() + + self._step_upload = step_upload.StepUpload( + self._api, + self._stats, + self._event_queue, + self.MAX_UPLOAD_JOBS, + file_stream=file_stream, + settings=settings, + ) + self._step_upload.start() + + self._stats_thread_stop = threading.Event() + if os.environ.get("WANDB_DEBUG"): + # debug thread to monitor and report file pusher stats + self._stats_thread = threading.Thread( + target=self._file_pusher_stats, + daemon=True, + name="FPStatsThread", + ) + self._stats_thread.start() + + def _file_pusher_stats(self) -> None: + while not self._stats_thread_stop.is_set(): + logger.info(f"FilePusher stats: {self._stats._stats}") + time.sleep(1) + + def get_status(self) -> Tuple[bool, stats.Summary]: + running = self.is_alive() + summary = self._stats.summary() + return running, summary + + def print_status(self, prefix: bool = True) -> None: + step = 0 + spinner_states = ["-", "\\", "|", "/"] + stop = False + while True: + if not self.is_alive(): + stop = True + summary = self._stats.summary() + line = f" {summary.uploaded_bytes / 1048576.0:.2f}MB of {summary.total_bytes / 1048576.0:.2f}MB uploaded ({summary.deduped_bytes / 1048576.0:.2f}MB deduped)\r" + line = spinner_states[step % 4] + line + step += 1 + wandb.termlog(line, newline=False, prefix=prefix) + if stop: + break + time.sleep(0.25) + dedupe_fraction = ( + summary.deduped_bytes / float(summary.total_bytes) + if summary.total_bytes > 0 + else 0 + ) + if dedupe_fraction > 0.01: + wandb.termlog( + "W&B sync reduced upload amount by %.1f%% " + % (dedupe_fraction * 100), + prefix=prefix, + ) + # clear progress line. + wandb.termlog(" " * 79, prefix=prefix) + + def file_counts_by_category(self) -> stats.FileCountsByCategory: + return self._stats.file_counts_by_category() + + def file_changed(self, save_name: LogicalPath, path: str, copy: bool = True): + """Tell the file pusher that a file's changed and should be uploaded. + + Args: + save_name: string logical location of the file relative to the run + directory. + path: actual string path of the file to upload on the filesystem. + """ + # Tests in linux were failing because wandb-events.jsonl didn't exist + if not os.path.exists(path) or not os.path.isfile(path): + return + if os.path.getsize(path) == 0: + return + + event = step_checksum.RequestUpload(path, save_name, copy) + self._incoming_queue.put(event) + + def store_manifest_files( + self, + manifest: "ArtifactManifest", + artifact_id: str, + save_fn: "SaveFn", + ) -> None: + event = step_checksum.RequestStoreManifestFiles(manifest, artifact_id, save_fn) + self._incoming_queue.put(event) + + def commit_artifact( + self, + artifact_id: str, + *, + finalize: bool = True, + before_commit: step_upload.PreCommitFn, + result_future: "concurrent.futures.Future[None]", + ): + event = step_checksum.RequestCommitArtifact( + artifact_id, finalize, before_commit, result_future + ) + self._incoming_queue.put(event) + + def finish(self, callback: Optional[step_upload.OnRequestFinishFn] = None): + logger.info("shutting down file pusher") + self._incoming_queue.put(step_checksum.RequestFinish(callback)) + self._stats_thread_stop.set() + + def join(self) -> None: + # NOTE: must have called finish before join + logger.info("waiting for file pusher") + while self.is_alive(): + time.sleep(0.5) + self._tempdir.cleanup() + + def is_alive(self) -> bool: + return self._step_checksum.is_alive() or self._step_upload.is_alive() diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/file_stream.py b/lib/python3.12/site-packages/wandb/sdk/internal/file_stream.py new file mode 100644 index 0000000000000000000000000000000000000000..644659c65281bcfd155d77c537ffaaf6f874d204 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/internal/file_stream.py @@ -0,0 +1,686 @@ +import functools +import itertools +import json +import logging +import os +import queue +import random +import sys +import threading +import time +from types import TracebackType +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + List, + NamedTuple, + Optional, + Set, + Tuple, + Type, + Union, +) + +if TYPE_CHECKING: + from typing import TypedDict + + class ProcessedChunk(TypedDict): + offset: int + content: List[str] + + class ProcessedBinaryChunk(TypedDict): + offset: int + content: str + encoding: str + + +import requests + +import wandb +from wandb import util +from wandb.sdk.internal import internal_api + +from ..lib import file_stream_utils + +logger = logging.getLogger(__name__) + + +class Chunk(NamedTuple): + filename: str + data: str + + +class DefaultFilePolicy: + def __init__(self, start_chunk_id: int = 0) -> None: + self._chunk_id = start_chunk_id + self.has_debug_log = False + + def process_chunks( + self, chunks: List[Chunk] + ) -> Union[bool, "ProcessedChunk", "ProcessedBinaryChunk", List["ProcessedChunk"]]: + chunk_id = self._chunk_id + self._chunk_id += len(chunks) + return {"offset": chunk_id, "content": [c.data for c in chunks]} + + # TODO: this is very inefficient, this is meant for temporary debugging and will be removed in future releases + def _debug_log(self, data: Any): + if self.has_debug_log or not os.environ.get("WANDB_DEBUG_FILESTREAM_LOG"): + return + + loaded = json.loads(data) + if not isinstance(loaded, dict): + return + + # get key size and convert to MB + key_sizes = [(k, len(json.dumps(v))) for k, v in loaded.items()] + key_msg = [f"{k}: {v / 1048576:.5f} MB" for k, v in key_sizes] + wandb.termerror(f"Step: {loaded['_step']} | {key_msg}", repeat=False) + self.has_debug_log = True + + +class JsonlFilePolicy(DefaultFilePolicy): + def process_chunks(self, chunks: List[Chunk]) -> "ProcessedChunk": + chunk_id = self._chunk_id + # TODO: chunk_id is getting reset on each request... + self._chunk_id += len(chunks) + chunk_data = [] + for chunk in chunks: + if len(chunk.data) > util.MAX_LINE_BYTES: + msg = f"Metric data exceeds maximum size of {util.to_human_size(util.MAX_LINE_BYTES)} ({util.to_human_size(len(chunk.data))})" + wandb.termerror(msg, repeat=False) + wandb._sentry.message(msg, repeat=False) + self._debug_log(chunk.data) + else: + chunk_data.append(chunk.data) + + return { + "offset": chunk_id, + "content": chunk_data, + } + + +class SummaryFilePolicy(DefaultFilePolicy): + def process_chunks(self, chunks: List[Chunk]) -> Union[bool, "ProcessedChunk"]: + data = chunks[-1].data + if len(data) > util.MAX_LINE_BYTES: + msg = f"Summary data exceeds maximum size of {util.to_human_size(util.MAX_LINE_BYTES)}. Dropping it." + wandb.termerror(msg, repeat=False) + wandb._sentry.message(msg, repeat=False) + self._debug_log(data) + return False + return {"offset": 0, "content": [data]} + + +class StreamCRState: + r"""Stream state that tracks carriage returns. + + There are two streams: stdout and stderr. We create two instances for each stream. + An instance holds state about: + found_cr: if a carriage return has been found in this stream. + cr: most recent offset (line number) where we found \r. + We update this offset with every progress bar update. + last_normal: most recent offset without a \r in this stream. + i.e. the most recent "normal" line. + """ + + found_cr: bool + cr: Optional[int] + last_normal: Optional[int] + + def __init__(self) -> None: + self.found_cr = False + self.cr = None + self.last_normal = None + + +class CRDedupeFilePolicy(DefaultFilePolicy): + r"""File stream policy for removing carriage-return erased characters. + + This is what a terminal does. We use it for console output to reduce the amount of + data we need to send over the network (eg. for progress bars), while preserving the + output's appearance in the web app. + + CR stands for "carriage return", for the character \r. It tells the terminal to move + the cursor back to the start of the current line. Progress bars (like tqdm) use \r + repeatedly to overwrite a line with newer updates. This gives the illusion of the + progress bar filling up in real-time. + """ + + def __init__(self, start_chunk_id: int = 0) -> None: + super().__init__(start_chunk_id=start_chunk_id) + self._prev_chunk = None + + self.global_offset = 0 + # cr refers to carriage return \r + self.stderr = StreamCRState() + self.stdout = StreamCRState() + + @staticmethod + def get_consecutive_offsets(console: Dict[int, str]) -> List[List[int]]: + """Compress consecutive line numbers into an interval. + + Args: + console: Dict[int, str] which maps offsets (line numbers) to lines of text. + It represents a mini version of our console dashboard on the UI. + + Returns: + A list of intervals (we compress consecutive line numbers into an interval). + + Example: + >>> console = {2: "", 3: "", 4: "", 5: "", 10: "", 11: "", 20: ""} + >>> get_consecutive_offsets(console) + [(2, 5), (10, 11), (20, 20)] + """ + offsets = sorted(list(console.keys())) + intervals: List = [] + for i, num in enumerate(offsets): + if i == 0: + intervals.append([num, num]) + continue + largest = intervals[-1][1] + if num == largest + 1: + intervals[-1][1] = num + else: + intervals.append([num, num]) + return intervals + + @staticmethod + def split_chunk(chunk: Chunk) -> Tuple[str, str]: + r"""Split chunks. + + Args: + chunk: object with two fields: filename (str) & data (str) + `chunk.data` is a str containing the lines we want. It usually contains \n or \r or both. + `chunk.data` has two possible formats (for the two streams - stdout and stderr): + - "2020-08-25T20:38:36.895321 this is my line of text\nsecond line\n" + - "ERROR 2020-08-25T20:38:36.895321 this is my line of text\nsecond line\nthird\n". + + Here's another example with a carriage return \r. + - "ERROR 2020-08-25T20:38:36.895321 \r progress bar\n" + + Returns: + A 2-tuple of strings. + First str is prefix, either "ERROR {timestamp} " or "{timestamp} ". + Second str is the rest of the string. + + Example: + >>> chunk = Chunk( + ... filename="output.log", + ... data="ERROR 2020-08-25T20:38 this is my line of text\n", + ... ) + >>> split_chunk(chunk) + ("ERROR 2020-08-25T20:38 ", "this is my line of text\n") + """ + prefix = "" + token, rest = chunk.data.split(" ", 1) + if token == "ERROR": + prefix += token + " " + token, rest = rest.split(" ", 1) + prefix += token + " " + return prefix, rest + + def process_chunks(self, chunks: List[Chunk]) -> List["ProcessedChunk"]: + r"""Process chunks. + + Args: + chunks: List of Chunk objects. See description of chunk above in `split_chunk(...)`. + + Returns: + List[Dict]. Each dict in the list contains two keys: an `offset` which holds the line number + and `content` which maps to a list of consecutive lines starting from that offset. + `offset` here means global line number in our console on the UI. + + Example: + >>> chunks = [ + Chunk("output.log", "ERROR 2020-08-25T20:38 this is my line of text\nboom\n"), + Chunk("output.log", "2020-08-25T20:38 this is test\n"), + ] + >>> process_chunks(chunks) + [ + {"offset": 0, "content": [ + "ERROR 2020-08-25T20:38 this is my line of text\n", + "ERROR 2020-08-25T20:38 boom\n", + "2020-08-25T20:38 this is test\n" + ] + } + ] + """ + # Dict[int->str], each offset (line number) mapped to a line. + # Represents a mini-version of our console pane on the UI. + console = {} + sep = os.linesep + + for c in chunks: + prefix, logs_str = self.split_chunk(c) + logs = logs_str.split(sep) + + for line in logs: + stream = self.stderr if prefix.startswith("ERROR ") else self.stdout + if line.startswith("\r"): + # line starting with \r will always overwrite a previous offset. + offset: int = ( + stream.cr + if (stream.found_cr and stream.cr is not None) + else (stream.last_normal or 0) + ) + stream.cr = offset + stream.found_cr = True + console[offset] = prefix + line[1:] + "\n" + + # Usually logs_str = "\r progress bar\n" for progress bar updates. + # If instead logs_str = "\r progress bar\n text\n text\n", + # treat this as the end of a progress bar and reset accordingly. + if ( + logs_str.count(sep) > 1 + and logs_str.replace(sep, "").count("\r") == 1 + ): + stream.found_cr = False + + elif line: + console[self.global_offset] = prefix + line + "\n" + stream.last_normal = self.global_offset + self.global_offset += 1 + + intervals = self.get_consecutive_offsets(console) + ret = [] + for a, b in intervals: + processed_chunk: ProcessedChunk = { + "offset": self._chunk_id + a, + "content": [console[i] for i in range(a, b + 1)], + } + ret.append(processed_chunk) + return ret + + +class FileStreamApi: + """Pushes chunks of files to our streaming endpoint. + + This class is used as a singleton. It has a thread that serializes access to + the streaming endpoint and performs rate-limiting and batching. + + TODO: Differentiate between binary/text encoding. + """ + + class Finish(NamedTuple): + exitcode: int + + class Preempting(NamedTuple): + pass + + class PushSuccess(NamedTuple): + artifact_id: str + save_name: str + + MAX_ITEMS_PER_PUSH = 10000 + + def __init__( + self, + api: "internal_api.Api", + run_id: str, + start_time: float, + timeout: float = 0, + settings: Optional[dict] = None, + ) -> None: + settings = settings or dict() + # NOTE: exc_info is set in thread_except_body context and readable by calling threads + self._exc_info: Optional[ + Union[ + Tuple[Type[BaseException], BaseException, TracebackType], + Tuple[None, None, None], + ] + ] = None + self._settings = settings + self._api = api + self._run_id = run_id + self._start_time = start_time + self._client = requests.Session() + timeout = timeout or 0 + if timeout > 0: + self._client.post = functools.partial(self._client.post, timeout=timeout) # type: ignore[method-assign] + self._client.auth = api.client.transport.session.auth + self._client.headers.update(api.client.transport.headers or {}) + self._client.cookies.update(api.client.transport.cookies or {}) # type: ignore[no-untyped-call] + self._client.proxies.update(api.client.transport.session.proxies or {}) + self._file_policies: Dict[str, DefaultFilePolicy] = {} + self._dropped_chunks: int = 0 + self._queue: queue.Queue = queue.Queue() + self._thread = threading.Thread(target=self._thread_except_body) + # It seems we need to make this a daemon thread to get sync.py's atexit handler to run, which + # cleans this thread up. + self._thread.name = "FileStreamThread" + self._thread.daemon = True + self._init_endpoint() + + def _init_endpoint(self) -> None: + settings = self._api.settings() + settings.update(self._settings) + self._endpoint = "{base}/files/{entity}/{project}/{run}/file_stream".format( + base=settings["base_url"], + entity=settings["entity"], + project=settings["project"], + run=self._run_id, + ) + + def start(self) -> None: + self._init_endpoint() + self._thread.start() + + def set_default_file_policy( + self, filename: str, file_policy: "DefaultFilePolicy" + ) -> None: + """Set an upload policy for a file unless one has already been set.""" + if filename not in self._file_policies: + self._file_policies[filename] = file_policy + + def set_file_policy(self, filename: str, file_policy: "DefaultFilePolicy") -> None: + self._file_policies[filename] = file_policy + + @property + def heartbeat_seconds(self) -> Union[int, float]: + # Defaults to 30 + heartbeat_seconds: Union[int, float] = self._api.dynamic_settings[ + "heartbeat_seconds" + ] + return heartbeat_seconds + + def rate_limit_seconds(self) -> Union[int, float]: + run_time = time.time() - self._start_time + if run_time < 60: + return max(1.0, self.heartbeat_seconds / 15) + elif run_time < 300: + return max(2.5, self.heartbeat_seconds / 3) + else: + return max(5.0, self.heartbeat_seconds) + + def _read_queue(self) -> List: + # called from the push thread (_thread_body), this does an initial read + # that'll block for up to rate_limit_seconds. Then it tries to read + # as much out of the queue as it can. We do this because the http post + # to the server happens within _thread_body, and can take longer than + # our rate limit. So next time we get a chance to read the queue we want + # read all the stuff that queue'd up since last time. + # + # If we have more than MAX_ITEMS_PER_PUSH in the queue then the push thread + # will get behind and data will buffer up in the queue. + return util.read_many_from_queue( + self._queue, self.MAX_ITEMS_PER_PUSH, self.rate_limit_seconds() + ) + + def _thread_body(self) -> None: + posted_data_time = time.time() + posted_anything_time = time.time() + ready_chunks = [] + uploaded: Set[str] = set() + finished: Optional[FileStreamApi.Finish] = None + while finished is None: + items = self._read_queue() + for item in items: + if isinstance(item, self.Finish): + finished = item + elif isinstance(item, self.Preempting): + request_with_retry( + self._client.post, + self._endpoint, + json={ + "complete": False, + "preempting": True, + "dropped": self._dropped_chunks, + "uploaded": list(uploaded), + }, + ) + uploaded = set() + elif isinstance(item, self.PushSuccess): + uploaded.add(item.save_name) + else: + # item is Chunk + ready_chunks.append(item) + + cur_time = time.time() + + if ready_chunks and ( + finished or cur_time - posted_data_time > self.rate_limit_seconds() + ): + posted_data_time = cur_time + posted_anything_time = cur_time + success = self._send(ready_chunks, uploaded=uploaded) + ready_chunks = [] + if success: + uploaded = set() + + # If there aren't ready chunks or uploaded files, we still want to + # send regular heartbeats so the backend doesn't erroneously mark this + # run as crashed. + if cur_time - posted_anything_time > self.heartbeat_seconds: + posted_anything_time = cur_time + + # If we encountered an error trying to publish the + # list of uploaded files, don't reset the `uploaded` + # list. Retry publishing the list on the next attempt. + if not isinstance( + request_with_retry( + self._client.post, + self._endpoint, + json={ + "complete": False, + "failed": False, + "dropped": self._dropped_chunks, + "uploaded": list(uploaded), + }, + ), + Exception, + ): + uploaded = set() + + # post the final close message. (item is self.Finish instance now) + request_with_retry( + self._client.post, + self._endpoint, + json={ + "complete": True, + "exitcode": int(finished.exitcode), + "dropped": self._dropped_chunks, + "uploaded": list(uploaded), + }, + ) + + def _thread_except_body(self) -> None: + # TODO: Consolidate with internal_util.ExceptionThread + try: + self._thread_body() + except Exception: + exc_info = sys.exc_info() + self._exc_info = exc_info + logger.exception("generic exception in filestream thread") + wandb._sentry.exception(exc_info) + raise + + def _handle_response(self, response: Union[Exception, "requests.Response"]) -> None: + """Log dropped chunks and updates dynamic settings.""" + if isinstance(response, Exception): + wandb.termerror( + "Dropped streaming file chunk (see wandb/debug-internal.log)" + ) + logger.exception(f"dropped chunk {response}") + self._dropped_chunks += 1 + else: + parsed: Optional[dict] = None + try: + parsed = response.json() + except Exception: + pass + if isinstance(parsed, dict): + limits = parsed.get("limits") + if isinstance(limits, dict): + self._api.dynamic_settings.update(limits) + + def _send(self, chunks: List[Chunk], uploaded: Optional[Set[str]] = None) -> bool: + uploaded_list = list(uploaded or []) + # create files dict. dict of pairs where chunks are a list of + # [chunk_id, chunk_data] tuples (as lists since this will be json). + files = {} + # Groupby needs group keys to be consecutive, so sort first. + chunks.sort(key=lambda c: c.filename) + for filename, file_chunks in itertools.groupby(chunks, lambda c: c.filename): + file_chunks_list = list(file_chunks) # groupby returns iterator + # Specific file policies are set by internal/sender.py + self.set_default_file_policy(filename, DefaultFilePolicy()) + files[filename] = self._file_policies[filename].process_chunks( + file_chunks_list + ) + if not files[filename]: + del files[filename] + + for fs in file_stream_utils.split_files(files, max_bytes=util.MAX_LINE_BYTES): + self._handle_response( + request_with_retry( + self._client.post, + self._endpoint, + json={"files": fs, "dropped": self._dropped_chunks}, + retry_callback=self._api.retry_callback, + ) + ) + + if uploaded_list: + if isinstance( + request_with_retry( + self._client.post, + self._endpoint, + json={ + "complete": False, + "failed": False, + "dropped": self._dropped_chunks, + "uploaded": uploaded_list, + }, + ), + Exception, + ): + return False + return True + + def stream_file(self, path: str) -> None: + name = path.split("/")[-1] + with open(path) as f: + self._send([Chunk(name, line) for line in f]) + + def enqueue_preempting(self) -> None: + self._queue.put(self.Preempting()) + + def push(self, filename: str, data: str) -> None: + """Push a chunk of a file to the streaming endpoint. + + Args: + filename: Name of file to append to. + data: Text to append to the file. + """ + self._queue.put(Chunk(filename, data)) + + def push_success(self, artifact_id: str, save_name: str) -> None: + """Notification that a file upload has been successfully completed. + + Args: + artifact_id: ID of artifact + save_name: saved name of the uploaded file + """ + self._queue.put(self.PushSuccess(artifact_id, save_name)) + + def finish(self, exitcode: int) -> None: + """Clean up. + + Anything pushed after finish will be dropped. + + Args: + exitcode: The exitcode of the watched process. + """ + logger.info("file stream finish called") + self._queue.put(self.Finish(exitcode)) + # TODO(jhr): join on a thread which exited with an exception is a noop, clean up this path + self._thread.join() + logger.info("file stream finish is done") + if self._exc_info: + logger.error("FileStream exception", exc_info=self._exc_info) + # re-raising the original exception, will get re-caught in internal.py for the sender thread + if self._exc_info[1] is not None: + raise self._exc_info[1].with_traceback(self._exc_info[2]) + + +MAX_SLEEP_SECONDS = 60 * 5 + + +def request_with_retry( + func: Callable, + *args: Any, + **kwargs: Any, +) -> Union["requests.Response", "requests.RequestException"]: + """Perform a requests http call, retrying with exponential backoff. + + Args: + func: An http-requesting function to call, like requests.post + max_retries: Maximum retries before giving up. + By default, we retry 30 times in ~2 hours before dropping the chunk + *args: passed through to func + **kwargs: passed through to func + """ + max_retries: int = kwargs.pop("max_retries", 30) + retry_callback: Optional[Callable] = kwargs.pop("retry_callback", None) + sleep = 2 + retry_count = 0 + while True: + try: + response: requests.Response = func(*args, **kwargs) + response.raise_for_status() + return response + except ( + requests.exceptions.ConnectionError, + requests.exceptions.HTTPError, + requests.exceptions.Timeout, + ) as e: + if isinstance(e, requests.exceptions.HTTPError): + # Non-retriable HTTP errors. + # + # We retry 500s just to be cautious, and because the back end + # returns them when there are infrastructure issues. If retrying + # some request winds up being problematic, we'll change the + # back end to indicate that it shouldn't be retried. + if e.response is not None and e.response.status_code in { + 400, + 403, + 404, + 409, + }: + return e + + if retry_count == max_retries: + return e + retry_count += 1 + delay = sleep + random.random() * 0.25 * sleep + if isinstance(e, requests.exceptions.HTTPError) and ( + e.response is not None and e.response.status_code == 429 + ): + err_str = ( + f"Filestream rate limit exceeded, retrying in {delay:.1f} seconds. " + ) + if retry_callback: + retry_callback(e.response.status_code, err_str) + logger.info(err_str) + else: + logger.warning( + "requests_with_retry encountered retryable exception: %s. func: %s, args: %s, kwargs: %s", + e, + func, + args, + kwargs, + ) + time.sleep(delay) + sleep *= 2 + if sleep > MAX_SLEEP_SECONDS: + sleep = MAX_SLEEP_SECONDS + except requests.exceptions.RequestException as e: + error_message = "unknown error" + try: + error_message = response.json()["error"] # todo: clean this up + except Exception: + pass + logger.exception(f"requests_with_retry error: {error_message}") + return e diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/handler.py b/lib/python3.12/site-packages/wandb/sdk/internal/handler.py new file mode 100644 index 0000000000000000000000000000000000000000..4defa1f2bc76f83bf267f5bc6f3fcd95780e950f --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/internal/handler.py @@ -0,0 +1,854 @@ +"""Handle Manager.""" + +import json +import logging +import math +import numbers +import time +from collections import defaultdict +from queue import Queue +from threading import Event +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + Iterable, + List, + Optional, + Sequence, + Tuple, + cast, +) + +from wandb.errors.links import url_registry +from wandb.proto.wandb_internal_pb2 import ( + HistoryRecord, + InternalMessages, + MetricRecord, + Record, + Result, + RunRecord, + SampledHistoryItem, + SummaryItem, + SummaryRecord, + SummaryRecordRequest, +) + +from ..interface.interface_queue import InterfaceQueue +from ..lib import handler_util, proto_util +from . import context, sample, tb_watcher +from .settings_static import SettingsStatic + +if TYPE_CHECKING: + from wandb.proto.wandb_internal_pb2 import MetricSummary + + +SummaryDict = Dict[str, Any] + +logger = logging.getLogger(__name__) + +# Update (March 5, 2024): Since ~2020/2021, when constructing the summary +# object, we had replaced the artifact path for media types with the latest +# artifact path. The primary purpose of this was to support live updating of +# media objects in the UI (since the default artifact path was fully qualified +# and would not update). However, in March of 2024, a bug was discovered with +# this approach which causes this path to be incorrect in cases where the media +# object is logged to another artifact before being logged to the run. Setting +# this to `False` disables this copy behavior. The impact is that users will +# need to refresh to see updates. Ironically, this updating behavior is not +# currently supported in the UI, so the impact of this change is minimal. +REPLACE_SUMMARY_ART_PATH_WITH_LATEST = False + + +def _dict_nested_set(target: Dict[str, Any], key_list: Sequence[str], v: Any) -> None: + # recurse down the dictionary structure: + + for k in key_list[:-1]: + target.setdefault(k, {}) + new_target = target.get(k) + if TYPE_CHECKING: + new_target = cast(Dict[str, Any], new_target) + target = new_target + # use the last element of the key to write the leaf: + target[key_list[-1]] = v + + +class HandleManager: + _consolidated_summary: SummaryDict + _sampled_history: Dict[str, sample.UniformSampleAccumulator] + _partial_history: Dict[str, Any] + _run_proto: Optional[RunRecord] + _settings: SettingsStatic + _record_q: "Queue[Record]" + _result_q: "Queue[Result]" + _stopped: Event + _writer_q: "Queue[Record]" + _interface: InterfaceQueue + _tb_watcher: Optional[tb_watcher.TBWatcher] + _metric_defines: Dict[str, MetricRecord] + _metric_globs: Dict[str, MetricRecord] + _metric_track: Dict[Tuple[str, ...], float] + _metric_copy: Dict[Tuple[str, ...], Any] + _track_time: Optional[float] + _accumulate_time: float + _run_start_time: Optional[float] + _context_keeper: context.ContextKeeper + + def __init__( + self, + settings: SettingsStatic, + record_q: "Queue[Record]", + result_q: "Queue[Result]", + stopped: Event, + writer_q: "Queue[Record]", + interface: InterfaceQueue, + context_keeper: context.ContextKeeper, + ) -> None: + self._settings = settings + self._record_q = record_q + self._result_q = result_q + self._stopped = stopped + self._writer_q = writer_q + self._interface = interface + self._context_keeper = context_keeper + + self._tb_watcher = None + self._step = 0 + + self._track_time = None + self._accumulate_time = 0 + self._run_start_time = None + + # keep track of summary from key/val updates + self._consolidated_summary = dict() + self._sampled_history = defaultdict(sample.UniformSampleAccumulator) + self._run_proto = None + self._partial_history = dict() + self._metric_defines = defaultdict(MetricRecord) + self._metric_globs = defaultdict(MetricRecord) + self._metric_track = dict() + self._metric_copy = dict() + self._internal_messages = InternalMessages() + + self._dropped_history = False + + def __len__(self) -> int: + return self._record_q.qsize() + + def handle(self, record: Record) -> None: + self._context_keeper.add_from_record(record) + record_type = record.WhichOneof("record_type") + assert record_type + handler_str = "handle_" + record_type + handler: Callable[[Record], None] = getattr(self, handler_str, None) # type: ignore + assert handler, f"unknown handle: {handler_str}" # type: ignore + handler(record) + + def handle_request(self, record: Record) -> None: + request_type = record.request.WhichOneof("request_type") + assert request_type + handler_str = "handle_request_" + request_type + handler: Callable[[Record], None] = getattr(self, handler_str, None) # type: ignore + if request_type != "network_status": + logger.debug(f"handle_request: {request_type}") + assert handler, f"unknown handle: {handler_str}" # type: ignore + handler(record) + + def _dispatch_record(self, record: Record, always_send: bool = False) -> None: + if always_send: + record.control.always_send = True + self._writer_q.put(record) + + def _respond_result(self, result: Result) -> None: + context_id = context.context_id_from_result(result) + self._context_keeper.release(context_id) + self._result_q.put(result) + + def debounce(self) -> None: + pass + + def handle_request_cancel(self, record: Record) -> None: + self._dispatch_record(record) + + def handle_request_defer(self, record: Record) -> None: + defer = record.request.defer + state = defer.state + + logger.info(f"handle defer: {state}") + if state == defer.FLUSH_TB: + if self._tb_watcher: + # shutdown tensorboard workers so we get all metrics flushed + self._tb_watcher.finish() + self._tb_watcher = None + elif state == defer.FLUSH_PARTIAL_HISTORY: + self._flush_partial_history() + elif state == defer.FLUSH_SUM: + self._save_summary(self._consolidated_summary, flush=True) + + # defer is used to drive the sender finish state machine + self._dispatch_record(record, always_send=True) + + def handle_request_python_packages(self, record: Record) -> None: + self._dispatch_record(record) + + def handle_run(self, record: Record) -> None: + if self._settings._offline: + self._run_proto = record.run + result = proto_util._result_from_record(record) + result.run_result.run.CopyFrom(record.run) + self._respond_result(result) + self._dispatch_record(record) + + def handle_stats(self, record: Record) -> None: + self._dispatch_record(record) + + def handle_config(self, record: Record) -> None: + self._dispatch_record(record) + + def handle_output(self, record: Record) -> None: + self._dispatch_record(record) + + def handle_output_raw(self, record: Record) -> None: + self._dispatch_record(record) + + def handle_files(self, record: Record) -> None: + self._dispatch_record(record) + + def handle_request_link_artifact(self, record: Record) -> None: + self._dispatch_record(record) + + def handle_use_artifact(self, record: Record) -> None: + self._dispatch_record(record) + + def handle_artifact(self, record: Record) -> None: + self._dispatch_record(record) + + def handle_alert(self, record: Record) -> None: + self._dispatch_record(record) + + def _save_summary(self, summary_dict: SummaryDict, flush: bool = False) -> None: + summary = SummaryRecord() + for k, v in summary_dict.items(): + update = summary.update.add() + update.key = k + update.value_json = json.dumps(v) + if flush: + record = Record(summary=summary) + self._dispatch_record(record) + elif not self._settings._offline: + # Send this summary update as a request since we aren't persisting every update + summary_record = SummaryRecordRequest(summary=summary) + request_record = self._interface._make_request( + summary_record=summary_record + ) + self._dispatch_record(request_record) + + def _save_history( + self, + history: HistoryRecord, + ) -> None: + for item in history.item: + # TODO(jhr) save nested keys? + k = item.key + v = json.loads(item.value_json) + if isinstance(v, numbers.Real): + self._sampled_history[k].add(v) + + def _update_summary_metrics( + self, + s: "MetricSummary", + kl: List[str], + v: "numbers.Real", + float_v: float, + goal_max: Optional[bool], + ) -> bool: + updated = False + best_key: Optional[Tuple[str, ...]] = None + if s.none: + return False + if s.copy: + # non-key list copy already done in _update_summary + if len(kl) > 1: + _dict_nested_set(self._consolidated_summary, kl, v) + return True + if s.last: + last_key = tuple(kl + ["last"]) + old_last = self._metric_track.get(last_key) + if old_last is None or float_v != old_last: + self._metric_track[last_key] = float_v + _dict_nested_set(self._consolidated_summary, last_key, v) + updated = True + if s.best: + best_key = tuple(kl + ["best"]) + if s.max or best_key and goal_max: + max_key = tuple(kl + ["max"]) + old_max = self._metric_track.get(max_key) + if old_max is None or float_v > old_max: + self._metric_track[max_key] = float_v + if s.max: + _dict_nested_set(self._consolidated_summary, max_key, v) + updated = True + if best_key: + _dict_nested_set(self._consolidated_summary, best_key, v) + updated = True + # defaulting to minimize if goal is not specified + if s.min or best_key and not goal_max: + min_key = tuple(kl + ["min"]) + old_min = self._metric_track.get(min_key) + if old_min is None or float_v < old_min: + self._metric_track[min_key] = float_v + if s.min: + _dict_nested_set(self._consolidated_summary, min_key, v) + updated = True + if best_key: + _dict_nested_set(self._consolidated_summary, best_key, v) + updated = True + if s.mean: + tot_key = tuple(kl + ["tot"]) + num_key = tuple(kl + ["num"]) + avg_key = tuple(kl + ["mean"]) + tot = self._metric_track.get(tot_key, 0.0) + num = self._metric_track.get(num_key, 0) + tot += float_v + num += 1 + self._metric_track[tot_key] = tot + self._metric_track[num_key] = num + _dict_nested_set(self._consolidated_summary, avg_key, tot / num) + updated = True + return updated + + def _update_summary_leaf( + self, + kl: List[str], + v: Any, + d: Optional[MetricRecord] = None, + ) -> bool: + has_summary = d and d.HasField("summary") + if len(kl) == 1: + copy_key = tuple(kl) + old_copy = self._metric_copy.get(copy_key) + if old_copy is None or v != old_copy: + self._metric_copy[copy_key] = v + # Store copy metric if not specified, or copy behavior + if not has_summary or (d and d.summary.copy): + self._consolidated_summary[kl[0]] = v + return True + if not d: + return False + if not has_summary: + return False + if not isinstance(v, numbers.Real): + return False + if math.isnan(v): + return False + float_v = float(v) + goal_max = None + if d.goal: + goal_max = d.goal == d.GOAL_MAXIMIZE + if self._update_summary_metrics( + d.summary, kl=kl, v=v, float_v=float_v, goal_max=goal_max + ): + return True + return False + + def _update_summary_list( + self, + kl: List[str], + v: Any, + d: Optional[MetricRecord] = None, + ) -> bool: + metric_key = ".".join([k.replace(".", "\\.") for k in kl]) + d = self._metric_defines.get(metric_key, d) + # if the dict has _type key, it's a wandb table object + if isinstance(v, dict) and not handler_util.metric_is_wandb_dict(v): + updated = False + for nk, nv in v.items(): + if self._update_summary_list(kl=kl[:] + [nk], v=nv, d=d): + updated = True + return updated + # If the dict is a media object, update the pointer to the latest alias + elif ( + REPLACE_SUMMARY_ART_PATH_WITH_LATEST + and isinstance(v, dict) + and handler_util.metric_is_wandb_dict(v) + ): + if "_latest_artifact_path" in v and "artifact_path" in v: + # TODO: Make non-destructive? + v["artifact_path"] = v["_latest_artifact_path"] + updated = self._update_summary_leaf(kl=kl, v=v, d=d) + return updated + + def _update_summary_media_objects(self, v: Dict[str, Any]) -> Dict[str, Any]: + # For now, non-recursive - just top level + for nk, nv in v.items(): + if REPLACE_SUMMARY_ART_PATH_WITH_LATEST and ( + isinstance(nv, dict) + and handler_util.metric_is_wandb_dict(nv) + and "_latest_artifact_path" in nv + and "artifact_path" in nv + ): + # TODO: Make non-destructive? + nv["artifact_path"] = nv["_latest_artifact_path"] + v[nk] = nv + return v + + def _update_summary(self, history_dict: Dict[str, Any]) -> List[str]: + # keep old behavior fast path if no define metrics have been used + if not self._metric_defines: + history_dict = self._update_summary_media_objects(history_dict) + self._consolidated_summary.update(history_dict) + return list(history_dict.keys()) + updated_keys = [] + for k, v in history_dict.items(): + if self._update_summary_list(kl=[k], v=v): + updated_keys.append(k) + return updated_keys + + def _history_assign_step( + self, + history: HistoryRecord, + history_dict: Dict[str, Any], + ) -> None: + has_step = history.HasField("step") + item = history.item.add() + item.key = "_step" + if has_step: + step = history.step.num + history_dict["_step"] = step + item.value_json = json.dumps(step) + self._step = step + 1 + else: + history_dict["_step"] = self._step + item.value_json = json.dumps(self._step) + self._step += 1 + + def _history_define_metric(self, hkey: str) -> Optional[MetricRecord]: + """Check for hkey match in glob metrics and return the defined metric.""" + # Dont define metric for internal metrics + if hkey.startswith("_"): + return None + for k, mglob in self._metric_globs.items(): + if k.endswith("*"): + if hkey.startswith(k[:-1]): + m = MetricRecord() + m.CopyFrom(mglob) + m.ClearField("glob_name") + m.options.defined = False + m.name = hkey + return m + return None + + def _history_update_leaf( + self, + kl: List[str], + v: Any, + history_dict: Dict[str, Any], + update_history: Dict[str, Any], + ) -> None: + hkey = ".".join([k.replace(".", "\\.") for k in kl]) + m = self._metric_defines.get(hkey) + if not m: + m = self._history_define_metric(hkey) + if not m: + return + mr = Record() + mr.metric.CopyFrom(m) + mr.control.local = True # Dont store this, just send it + self._handle_defined_metric(mr) + + if m.options.step_sync and m.step_metric: + if m.step_metric not in history_dict: + copy_key = tuple([m.step_metric]) + step = self._metric_copy.get(copy_key) + if step is not None: + update_history[m.step_metric] = step + + def _history_update_list( + self, + kl: List[str], + v: Any, + history_dict: Dict[str, Any], + update_history: Dict[str, Any], + ) -> None: + if isinstance(v, dict): + for nk, nv in v.items(): + self._history_update_list( + kl=kl[:] + [nk], + v=nv, + history_dict=history_dict, + update_history=update_history, + ) + return + self._history_update_leaf( + kl=kl, v=v, history_dict=history_dict, update_history=update_history + ) + + def _history_update( + self, + history: HistoryRecord, + history_dict: Dict[str, Any], + ) -> None: + # if syncing an old run, we can skip this logic + if history_dict.get("_step") is None: + self._history_assign_step(history, history_dict) + + update_history: Dict[str, Any] = {} + # Look for metric matches + if self._metric_defines or self._metric_globs: + for hkey, hval in history_dict.items(): + self._history_update_list([hkey], hval, history_dict, update_history) + + if update_history: + history_dict.update(update_history) + for k, v in update_history.items(): + item = history.item.add() + item.key = k + item.value_json = json.dumps(v) + + def handle_history(self, record: Record) -> None: + history_dict = proto_util.dict_from_proto_list(record.history.item) + + # Inject _runtime if it is not present + if history_dict is not None: + if "_runtime" not in history_dict: + self._history_assign_runtime(record.history, history_dict) + + self._history_update(record.history, history_dict) + self._dispatch_record(record) + self._save_history(record.history) + # update summary from history + updated_keys = self._update_summary(history_dict) + if updated_keys: + updated_items = {k: self._consolidated_summary[k] for k in updated_keys} + self._save_summary(updated_items) + + def _flush_partial_history( + self, + step: Optional[int] = None, + ) -> None: + if not self._partial_history: + return + + history = HistoryRecord() + for k, v in self._partial_history.items(): + item = history.item.add() + item.key = k + item.value_json = json.dumps(v) + if step is not None: + history.step.num = step + self.handle_history(Record(history=history)) + self._partial_history = {} + + def handle_request_sender_mark_report(self, record: Record) -> None: + self._dispatch_record(record, always_send=True) + + def handle_request_status_report(self, record: Record) -> None: + self._dispatch_record(record, always_send=True) + + def handle_request_partial_history(self, record: Record) -> None: + partial_history = record.request.partial_history + + flush = None + if partial_history.HasField("action"): + flush = partial_history.action.flush + + step = None + if partial_history.HasField("step"): + step = partial_history.step.num + + history_dict = proto_util.dict_from_proto_list(partial_history.item) + if step is not None: + if step < self._step: + if not self._dropped_history: + message = ( + "Step only supports monotonically increasing values, use define_metric to set a custom x " + f"axis. For details see: {url_registry.url('define-metric')}" + ) + self._internal_messages.warning.append(message) + self._dropped_history = True + message = ( + f"(User provided step: {step} is less than current step: {self._step}. " + f"Dropping entry: {history_dict})." + ) + self._internal_messages.warning.append(message) + return + elif step > self._step: + self._flush_partial_history() + self._step = step + elif flush is None: + flush = True + + self._partial_history.update(history_dict) + + if flush: + self._flush_partial_history(self._step) + + def handle_summary(self, record: Record) -> None: + summary = record.summary + for item in summary.update: + if len(item.nested_key) > 0: + # we use either key or nested_key -- not both + assert item.key == "" + key = tuple(item.nested_key) + else: + # no counter-assertion here, because technically + # summary[""] is valid + key = (item.key,) + + target = self._consolidated_summary + + # recurse down the dictionary structure: + for prop in key[:-1]: + target = target[prop] + + # use the last element of the key to write the leaf: + target[key[-1]] = json.loads(item.value_json) + + for item in summary.remove: + if len(item.nested_key) > 0: + # we use either key or nested_key -- not both + assert item.key == "" + key = tuple(item.nested_key) + else: + # no counter-assertion here, because technically + # summary[""] is valid + key = (item.key,) + + target = self._consolidated_summary + + # recurse down the dictionary structure: + for prop in key[:-1]: + target = target[prop] + + # use the last element of the key to erase the leaf: + del target[key[-1]] + + self._save_summary(self._consolidated_summary) + + def handle_exit(self, record: Record) -> None: + if self._track_time is not None: + self._accumulate_time += time.time() - self._track_time + record.exit.runtime = int(self._accumulate_time) + self._dispatch_record(record, always_send=True) + + def handle_final(self, record: Record) -> None: + self._dispatch_record(record, always_send=True) + + def handle_preempting(self, record: Record) -> None: + self._dispatch_record(record) + + def handle_header(self, record: Record) -> None: + self._dispatch_record(record) + + def handle_footer(self, record: Record) -> None: + self._dispatch_record(record) + + def handle_metadata(self, record: Record) -> None: + self._dispatch_record(record) + + def handle_request_attach(self, record: Record) -> None: + result = proto_util._result_from_record(record) + attach_id = record.request.attach.attach_id + assert attach_id + assert self._run_proto + result.response.attach_response.run.CopyFrom(self._run_proto) + self._respond_result(result) + + def handle_request_log_artifact(self, record: Record) -> None: + self._dispatch_record(record) + + def handle_telemetry(self, record: Record) -> None: + self._dispatch_record(record) + + def handle_request_run_start(self, record: Record) -> None: + run_start = record.request.run_start + assert run_start + assert run_start.run + + self._run_proto = run_start.run + + self._run_start_time = run_start.run.start_time.ToMicroseconds() / 1e6 + + self._track_time = time.time() + if run_start.run.resumed and run_start.run.runtime: + self._accumulate_time = run_start.run.runtime + else: + self._accumulate_time = 0 + + self._tb_watcher = tb_watcher.TBWatcher( + self._settings, interface=self._interface, run_proto=run_start.run + ) + + if run_start.run.resumed or run_start.run.forked: + self._step = run_start.run.starting_step + result = proto_util._result_from_record(record) + self._respond_result(result) + + def handle_request_resume(self, record: Record) -> None: + if self._track_time is not None: + self._accumulate_time += time.time() - self._track_time + self._track_time = time.time() + + def handle_request_pause(self, record: Record) -> None: + if self._track_time is not None: + self._accumulate_time += time.time() - self._track_time + self._track_time = None + + def handle_request_poll_exit(self, record: Record) -> None: + self._dispatch_record(record, always_send=True) + + def handle_request_stop_status(self, record: Record) -> None: + self._dispatch_record(record) + + def handle_request_network_status(self, record: Record) -> None: + self._dispatch_record(record) + + def handle_request_internal_messages(self, record: Record) -> None: + result = proto_util._result_from_record(record) + result.response.internal_messages_response.messages.CopyFrom( + self._internal_messages + ) + self._internal_messages.Clear() + self._respond_result(result) + + def handle_request_status(self, record: Record) -> None: + result = proto_util._result_from_record(record) + self._respond_result(result) + + def handle_request_get_summary(self, record: Record) -> None: + result = proto_util._result_from_record(record) + for key, value in self._consolidated_summary.items(): + item = SummaryItem() + item.key = key + item.value_json = json.dumps(value) + result.response.get_summary_response.item.append(item) + self._respond_result(result) + + def handle_tbrecord(self, record: Record) -> None: + logger.info("handling tbrecord: %s", record) + if self._tb_watcher: + tbrecord = record.tbrecord + self._tb_watcher.add(tbrecord.log_dir, tbrecord.save, tbrecord.root_dir) + self._dispatch_record(record) + + def _handle_defined_metric(self, record: Record) -> None: + metric = record.metric + if metric._control.overwrite: + self._metric_defines[metric.name].CopyFrom(metric) + else: + self._metric_defines[metric.name].MergeFrom(metric) + + # before dispatching, make sure step_metric is defined, if not define it and + # dispatch it locally first + metric = self._metric_defines[metric.name] + if metric.step_metric and metric.step_metric not in self._metric_defines: + m = MetricRecord(name=metric.step_metric) + self._metric_defines[metric.step_metric] = m + mr = Record() + mr.metric.CopyFrom(m) + mr.control.local = True # Don't store this, just send it + self._dispatch_record(mr) + + self._dispatch_record(record) + + def _handle_glob_metric(self, record: Record) -> None: + metric = record.metric + if metric._control.overwrite: + self._metric_globs[metric.glob_name].CopyFrom(metric) + else: + self._metric_globs[metric.glob_name].MergeFrom(metric) + self._dispatch_record(record) + + def handle_metric(self, record: Record) -> None: + """Handle MetricRecord. + + Walkthrough of the life of a MetricRecord: + + Metric defined: + - run.define_metric() parses arguments create wandb_metric.Metric + - build MetricRecord publish to interface + - handler (this function) keeps list of metrics published: + - self._metric_defines: Fully defined metrics + - self._metric_globs: metrics that have a wildcard + - dispatch writer and sender thread + - writer: records are saved to persistent store + - sender: fully defined metrics get mapped into metadata for UI + + History logged: + - handle_history + - check if metric matches _metric_defines + - if not, check if metric matches _metric_globs + - if _metric globs match, generate defined metric and call _handle_metric + + Args: + record (Record): Metric record to process + """ + if record.metric.name: + self._handle_defined_metric(record) + elif record.metric.glob_name: + self._handle_glob_metric(record) + + def handle_request_sampled_history(self, record: Record) -> None: + result = proto_util._result_from_record(record) + for key, sampled in self._sampled_history.items(): + item = SampledHistoryItem() + item.key = key + values: Iterable[Any] = sampled.get() + if all(isinstance(i, numbers.Integral) for i in values): + try: + item.values_int.extend(values) + except ValueError: + # it is safe to ignore these as this is for display information + pass + elif all(isinstance(i, numbers.Real) for i in values): + item.values_float.extend(values) + result.response.sampled_history_response.item.append(item) + self._respond_result(result) + + def handle_request_keepalive(self, record: Record) -> None: + """Handle a keepalive request. + + Keepalive is a noop, we just want to verify transport is alive. + """ + + def handle_request_run_status(self, record: Record) -> None: + self._dispatch_record(record, always_send=True) + + def handle_request_shutdown(self, record: Record) -> None: + # TODO(jhr): should we drain things and stop new requests from coming in? + result = proto_util._result_from_record(record) + self._respond_result(result) + self._stopped.set() + + def handle_request_operations(self, record: Record) -> None: + """No-op. Not implemented for the legacy-service.""" + self._respond_result(proto_util._result_from_record(record)) + + def finish(self) -> None: + logger.info("shutting down handler") + if self._tb_watcher: + self._tb_watcher.finish() + # self._context_keeper._debug_print_orphans() + + def __next__(self) -> Record: + return self._record_q.get(block=True) + + next = __next__ + + def _history_assign_runtime( + self, + history: HistoryRecord, + history_dict: Dict[str, Any], + ) -> None: + # _runtime calculation is meaningless if there is no _timestamp + if "_timestamp" not in history_dict: + return + # if it is offline sync, self._run_start_time is None + # in that case set it to the first tfevent timestamp + if self._run_start_time is None: + self._run_start_time = history_dict["_timestamp"] + history_dict["_runtime"] = history_dict["_timestamp"] - self._run_start_time + item = history.item.add() + item.key = "_runtime" + item.value_json = json.dumps(history_dict[item.key]) diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/incremental_table_util.py b/lib/python3.12/site-packages/wandb/sdk/internal/incremental_table_util.py new file mode 100644 index 0000000000000000000000000000000000000000..b1d53f7eac0cdc19328f05937427bbd173e0bdfd --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/internal/incremental_table_util.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +import time +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from wandb import Table + from wandb.sdk.artifacts.artifact import Artifact + + from ..wandb_run import Run as LocalRun + +ART_TYPE = "wandb-run-incremental-table" + + +def _get_artifact_name(run: LocalRun, key: str) -> str: + from wandb.sdk.artifacts._internal_artifact import sanitize_artifact_name + + return sanitize_artifact_name(f"run-{run.id}-incr-{key}") + + +def init_artifact(run: LocalRun, sanitized_key: str) -> Artifact: + """Initialize a new artifact for an incremental table. + + Args: + run: The wandb run associated with this artifact + sanitized_key: Sanitized string key to identify the table + + Returns: + A wandb Artifact configured for incremental table storage + """ + from wandb.sdk.artifacts._internal_artifact import InternalArtifact + + artifact = InternalArtifact( + _get_artifact_name(run, sanitized_key), + ART_TYPE, + incremental=True, + ) + return artifact + + +def get_entry_name(incr_table: Table, key: str) -> str: + """Generate a unique entry name for a table increment. + + Args: + run: The wandb run associated with this table + incr_table: The incremental table being updated + key: String key for the table entry + + Returns: + A unique string name for the table entry + """ + epoch = time.time_ns() // 1_000_000 + return f"{incr_table._increment_num}-{epoch}.{key}" diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/internal_api.py b/lib/python3.12/site-packages/wandb/sdk/internal/internal_api.py new file mode 100644 index 0000000000000000000000000000000000000000..29c09d59df39fcc952e684bc43aef26d99104ac1 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/internal/internal_api.py @@ -0,0 +1,4717 @@ +import base64 +import datetime +import functools +import http.client +import json +import logging +import os +import re +import socket +import sys +import threading +from copy import deepcopy +from pathlib import Path +from typing import ( + IO, + TYPE_CHECKING, + Any, + Callable, + Dict, + Iterable, + List, + Literal, + Mapping, + MutableMapping, + NamedTuple, + Optional, + Sequence, + TextIO, + Tuple, + Union, +) + +import click +import requests +import yaml +from wandb_gql import Client, gql +from wandb_gql.client import RetryError +from wandb_graphql.language.ast import Document + +import wandb +from wandb import env, util +from wandb.apis.normalize import normalize_exceptions, parse_backend_error_messages +from wandb.errors import AuthenticationError, CommError, UnsupportedError, UsageError +from wandb.integration.sagemaker import parse_sm_secrets +from wandb.old.settings import Settings +from wandb.proto.wandb_internal_pb2 import ServerFeature +from wandb.sdk.artifacts._validators import is_artifact_registry_project +from wandb.sdk.internal._generated import SERVER_FEATURES_QUERY_GQL, ServerFeaturesQuery +from wandb.sdk.internal.thread_local_settings import _thread_local_api_settings +from wandb.sdk.lib.gql_request import GraphQLSession +from wandb.sdk.lib.hashutil import B64MD5, md5_file_b64 + +from ..lib import credentials, retry +from ..lib.filenames import DIFF_FNAME, METADATA_FNAME +from ..lib.gitlib import GitRepo +from . import context +from .progress import Progress + +logger = logging.getLogger(__name__) + +LAUNCH_DEFAULT_PROJECT = "model-registry" + +if TYPE_CHECKING: + from typing import Literal, TypedDict + + from .progress import ProgressFn + + class CreateArtifactFileSpecInput(TypedDict, total=False): + """Corresponds to `type CreateArtifactFileSpecInput` in schema.graphql.""" + + artifactID: str + name: str + md5: str + mimetype: Optional[str] + artifactManifestID: Optional[str] + uploadPartsInput: Optional[List[Dict[str, object]]] + + class CreateArtifactFilesResponseFile(TypedDict): + id: str + name: str + displayName: str + uploadUrl: Optional[str] + uploadHeaders: Sequence[str] + uploadMultipartUrls: "UploadPartsResponse" + storagePath: str + artifact: "CreateArtifactFilesResponseFileNode" + + class CreateArtifactFilesResponseFileNode(TypedDict): + id: str + + class UploadPartsResponse(TypedDict): + uploadUrlParts: List["UploadUrlParts"] + uploadID: str + + class UploadUrlParts(TypedDict): + partNumber: int + uploadUrl: str + + class CompleteMultipartUploadArtifactInput(TypedDict): + """Corresponds to `type CompleteMultipartUploadArtifactInput` in schema.graphql.""" + + completeMultipartAction: str + completedParts: Dict[int, str] + artifactID: str + storagePath: str + uploadID: str + md5: str + + class CompleteMultipartUploadArtifactResponse(TypedDict): + digest: str + + class DefaultSettings(TypedDict): + section: str + git_remote: str + ignore_globs: Optional[List[str]] + base_url: Optional[str] + root_dir: Optional[str] + api_key: Optional[str] + entity: Optional[str] + organization: Optional[str] + project: Optional[str] + _extra_http_headers: Optional[Mapping[str, str]] + _proxies: Optional[Mapping[str, str]] + + _Response = MutableMapping + SweepState = Literal["RUNNING", "PAUSED", "CANCELED", "FINISHED"] + Number = Union[int, float] + +# class _MappingSupportsCopy(Protocol): +# def copy(self) -> "_MappingSupportsCopy": ... +# def keys(self) -> Iterable: ... +# def __getitem__(self, name: str) -> Any: ... + +httpclient_logger = logging.getLogger("http.client") +if os.environ.get("WANDB_DEBUG"): + httpclient_logger.setLevel(logging.DEBUG) + + +def check_httpclient_logger_handler() -> None: + # Only enable http.client logging if WANDB_DEBUG is set + if not os.environ.get("WANDB_DEBUG"): + return + if httpclient_logger.handlers: + return + + # Enable HTTPConnection debug logging to the logging framework + level = logging.DEBUG + + def httpclient_log(*args: Any) -> None: + httpclient_logger.log(level, " ".join(args)) + + # mask the print() built-in in the http.client module to use logging instead + http.client.print = httpclient_log # type: ignore[attr-defined] + # enable debugging + http.client.HTTPConnection.debuglevel = 1 + + root_logger = logging.getLogger("wandb") + if root_logger.handlers: + httpclient_logger.addHandler(root_logger.handlers[0]) + + +class _ThreadLocalData(threading.local): + context: Optional[context.Context] + + def __init__(self) -> None: + self.context = None + + +class _OrgNames(NamedTuple): + entity_name: str + display_name: str + + +def _match_org_with_fetched_org_entities( + organization: str, orgs: Sequence[_OrgNames] +) -> str: + """Match the organization provided in the path with the org entity or org name of the input entity. + + Args: + organization: The organization name to match + orgs: List of tuples containing (org_entity_name, org_display_name) + + Returns: + str: The matched org entity name + + Raises: + ValueError: If no matching organization is found or if multiple orgs exist without a match + """ + for org_names in orgs: + if organization in org_names: + return org_names.entity_name + + if len(orgs) == 1: + raise ValueError( + f"Expecting the organization name or entity name to match {orgs[0].display_name!r} " + f"and cannot be linked/fetched with {organization!r}. " + "Please update the target path with the correct organization name." + ) + + raise ValueError( + "Personal entity belongs to multiple organizations " + f"and cannot be linked/fetched with {organization!r}. " + "Please update the target path with the correct organization name " + "or use a team entity in the entity settings." + ) + + +class Api: + """W&B Internal Api wrapper. + + Note: + Settings are automatically overridden by looking for + a `wandb/settings` file in the current working directory or its parent + directory. If none can be found, we look in the current user's home + directory. + + Args: + default_settings(dict, optional): If you aren't using a settings + file, or you wish to override the section to use in the settings file + Override the settings here. + """ + + HTTP_TIMEOUT = env.get_http_timeout(20) + FILE_PUSHER_TIMEOUT = env.get_file_pusher_timeout() + _global_context: context.Context + _local_data: _ThreadLocalData + + def __init__( + self, + default_settings: Optional[ + Union[ + "wandb.sdk.wandb_settings.Settings", + "wandb.sdk.internal.settings_static.SettingsStatic", + Settings, + dict, + ] + ] = None, + load_settings: bool = True, + retry_timedelta: datetime.timedelta = datetime.timedelta( # okay because it's immutable + days=7 + ), + environ: MutableMapping = os.environ, + retry_callback: Optional[Callable[[int, str], Any]] = None, + api_key: Optional[str] = None, + ) -> None: + self._environ = environ + self._global_context = context.Context() + self._local_data = _ThreadLocalData() + self.default_settings: DefaultSettings = { + "section": "default", + "git_remote": "origin", + "ignore_globs": [], + "base_url": "https://api.wandb.ai", + "root_dir": None, + "api_key": None, + "entity": None, + "organization": None, + "project": None, + "_extra_http_headers": None, + "_proxies": None, + } + self.retry_timedelta = retry_timedelta + # todo: Old Settings do not follow the SupportsKeysAndGetItem Protocol + default_settings = default_settings or {} + self.default_settings.update(default_settings) # type: ignore + self.retry_uploads = 10 + self._settings = Settings( + load_settings=load_settings, + root_dir=self.default_settings.get("root_dir"), + ) + self.git = GitRepo(remote=self.settings("git_remote")) + # Mutable settings set by the _file_stream_api + self.dynamic_settings = { + "system_sample_seconds": 2, + "system_samples": 15, + "heartbeat_seconds": 30, + } + + # todo: remove these hacky hacks after settings refactor is complete + # keeping this code here to limit scope and so that it is easy to remove later + self._extra_http_headers = self.settings("_extra_http_headers") or json.loads( + self._environ.get("WANDB__EXTRA_HTTP_HEADERS", "{}") + ) + self._extra_http_headers.update(_thread_local_api_settings.headers or {}) + + auth = None + if api_key: + auth = ("api", api_key) + elif self.access_token is not None: + self._extra_http_headers["Authorization"] = f"Bearer {self.access_token}" + elif _thread_local_api_settings.cookies is None: + auth = ("api", self.api_key or "") + + proxies = self.settings("_proxies") or json.loads( + self._environ.get("WANDB__PROXIES", "{}") + ) + + self.client = Client( + transport=GraphQLSession( + headers={ + "User-Agent": self.user_agent, + "X-WANDB-USERNAME": env.get_username(env=self._environ), + "X-WANDB-USER-EMAIL": env.get_user_email(env=self._environ), + **self._extra_http_headers, + }, + use_json=True, + # this timeout won't apply when the DNS lookup fails. in that case, it will be 60s + # https://bugs.python.org/issue22889 + timeout=self.HTTP_TIMEOUT, + auth=auth, + url=f"{self.settings('base_url')}/graphql", + cookies=_thread_local_api_settings.cookies, + proxies=proxies, + ) + ) + + self.retry_callback = retry_callback + self._retry_gql = retry.Retry( + self.execute, + retry_timedelta=retry_timedelta, + check_retry_fn=util.no_retry_auth, + retryable_exceptions=(RetryError, requests.RequestException), + retry_callback=retry_callback, + ) + self._current_run_id: Optional[str] = None + self._file_stream_api = None + self._upload_file_session = requests.Session() + if self.FILE_PUSHER_TIMEOUT: + self._upload_file_session.put = functools.partial( # type: ignore + self._upload_file_session.put, + timeout=self.FILE_PUSHER_TIMEOUT, + ) + if proxies: + self._upload_file_session.proxies.update(proxies) + # This Retry class is initialized once for each Api instance, so this + # defaults to retrying 1 million times per process or 7 days + self.upload_file_retry = normalize_exceptions( + retry.retriable(retry_timedelta=retry_timedelta)(self.upload_file) + ) + self.upload_multipart_file_chunk_retry = normalize_exceptions( + retry.retriable(retry_timedelta=retry_timedelta)( + self.upload_multipart_file_chunk + ) + ) + self._client_id_mapping: Dict[str, str] = {} + # Large file uploads to azure can optionally use their SDK + self._azure_blob_module = util.get_module("azure.storage.blob") + + self.query_types: Optional[List[str]] = None + self.mutation_types: Optional[List[str]] = None + self.server_info_types: Optional[List[str]] = None + self.server_use_artifact_input_info: Optional[List[str]] = None + self.server_create_artifact_input_info: Optional[List[str]] = None + self.server_artifact_fields_info: Optional[List[str]] = None + self.server_organization_type_fields_info: Optional[List[str]] = None + self.server_supports_enabling_artifact_usage_tracking: Optional[bool] = None + self._max_cli_version: Optional[str] = None + self._server_settings_type: Optional[List[str]] = None + self.fail_run_queue_item_input_info: Optional[List[str]] = None + self.create_launch_agent_input_info: Optional[List[str]] = None + self.server_create_run_queue_supports_drc: Optional[bool] = None + self.server_create_run_queue_supports_priority: Optional[bool] = None + self.server_supports_template_variables: Optional[bool] = None + self.server_push_to_run_queue_supports_priority: Optional[bool] = None + + self._server_features_cache: Optional[Dict[str, bool]] = None + + def gql(self, *args: Any, **kwargs: Any) -> Any: + ret = self._retry_gql( + *args, + retry_cancel_event=self.context.cancel_event, + **kwargs, + ) + return ret + + def set_local_context(self, api_context: Optional[context.Context]) -> None: + self._local_data.context = api_context + + def clear_local_context(self) -> None: + self._local_data.context = None + + @property + def context(self) -> context.Context: + return self._local_data.context or self._global_context + + def reauth(self) -> None: + """Ensure the current api key is set in the transport.""" + self.client.transport.session.auth = ("api", self.api_key or "") + + def relocate(self) -> None: + """Ensure the current api points to the right server.""" + self.client.transport.url = "{}/graphql".format(self.settings("base_url")) + + def execute(self, *args: Any, **kwargs: Any) -> "_Response": + """Wrapper around execute that logs in cases of failure.""" + try: + return self.client.execute(*args, **kwargs) # type: ignore + except requests.exceptions.HTTPError as err: + response = err.response + assert response is not None + logger.exception("Error executing GraphQL.") + for error in parse_backend_error_messages(response): + wandb.termerror(f"Error while calling W&B API: {error} ({response})") + raise + + def validate_api_key(self) -> bool: + """Returns whether the API key stored on initialization is valid.""" + res = self.execute(gql("query { viewer { id } }")) + return res is not None and res["viewer"] is not None + + def set_current_run_id(self, run_id: str) -> None: + self._current_run_id = run_id + + @property + def current_run_id(self) -> Optional[str]: + return self._current_run_id + + @property + def user_agent(self) -> str: + return f"W&B Internal Client {wandb.__version__}" + + @property + def api_key(self) -> Optional[str]: + if _thread_local_api_settings.api_key: + return _thread_local_api_settings.api_key + auth = requests.utils.get_netrc_auth(self.api_url) + key = None + if auth: + key = auth[-1] + + # Environment should take precedence + env_key: Optional[str] = self._environ.get(env.API_KEY) + sagemaker_key: Optional[str] = parse_sm_secrets().get(env.API_KEY) + default_key: Optional[str] = self.default_settings.get("api_key") + return env_key or key or sagemaker_key or default_key + + @property + def access_token(self) -> Optional[str]: + """Retrieves an access token for authentication. + + This function attempts to exchange an identity token for a temporary + access token from the server, and save it to the credentials file. + It uses the path to the identity token as defined in the environment + variables. If the environment variable is not set, it returns None. + + Returns: + Optional[str]: The access token if available, otherwise None if + no identity token is supplied. + Raises: + AuthenticationError: If the path to the identity token is not found. + """ + token_file_str = self._environ.get(env.IDENTITY_TOKEN_FILE) + if not token_file_str: + return None + + token_file = Path(token_file_str) + if not token_file.exists(): + raise AuthenticationError(f"Identity token file not found: {token_file}") + + base_url = self.settings("base_url") + credentials_file = env.get_credentials_file( + str(credentials.DEFAULT_WANDB_CREDENTIALS_FILE), self._environ + ) + return credentials.access_token(base_url, token_file, credentials_file) + + @property + def api_url(self) -> str: + return self.settings("base_url") # type: ignore + + @property + def app_url(self) -> str: + return wandb.util.app_url(self.api_url) + + @property + def default_entity(self) -> str: + return self.viewer().get("entity") # type: ignore + + def settings(self, key: Optional[str] = None, section: Optional[str] = None) -> Any: + """The settings overridden from the wandb/settings file. + + Args: + key (str, optional): If provided only this setting is returned + section (str, optional): If provided this section of the setting file is + used, defaults to "default" + + Returns: + A dict with the current settings + + { + "entity": "models", + "base_url": "https://api.wandb.ai", + "project": None, + "organization": "my-org", + } + """ + result = self.default_settings.copy() + result.update(self._settings.items(section=section)) # type: ignore + result.update( + { + "entity": env.get_entity( + self._settings.get( + Settings.DEFAULT_SECTION, + "entity", + fallback=result.get("entity"), + ), + env=self._environ, + ), + "organization": env.get_organization( + self._settings.get( + Settings.DEFAULT_SECTION, + "organization", + fallback=result.get("organization"), + ), + env=self._environ, + ), + "project": env.get_project( + self._settings.get( + Settings.DEFAULT_SECTION, + "project", + fallback=result.get("project"), + ), + env=self._environ, + ), + "base_url": env.get_base_url( + self._settings.get( + Settings.DEFAULT_SECTION, + "base_url", + fallback=result.get("base_url"), + ), + env=self._environ, + ), + "ignore_globs": env.get_ignore( + self._settings.get( + Settings.DEFAULT_SECTION, + "ignore_globs", + fallback=result.get("ignore_globs"), + ), + env=self._environ, + ), + } + ) + + return result if key is None else result[key] # type: ignore + + def clear_setting( + self, key: str, globally: bool = False, persist: bool = False + ) -> None: + self._settings.clear( + Settings.DEFAULT_SECTION, key, globally=globally, persist=persist + ) + + def set_setting( + self, key: str, value: Any, globally: bool = False, persist: bool = False + ) -> None: + self._settings.set( + Settings.DEFAULT_SECTION, key, value, globally=globally, persist=persist + ) + if key == "entity": + env.set_entity(value, env=self._environ) + elif key == "project": + env.set_project(value, env=self._environ) + elif key == "base_url": + self.relocate() + + def parse_slug( + self, slug: str, project: Optional[str] = None, run: Optional[str] = None + ) -> Tuple[str, str]: + """Parse a slug into a project and run. + + Args: + slug (str): The slug to parse + project (str, optional): The project to use, if not provided it will be + inferred from the slug + run (str, optional): The run to use, if not provided it will be inferred + from the slug + + Returns: + A dict with the project and run + """ + if slug and "/" in slug: + parts = slug.split("/") + project = parts[0] + run = parts[1] + else: + project = project or self.settings().get("project") + if project is None: + raise CommError("No default project configured.") + run = run or slug or self.current_run_id or env.get_run(env=self._environ) + assert run, "run must be specified" + return project, run + + @normalize_exceptions + def server_info_introspection(self) -> Tuple[List[str], List[str], List[str]]: + query_string = """ + query ProbeServerCapabilities { + QueryType: __type(name: "Query") { + ...fieldData + } + MutationType: __type(name: "Mutation") { + ...fieldData + } + ServerInfoType: __type(name: "ServerInfo") { + ...fieldData + } + } + + fragment fieldData on __Type { + fields { + name + } + } + """ + if ( + self.query_types is None + or self.mutation_types is None + or self.server_info_types is None + ): + query = gql(query_string) + res = self.gql(query) + + self.query_types = [ + field.get("name", "") + for field in res.get("QueryType", {}).get("fields", [{}]) + ] + self.mutation_types = [ + field.get("name", "") + for field in res.get("MutationType", {}).get("fields", [{}]) + ] + self.server_info_types = [ + field.get("name", "") + for field in res.get("ServerInfoType", {}).get("fields", [{}]) + ] + return self.query_types, self.server_info_types, self.mutation_types + + @normalize_exceptions + def server_settings_introspection(self) -> None: + query_string = """ + query ProbeServerSettings { + ServerSettingsType: __type(name: "ServerSettings") { + ...fieldData + } + } + + fragment fieldData on __Type { + fields { + name + } + } + """ + if self._server_settings_type is None: + query = gql(query_string) + res = self.gql(query) + self._server_settings_type = ( + [ + field.get("name", "") + for field in res.get("ServerSettingsType", {}).get("fields", [{}]) + ] + if res + else [] + ) + + def server_use_artifact_input_introspection(self) -> List: + query_string = """ + query ProbeServerUseArtifactInput { + UseArtifactInputInfoType: __type(name: "UseArtifactInput") { + name + inputFields { + name + } + } + } + """ + + if self.server_use_artifact_input_info is None: + query = gql(query_string) + res = self.gql(query) + self.server_use_artifact_input_info = [ + field.get("name", "") + for field in res.get("UseArtifactInputInfoType", {}).get( + "inputFields", [{}] + ) + ] + return self.server_use_artifact_input_info + + @normalize_exceptions + def launch_agent_introspection(self) -> Optional[str]: + query = gql( + """ + query LaunchAgentIntrospection { + LaunchAgentType: __type(name: "LaunchAgent") { + name + } + } + """ + ) + + res = self.gql(query) + return res.get("LaunchAgentType") or None + + @normalize_exceptions + def create_run_queue_introspection(self) -> Tuple[bool, bool, bool]: + _, _, mutations = self.server_info_introspection() + query_string = """ + query ProbeCreateRunQueueInput { + CreateRunQueueInputType: __type(name: "CreateRunQueueInput") { + name + inputFields { + name + } + } + } + """ + if ( + self.server_create_run_queue_supports_drc is None + or self.server_create_run_queue_supports_priority is None + ): + query = gql(query_string) + res = self.gql(query) + if res is None: + raise CommError("Could not get CreateRunQueue input from GQL.") + self.server_create_run_queue_supports_drc = "defaultResourceConfigID" in [ + x["name"] + for x in ( + res.get("CreateRunQueueInputType", {}).get("inputFields", [{}]) + ) + ] + self.server_create_run_queue_supports_priority = "prioritizationMode" in [ + x["name"] + for x in ( + res.get("CreateRunQueueInputType", {}).get("inputFields", [{}]) + ) + ] + return ( + "createRunQueue" in mutations, + self.server_create_run_queue_supports_drc, + self.server_create_run_queue_supports_priority, + ) + + @normalize_exceptions + def upsert_run_queue_introspection(self) -> bool: + _, _, mutations = self.server_info_introspection() + return "upsertRunQueue" in mutations + + @normalize_exceptions + def push_to_run_queue_introspection(self) -> Tuple[bool, bool]: + query_string = """ + query ProbePushToRunQueueInput { + PushToRunQueueInputType: __type(name: "PushToRunQueueInput") { + name + inputFields { + name + } + } + } + """ + + if ( + self.server_supports_template_variables is None + or self.server_push_to_run_queue_supports_priority is None + ): + query = gql(query_string) + res = self.gql(query) + self.server_supports_template_variables = "templateVariableValues" in [ + x["name"] + for x in ( + res.get("PushToRunQueueInputType", {}).get("inputFields", [{}]) + ) + ] + self.server_push_to_run_queue_supports_priority = "priority" in [ + x["name"] + for x in ( + res.get("PushToRunQueueInputType", {}).get("inputFields", [{}]) + ) + ] + + return ( + self.server_supports_template_variables, + self.server_push_to_run_queue_supports_priority, + ) + + @normalize_exceptions + def create_default_resource_config_introspection(self) -> bool: + _, _, mutations = self.server_info_introspection() + return "createDefaultResourceConfig" in mutations + + @normalize_exceptions + def fail_run_queue_item_introspection(self) -> bool: + _, _, mutations = self.server_info_introspection() + return "failRunQueueItem" in mutations + + @normalize_exceptions + def fail_run_queue_item_fields_introspection(self) -> List: + if self.fail_run_queue_item_input_info: + return self.fail_run_queue_item_input_info + query_string = """ + query ProbeServerFailRunQueueItemInput { + FailRunQueueItemInputInfoType: __type(name:"FailRunQueueItemInput") { + inputFields{ + name + } + } + } + """ + + query = gql(query_string) + res = self.gql(query) + + self.fail_run_queue_item_input_info = [ + field.get("name", "") + for field in res.get("FailRunQueueItemInputInfoType", {}).get( + "inputFields", [{}] + ) + ] + return self.fail_run_queue_item_input_info + + @normalize_exceptions + def fail_run_queue_item( + self, + run_queue_item_id: str, + message: str, + stage: str, + file_paths: Optional[List[str]] = None, + ) -> bool: + if not self.fail_run_queue_item_introspection(): + return False + variable_values: Dict[str, Union[str, Optional[List[str]]]] = { + "runQueueItemId": run_queue_item_id, + } + if "message" in self.fail_run_queue_item_fields_introspection(): + variable_values.update({"message": message, "stage": stage}) + if file_paths is not None: + variable_values["filePaths"] = file_paths + mutation_string = """ + mutation failRunQueueItem($runQueueItemId: ID!, $message: String!, $stage: String!, $filePaths: [String!]) { + failRunQueueItem( + input: { + runQueueItemId: $runQueueItemId + message: $message + stage: $stage + filePaths: $filePaths + } + ) { + success + } + } + """ + else: + mutation_string = """ + mutation failRunQueueItem($runQueueItemId: ID!) { + failRunQueueItem( + input: { + runQueueItemId: $runQueueItemId + } + ) { + success + } + } + """ + + mutation = gql(mutation_string) + response = self.gql(mutation, variable_values=variable_values) + result: bool = response["failRunQueueItem"]["success"] + return result + + @normalize_exceptions + def update_run_queue_item_warning_introspection(self) -> bool: + _, _, mutations = self.server_info_introspection() + return "updateRunQueueItemWarning" in mutations + + def _server_features(self) -> Dict[str, bool]: + # NOTE: Avoid caching via `@cached_property`, due to undocumented + # locking behavior before Python 3.12. + # See: https://github.com/python/cpython/issues/87634 + query = gql(SERVER_FEATURES_QUERY_GQL) + try: + response = self.gql(query) + except Exception as e: + # Unfortunately we currently have to match on the text of the error message, + # as the `gql` client raises `Exception` rather than a more specific error. + if 'Cannot query field "features" on type "ServerInfo".' in str(e): + self._server_features_cache = {} + else: + raise + else: + info = ServerFeaturesQuery.model_validate(response).server_info + if info and (feats := info.features): + self._server_features_cache = {f.name: f.is_enabled for f in feats if f} + else: + self._server_features_cache = {} + return self._server_features_cache + + def _server_supports(self, feature: Union[int, str]) -> bool: + """Return whether the current server supports the given feature. + + This also caches the underlying lookup of server feature flags, + and it maps {feature_name (str) -> is_enabled (bool)}. + + Good to use for features that have a fallback mechanism for older servers. + """ + # If we're given the protobuf enum value, convert to a string name. + # NOTE: We deliberately use names (str) instead of enum values (int) + # as the keys here, since: + # - the server identifies features by their name, rather than (client-side) enum value + # - the defined list of client-side flags may be behind the server-side list of flags + key = ServerFeature.Name(feature) if isinstance(feature, int) else feature + return self._server_features().get(key) or False + + @normalize_exceptions + def update_run_queue_item_warning( + self, + run_queue_item_id: str, + message: str, + stage: str, + file_paths: Optional[List[str]] = None, + ) -> bool: + if not self.update_run_queue_item_warning_introspection(): + return False + mutation = gql( + """ + mutation updateRunQueueItemWarning($runQueueItemId: ID!, $message: String!, $stage: String!, $filePaths: [String!]) { + updateRunQueueItemWarning( + input: { + runQueueItemId: $runQueueItemId + message: $message + stage: $stage + filePaths: $filePaths + } + ) { + success + } + } + """ + ) + response = self.gql( + mutation, + variable_values={ + "runQueueItemId": run_queue_item_id, + "message": message, + "stage": stage, + "filePaths": file_paths, + }, + ) + result: bool = response["updateRunQueueItemWarning"]["success"] + return result + + @normalize_exceptions + def viewer(self) -> Dict[str, Any]: + query = gql( + """ + query Viewer{ + viewer { + id + entity + username + flags + teams { + edges { + node { + name + } + } + } + } + } + """ + ) + res = self.gql(query) + return res.get("viewer") or {} + + @normalize_exceptions + def max_cli_version(self) -> Optional[str]: + if self._max_cli_version is not None: + return self._max_cli_version + + query_types, server_info_types, _ = self.server_info_introspection() + cli_version_exists = ( + "serverInfo" in query_types and "cliVersionInfo" in server_info_types + ) + if not cli_version_exists: + return None + + _, server_info = self.viewer_server_info() + self._max_cli_version = server_info.get("cliVersionInfo", {}).get( + "max_cli_version" + ) + return self._max_cli_version + + @normalize_exceptions + def viewer_server_info(self) -> Tuple[Dict[str, Any], Dict[str, Any]]: + local_query = """ + latestLocalVersionInfo { + outOfDate + latestVersionString + versionOnThisInstanceString + } + """ + cli_query = """ + serverInfo { + cliVersionInfo + _LOCAL_QUERY_ + } + """ + query_template = """ + query Viewer{ + viewer { + id + entity + username + email + flags + teams { + edges { + node { + name + } + } + } + } + _CLI_QUERY_ + } + """ + query_types, server_info_types, _ = self.server_info_introspection() + + cli_version_exists = ( + "serverInfo" in query_types and "cliVersionInfo" in server_info_types + ) + + local_version_exists = ( + "serverInfo" in query_types + and "latestLocalVersionInfo" in server_info_types + ) + + cli_query_string = "" if not cli_version_exists else cli_query + local_query_string = "" if not local_version_exists else local_query + + query_string = query_template.replace("_CLI_QUERY_", cli_query_string).replace( + "_LOCAL_QUERY_", local_query_string + ) + query = gql(query_string) + res = self.gql(query) + return res.get("viewer") or {}, res.get("serverInfo") or {} + + @normalize_exceptions + def list_projects(self, entity: Optional[str] = None) -> List[Dict[str, str]]: + """List projects in W&B scoped by entity. + + Args: + entity (str, optional): The entity to scope this project to. + + Returns: + [{"id","name","description"}] + """ + query = gql( + """ + query EntityProjects($entity: String) { + models(first: 10, entityName: $entity) { + edges { + node { + id + name + description + } + } + } + } + """ + ) + project_list: List[Dict[str, str]] = self._flatten_edges( + self.gql( + query, variable_values={"entity": entity or self.settings("entity")} + )["models"] + ) + return project_list + + @normalize_exceptions + def project(self, project: str, entity: Optional[str] = None) -> "_Response": + """Retrieve project. + + Args: + project (str): The project to get details for + entity (str, optional): The entity to scope this project to. + + Returns: + [{"id","name","repo","dockerImage","description"}] + """ + query = gql( + """ + query ProjectDetails($entity: String, $project: String) { + model(name: $project, entityName: $entity) { + id + name + repo + dockerImage + description + } + } + """ + ) + response: _Response = self.gql( + query, variable_values={"entity": entity, "project": project} + )["model"] + return response + + @normalize_exceptions + def sweep( + self, + sweep: str, + specs: str, + project: Optional[str] = None, + entity: Optional[str] = None, + ) -> Dict[str, Any]: + """Retrieve sweep. + + Args: + sweep (str): The sweep to get details for + specs (str): history specs + project (str, optional): The project to scope this sweep to. + entity (str, optional): The entity to scope this sweep to. + + Returns: + [{"id","name","repo","dockerImage","description"}] + """ + query = gql( + """ + query SweepWithRuns($entity: String, $project: String, $sweep: String!, $specs: [JSONString!]!) { + project(name: $project, entityName: $entity) { + sweep(sweepName: $sweep) { + id + name + method + state + description + config + createdAt + heartbeatAt + updatedAt + earlyStopJobRunning + bestLoss + controller + scheduler + runs { + edges { + node { + name + state + config + exitcode + heartbeatAt + shouldStop + failed + stopped + running + summaryMetrics + sampledHistory(specs: $specs) + } + } + } + } + } + } + """ + ) + entity = entity or self.settings("entity") + project = project or self.settings("project") + response = self.gql( + query, + variable_values={ + "entity": entity, + "project": project, + "sweep": sweep, + "specs": specs, + }, + ) + if response["project"] is None or response["project"]["sweep"] is None: + raise ValueError(f"Sweep {entity}/{project}/{sweep} not found") + data: Dict[str, Any] = response["project"]["sweep"] + if data: + data["runs"] = self._flatten_edges(data["runs"]) + return data + + @normalize_exceptions + def list_runs( + self, project: str, entity: Optional[str] = None + ) -> List[Dict[str, str]]: + """List runs in W&B scoped by project. + + Args: + project (str): The project to scope the runs to + entity (str, optional): The entity to scope this project to. Defaults to public models + + Returns: + [{"id","name","description"}] + """ + query = gql( + """ + query ProjectRuns($model: String!, $entity: String) { + model(name: $model, entityName: $entity) { + buckets(first: 10) { + edges { + node { + id + name + displayName + description + } + } + } + } + } + """ + ) + return self._flatten_edges( + self.gql( + query, + variable_values={ + "entity": entity or self.settings("entity"), + "model": project or self.settings("project"), + }, + )["model"]["buckets"] + ) + + @normalize_exceptions + def run_config( + self, project: str, run: Optional[str] = None, entity: Optional[str] = None + ) -> Tuple[str, Dict[str, Any], Optional[str], Dict[str, Any]]: + """Get the relevant configs for a run. + + Args: + project (str): The project to download, (can include bucket) + run (str, optional): The run to download + entity (str, optional): The entity to scope this project to. + """ + check_httpclient_logger_handler() + + query = gql( + """ + query RunConfigs( + $name: String!, + $entity: String, + $run: String!, + $pattern: String!, + $includeConfig: Boolean!, + ) { + model(name: $name, entityName: $entity) { + bucket(name: $run) { + config @include(if: $includeConfig) + commit @include(if: $includeConfig) + files(pattern: $pattern) { + pageInfo { + hasNextPage + endCursor + } + edges { + node { + name + directUrl + } + } + } + } + } + } + """ + ) + + variable_values = { + "name": project, + "run": run, + "entity": entity, + "includeConfig": True, + } + + commit: str = "" + config: Dict[str, Any] = {} + patch: Optional[str] = None + metadata: Dict[str, Any] = {} + + # If we use the `names` parameter on the `files` node, then the server + # will helpfully give us and 'open' file handle to the files that don't + # exist. This is so that we can upload data to it. However, in this + # case, we just want to download that file and not upload to it, so + # let's instead query for the files that do exist using `pattern` + # (with no wildcards). + # + # Unfortunately we're unable to construct a single pattern that matches + # our 2 files, we would need something like regex for that. + for filename in [DIFF_FNAME, METADATA_FNAME]: + variable_values["pattern"] = filename + response = self.gql(query, variable_values=variable_values) + if response["model"] is None: + raise CommError(f"Run {entity}/{project}/{run} not found") + run_obj: Dict = response["model"]["bucket"] + # we only need to fetch this config once + if variable_values["includeConfig"]: + commit = run_obj["commit"] + config = json.loads(run_obj["config"] or "{}") + variable_values["includeConfig"] = False + if run_obj["files"] is not None: + for file_edge in run_obj["files"]["edges"]: + name = file_edge["node"]["name"] + url = file_edge["node"]["directUrl"] + res = requests.get(url) + res.raise_for_status() + if name == METADATA_FNAME: + metadata = res.json() + elif name == DIFF_FNAME: + patch = res.text + + return commit, config, patch, metadata + + @normalize_exceptions + def run_resume_status( + self, entity: str, project_name: str, name: str + ) -> Optional[Dict[str, Any]]: + """Check if a run exists and get resume information. + + Args: + entity (str): The entity to scope this project to. + project_name (str): The project to download, (can include bucket) + name (str): The run to download + """ + # Pulling wandbConfig.start_time is required so that we can determine if a run has actually started + query = gql( + """ + query RunResumeStatus($project: String, $entity: String, $name: String!) { + model(name: $project, entityName: $entity) { + id + name + entity { + id + name + } + + bucket(name: $name, missingOk: true) { + id + name + summaryMetrics + displayName + logLineCount + historyLineCount + eventsLineCount + historyTail + eventsTail + config + tags + wandbConfig(keys: ["t"]) + } + } + } + """ + ) + + response = self.gql( + query, + variable_values={ + "entity": entity, + "project": project_name, + "name": name, + }, + ) + + if "model" not in response or "bucket" not in (response["model"] or {}): + return None + + project = response["model"] + self.set_setting("project", project_name) + if "entity" in project: + self.set_setting("entity", project["entity"]["name"]) + + result: Dict[str, Any] = project["bucket"] + + return result + + @normalize_exceptions + def check_stop_requested( + self, project_name: str, entity_name: str, run_id: str + ) -> bool: + query = gql( + """ + query RunStoppedStatus($projectName: String, $entityName: String, $runId: String!) { + project(name:$projectName, entityName:$entityName) { + run(name:$runId) { + stopped + } + } + } + """ + ) + + response = self.gql( + query, + variable_values={ + "projectName": project_name, + "entityName": entity_name, + "runId": run_id, + }, + ) + + project = response.get("project", None) + if not project: + return False + run = project.get("run", None) + if not run: + return False + + status: bool = run["stopped"] + return status + + def format_project(self, project: str) -> str: + return re.sub(r"\W+", "-", project.lower()).strip("-_") + + @normalize_exceptions + def upsert_project( + self, + project: str, + id: Optional[str] = None, + description: Optional[str] = None, + entity: Optional[str] = None, + ) -> Dict[str, Any]: + """Create a new project. + + Args: + project (str): The project to create + description (str, optional): A description of this project + entity (str, optional): The entity to scope this project to. + """ + mutation = gql( + """ + mutation UpsertModel($name: String!, $id: String, $entity: String!, $description: String, $repo: String) { + upsertModel(input: { id: $id, name: $name, entityName: $entity, description: $description, repo: $repo }) { + model { + name + description + } + } + } + """ + ) + response = self.gql( + mutation, + variable_values={ + "name": self.format_project(project), + "entity": entity or self.settings("entity"), + "description": description, + "id": id, + }, + ) + # TODO(jhr): Commenting out 'repo' field for cling, add back + # 'description': description, 'repo': self.git.remote_url, 'id': id}) + result: Dict[str, Any] = response["upsertModel"]["model"] + return result + + @normalize_exceptions + def entity_is_team(self, entity: str) -> bool: + query = gql( + """ + query EntityIsTeam($entity: String!) { + entity(name: $entity) { + id + isTeam + } + } + """ + ) + variable_values = { + "entity": entity, + } + + res = self.gql(query, variable_values) + if res.get("entity") is None: + raise Exception( + f"Error fetching entity {entity} " + "check that you have access to this entity" + ) + + is_team: bool = res["entity"]["isTeam"] + return is_team + + @normalize_exceptions + def get_project_run_queues(self, entity: str, project: str) -> List[Dict[str, str]]: + query = gql( + """ + query ProjectRunQueues($entity: String!, $projectName: String!){ + project(entityName: $entity, name: $projectName) { + runQueues { + id + name + createdBy + access + } + } + } + """ + ) + variable_values = { + "projectName": project, + "entity": entity, + } + + res = self.gql(query, variable_values) + if res.get("project") is None: + # circular dependency: (LAUNCH_DEFAULT_PROJECT = model-registry) + if project == "model-registry": + msg = ( + f"Error fetching run queues for {entity} " + "check that you have access to this entity and project" + ) + else: + msg = ( + f"Error fetching run queues for {entity}/{project} " + "check that you have access to this entity and project" + ) + + raise Exception(msg) + + project_run_queues: List[Dict[str, str]] = res["project"]["runQueues"] + return project_run_queues + + @normalize_exceptions + def create_default_resource_config( + self, + entity: str, + resource: str, + config: str, + template_variables: Optional[Dict[str, Union[float, int, str]]], + ) -> Optional[Dict[str, Any]]: + if not self.create_default_resource_config_introspection(): + raise Exception() + supports_template_vars, _ = self.push_to_run_queue_introspection() + + mutation_params = """ + $entityName: String!, + $resource: String!, + $config: JSONString! + """ + mutation_inputs = """ + entityName: $entityName, + resource: $resource, + config: $config + """ + + if supports_template_vars: + mutation_params += ", $templateVariables: JSONString" + mutation_inputs += ", templateVariables: $templateVariables" + else: + if template_variables is not None: + raise UnsupportedError( + "server does not support template variables, please update server instance to >=0.46" + ) + + variable_values = { + "entityName": entity, + "resource": resource, + "config": config, + } + if supports_template_vars: + if template_variables is not None: + variable_values["templateVariables"] = json.dumps(template_variables) + else: + variable_values["templateVariables"] = "{}" + + query = gql( + f""" + mutation createDefaultResourceConfig( + {mutation_params} + ) {{ + createDefaultResourceConfig( + input: {{ + {mutation_inputs} + }} + ) {{ + defaultResourceConfigID + success + }} + }} + """ + ) + + result: Optional[Dict[str, Any]] = self.gql(query, variable_values)[ + "createDefaultResourceConfig" + ] + return result + + @normalize_exceptions + def create_run_queue( + self, + entity: str, + project: str, + queue_name: str, + access: str, + prioritization_mode: Optional[str] = None, + config_id: Optional[str] = None, + ) -> Optional[Dict[str, Any]]: + ( + create_run_queue, + supports_drc, + supports_prioritization, + ) = self.create_run_queue_introspection() + if not create_run_queue: + raise UnsupportedError( + "run queue creation is not supported by this version of " + "wandb server. Consider updating to the latest version." + ) + if not supports_drc and config_id is not None: + raise UnsupportedError( + "default resource configurations are not supported by this version " + "of wandb server. Consider updating to the latest version." + ) + if not supports_prioritization and prioritization_mode is not None: + raise UnsupportedError( + "launch prioritization is not supported by this version of " + "wandb server. Consider updating to the latest version." + ) + + if supports_prioritization: + query = gql( + """ + mutation createRunQueue( + $entity: String!, + $project: String!, + $queueName: String!, + $access: RunQueueAccessType!, + $prioritizationMode: RunQueuePrioritizationMode, + $defaultResourceConfigID: ID, + ) { + createRunQueue( + input: { + entityName: $entity, + projectName: $project, + queueName: $queueName, + access: $access, + prioritizationMode: $prioritizationMode + defaultResourceConfigID: $defaultResourceConfigID + } + ) { + success + queueID + } + } + """ + ) + variable_values = { + "entity": entity, + "project": project, + "queueName": queue_name, + "access": access, + "prioritizationMode": prioritization_mode, + "defaultResourceConfigID": config_id, + } + else: + query = gql( + """ + mutation createRunQueue( + $entity: String!, + $project: String!, + $queueName: String!, + $access: RunQueueAccessType!, + $defaultResourceConfigID: ID, + ) { + createRunQueue( + input: { + entityName: $entity, + projectName: $project, + queueName: $queueName, + access: $access, + defaultResourceConfigID: $defaultResourceConfigID + } + ) { + success + queueID + } + } + """ + ) + variable_values = { + "entity": entity, + "project": project, + "queueName": queue_name, + "access": access, + "defaultResourceConfigID": config_id, + } + + result: Optional[Dict[str, Any]] = self.gql(query, variable_values)[ + "createRunQueue" + ] + return result + + @normalize_exceptions + def upsert_run_queue( + self, + queue_name: str, + entity: str, + resource_type: str, + resource_config: dict, + project: str = LAUNCH_DEFAULT_PROJECT, + prioritization_mode: Optional[str] = None, + template_variables: Optional[dict] = None, + external_links: Optional[dict] = None, + ) -> Optional[Dict[str, Any]]: + if not self.upsert_run_queue_introspection(): + raise UnsupportedError( + "upserting run queues is not supported by this version of " + "wandb server. Consider updating to the latest version." + ) + query = gql( + """ + mutation upsertRunQueue( + $entityName: String! + $projectName: String! + $queueName: String! + $resourceType: String! + $resourceConfig: JSONString! + $templateVariables: JSONString + $prioritizationMode: RunQueuePrioritizationMode + $externalLinks: JSONString + $clientMutationId: String + ) { + upsertRunQueue( + input: { + entityName: $entityName + projectName: $projectName + queueName: $queueName + resourceType: $resourceType + resourceConfig: $resourceConfig + templateVariables: $templateVariables + prioritizationMode: $prioritizationMode + externalLinks: $externalLinks + clientMutationId: $clientMutationId + } + ) { + success + configSchemaValidationErrors + } + } + """ + ) + variable_values = { + "entityName": entity, + "projectName": project, + "queueName": queue_name, + "resourceType": resource_type, + "resourceConfig": json.dumps(resource_config), + "templateVariables": ( + json.dumps(template_variables) if template_variables else None + ), + "prioritizationMode": prioritization_mode, + "externalLinks": json.dumps(external_links) if external_links else None, + } + result: Dict[str, Any] = self.gql(query, variable_values) + return result["upsertRunQueue"] + + @normalize_exceptions + def push_to_run_queue_by_name( + self, + entity: str, + project: str, + queue_name: str, + run_spec: str, + template_variables: Optional[Dict[str, Union[int, float, str]]], + priority: Optional[int] = None, + ) -> Optional[Dict[str, Any]]: + self.push_to_run_queue_introspection() + """Queryless mutation, should be used before legacy fallback method.""" + + mutation_params = """ + $entityName: String!, + $projectName: String!, + $queueName: String!, + $runSpec: JSONString! + """ + + mutation_input = """ + entityName: $entityName, + projectName: $projectName, + queueName: $queueName, + runSpec: $runSpec + """ + + variables: Dict[str, Any] = { + "entityName": entity, + "projectName": project, + "queueName": queue_name, + "runSpec": run_spec, + } + if self.server_push_to_run_queue_supports_priority: + if priority is not None: + variables["priority"] = priority + mutation_params += ", $priority: Int" + mutation_input += ", priority: $priority" + else: + if priority is not None: + raise UnsupportedError( + "server does not support priority, please update server instance to >=0.46" + ) + + if self.server_supports_template_variables: + if template_variables is not None: + variables.update( + {"templateVariableValues": json.dumps(template_variables)} + ) + mutation_params += ", $templateVariableValues: JSONString" + mutation_input += ", templateVariableValues: $templateVariableValues" + else: + if template_variables is not None: + raise UnsupportedError( + "server does not support template variables, please update server instance to >=0.46" + ) + + mutation = gql( + f""" + mutation pushToRunQueueByName( + {mutation_params} + ) {{ + pushToRunQueueByName( + input: {{ + {mutation_input} + }} + ) {{ + runQueueItemId + runSpec + }} + }} + """ + ) + + try: + result: Optional[Dict[str, Any]] = self.gql( + mutation, variables, check_retry_fn=util.no_retry_4xx + ).get("pushToRunQueueByName") + if not result: + return None + + if result.get("runSpec"): + run_spec = json.loads(str(result["runSpec"])) + result["runSpec"] = run_spec + + return result + except Exception as e: + if ( + 'Cannot query field "runSpec" on type "PushToRunQueueByNamePayload"' + not in str(e) + ): + return None + + mutation_no_runspec = gql( + """ + mutation pushToRunQueueByName( + $entityName: String!, + $projectName: String!, + $queueName: String!, + $runSpec: JSONString!, + ) { + pushToRunQueueByName( + input: { + entityName: $entityName, + projectName: $projectName, + queueName: $queueName, + runSpec: $runSpec + } + ) { + runQueueItemId + } + } + """ + ) + + try: + result = self.gql( + mutation_no_runspec, variables, check_retry_fn=util.no_retry_4xx + ).get("pushToRunQueueByName") + except Exception: + result = None + + return result + + @normalize_exceptions + def push_to_run_queue( + self, + queue_name: str, + launch_spec: Dict[str, str], + template_variables: Optional[dict], + project_queue: str, + priority: Optional[int] = None, + ) -> Optional[Dict[str, Any]]: + self.push_to_run_queue_introspection() + entity = launch_spec.get("queue_entity") or launch_spec["entity"] + run_spec = json.dumps(launch_spec) + + push_result = self.push_to_run_queue_by_name( + entity, project_queue, queue_name, run_spec, template_variables, priority + ) + + if push_result: + return push_result + + if priority is not None: + # Cannot proceed with legacy method if priority is set + return None + + """ Legacy Method """ + queues_found = self.get_project_run_queues(entity, project_queue) + matching_queues = [ + q + for q in queues_found + if q["name"] == queue_name + # ensure user has access to queue + and ( + # TODO: User created queues in the UI have USER access + q["access"] in ["PROJECT", "USER"] + or q["createdBy"] == self.default_entity + ) + ] + if not matching_queues: + # in the case of a missing default queue. create it + if queue_name == "default": + wandb.termlog( + f"No default queue existing for entity: {entity} in project: {project_queue}, creating one." + ) + res = self.create_run_queue( + launch_spec["entity"], + project_queue, + queue_name, + access="PROJECT", + ) + + if res is None or res.get("queueID") is None: + wandb.termerror( + f"Unable to create default queue for entity: {entity} on project: {project_queue}. Run could not be added to a queue" + ) + return None + queue_id = res["queueID"] + + else: + if project_queue == "model-registry": + _msg = f"Unable to push to run queue {queue_name}. Queue not found." + else: + _msg = f"Unable to push to run queue {project_queue}/{queue_name}. Queue not found." + wandb.termwarn(_msg) + return None + elif len(matching_queues) > 1: + wandb.termerror( + f"Unable to push to run queue {queue_name}. More than one queue found with this name." + ) + return None + else: + queue_id = matching_queues[0]["id"] + spec_json = json.dumps(launch_spec) + variables = {"queueID": queue_id, "runSpec": spec_json} + + mutation_params = """ + $queueID: ID!, + $runSpec: JSONString! + """ + mutation_input = """ + queueID: $queueID, + runSpec: $runSpec + """ + if self.server_supports_template_variables: + if template_variables is not None: + mutation_params += ", $templateVariableValues: JSONString" + mutation_input += ", templateVariableValues: $templateVariableValues" + variables.update( + {"templateVariableValues": json.dumps(template_variables)} + ) + else: + if template_variables is not None: + raise UnsupportedError( + "server does not support template variables, please update server instance to >=0.46" + ) + + mutation = gql( + f""" + mutation pushToRunQueue( + {mutation_params} + ) {{ + pushToRunQueue( + input: {{{mutation_input}}} + ) {{ + runQueueItemId + }} + }} + """ + ) + + response = self.gql(mutation, variable_values=variables) + if not response.get("pushToRunQueue"): + raise CommError(f"Error pushing run queue item to queue {queue_name}.") + + result: Optional[Dict[str, Any]] = response["pushToRunQueue"] + return result + + @normalize_exceptions + def pop_from_run_queue( + self, + queue_name: str, + entity: Optional[str] = None, + project: Optional[str] = None, + agent_id: Optional[str] = None, + ) -> Optional[Dict[str, Any]]: + mutation = gql( + """ + mutation popFromRunQueue($entity: String!, $project: String!, $queueName: String!, $launchAgentId: ID) { + popFromRunQueue(input: { + entityName: $entity, + projectName: $project, + queueName: $queueName, + launchAgentId: $launchAgentId + }) { + runQueueItemId + runSpec + } + } + """ + ) + response = self.gql( + mutation, + variable_values={ + "entity": entity, + "project": project, + "queueName": queue_name, + "launchAgentId": agent_id, + }, + ) + result: Optional[Dict[str, Any]] = response["popFromRunQueue"] + return result + + @normalize_exceptions + def ack_run_queue_item(self, item_id: str, run_id: Optional[str] = None) -> bool: + mutation = gql( + """ + mutation ackRunQueueItem($itemId: ID!, $runId: String!) { + ackRunQueueItem(input: { runQueueItemId: $itemId, runName: $runId }) { + success + } + } + """ + ) + response = self.gql( + mutation, variable_values={"itemId": item_id, "runId": str(run_id)} + ) + if not response["ackRunQueueItem"]["success"]: + raise CommError( + "Error acking run queue item. Item may have already been acknowledged by another process" + ) + result: bool = response["ackRunQueueItem"]["success"] + return result + + @normalize_exceptions + def create_launch_agent_fields_introspection(self) -> List: + if self.create_launch_agent_input_info: + return self.create_launch_agent_input_info + query_string = """ + query ProbeServerCreateLaunchAgentInput { + CreateLaunchAgentInputInfoType: __type(name:"CreateLaunchAgentInput") { + inputFields{ + name + } + } + } + """ + + query = gql(query_string) + res = self.gql(query) + + self.create_launch_agent_input_info = [ + field.get("name", "") + for field in res.get("CreateLaunchAgentInputInfoType", {}).get( + "inputFields", [{}] + ) + ] + return self.create_launch_agent_input_info + + @normalize_exceptions + def create_launch_agent( + self, + entity: str, + project: str, + queues: List[str], + agent_config: Dict[str, Any], + version: str, + gorilla_agent_support: bool, + ) -> dict: + project_queues = self.get_project_run_queues(entity, project) + if not project_queues: + # create default queue if it doesn't already exist + default = self.create_run_queue( + entity, project, "default", access="PROJECT" + ) + if default is None or default.get("queueID") is None: + raise CommError( + f"Unable to create default queue for {entity}/{project}. No queues for agent to poll" + ) + project_queues = [{"id": default["queueID"], "name": "default"}] + polling_queue_ids = [ + q["id"] for q in project_queues if q["name"] in queues + ] # filter to poll specified queues + if len(polling_queue_ids) != len(queues): + raise CommError( + f"Could not start launch agent: Not all of requested queues ({', '.join(queues)}) found. " + f"Available queues for this project: {','.join([q['name'] for q in project_queues])}" + ) + + if not gorilla_agent_support: + # if gorilla doesn't support launch agents, return a client-generated id + return { + "success": True, + "launchAgentId": None, + } + + hostname = socket.gethostname() + + variable_values = { + "entity": entity, + "project": project, + "queues": polling_queue_ids, + "hostname": hostname, + } + + mutation_params = """ + $entity: String!, + $project: String!, + $queues: [ID!]!, + $hostname: String! + """ + + mutation_input = """ + entityName: $entity, + projectName: $project, + runQueues: $queues, + hostname: $hostname + """ + + if "agentConfig" in self.create_launch_agent_fields_introspection(): + variable_values["agentConfig"] = json.dumps(agent_config) + mutation_params += ", $agentConfig: JSONString" + mutation_input += ", agentConfig: $agentConfig" + if "version" in self.create_launch_agent_fields_introspection(): + variable_values["version"] = version + mutation_params += ", $version: String" + mutation_input += ", version: $version" + + mutation = gql( + f""" + mutation createLaunchAgent( + {mutation_params} + ) {{ + createLaunchAgent( + input: {{ + {mutation_input} + }} + ) {{ + launchAgentId + }} + }} + """ + ) + result: dict = self.gql(mutation, variable_values)["createLaunchAgent"] + return result + + @normalize_exceptions + def update_launch_agent_status( + self, + agent_id: str, + status: str, + gorilla_agent_support: bool, + ) -> dict: + if not gorilla_agent_support: + # if gorilla doesn't support launch agents, this is a no-op + return { + "success": True, + } + + mutation = gql( + """ + mutation updateLaunchAgent($agentId: ID!, $agentStatus: String){ + updateLaunchAgent( + input: { + launchAgentId: $agentId + agentStatus: $agentStatus + } + ) { + success + } + } + """ + ) + variable_values = { + "agentId": agent_id, + "agentStatus": status, + } + result: dict = self.gql(mutation, variable_values)["updateLaunchAgent"] + return result + + @normalize_exceptions + def get_launch_agent(self, agent_id: str, gorilla_agent_support: bool) -> dict: + if not gorilla_agent_support: + return { + "id": None, + "name": "", + "stopPolling": False, + } + query = gql( + """ + query LaunchAgent($agentId: ID!) { + launchAgent(id: $agentId) { + id + name + runQueues + hostname + agentStatus + stopPolling + heartbeatAt + } + } + """ + ) + variable_values = { + "agentId": agent_id, + } + result: dict = self.gql(query, variable_values)["launchAgent"] + return result + + @normalize_exceptions + def upsert_run( + self, + id: Optional[str] = None, + name: Optional[str] = None, + project: Optional[str] = None, + host: Optional[str] = None, + group: Optional[str] = None, + tags: Optional[List[str]] = None, + config: Optional[dict] = None, + description: Optional[str] = None, + entity: Optional[str] = None, + state: Optional[str] = None, + display_name: Optional[str] = None, + notes: Optional[str] = None, + repo: Optional[str] = None, + job_type: Optional[str] = None, + program_path: Optional[str] = None, + commit: Optional[str] = None, + sweep_name: Optional[str] = None, + summary_metrics: Optional[str] = None, + num_retries: Optional[int] = None, + ) -> Tuple[dict, bool, Optional[List]]: + """Update a run. + + Args: + id (str, optional): The existing run to update + name (str, optional): The name of the run to create + group (str, optional): Name of the group this run is a part of + project (str, optional): The name of the project + host (str, optional): The name of the host + tags (list, optional): A list of tags to apply to the run + config (dict, optional): The latest config params + description (str, optional): A description of this project + entity (str, optional): The entity to scope this project to. + display_name (str, optional): The display name of this project + notes (str, optional): Notes about this run + repo (str, optional): Url of the program's repository. + state (str, optional): State of the program. + job_type (str, optional): Type of job, e.g 'train'. + program_path (str, optional): Path to the program. + commit (str, optional): The Git SHA to associate the run with + sweep_name (str, optional): The name of the sweep this run is a part of + summary_metrics (str, optional): The JSON summary metrics + num_retries (int, optional): Number of retries + """ + query_string = """ + mutation UpsertBucket( + $id: String, + $name: String, + $project: String, + $entity: String, + $groupName: String, + $description: String, + $displayName: String, + $notes: String, + $commit: String, + $config: JSONString, + $host: String, + $debug: Boolean, + $program: String, + $repo: String, + $jobType: String, + $state: String, + $sweep: String, + $tags: [String!], + $summaryMetrics: JSONString, + ) { + upsertBucket(input: { + id: $id, + name: $name, + groupName: $groupName, + modelName: $project, + entityName: $entity, + description: $description, + displayName: $displayName, + notes: $notes, + config: $config, + commit: $commit, + host: $host, + debug: $debug, + jobProgram: $program, + jobRepo: $repo, + jobType: $jobType, + state: $state, + sweep: $sweep, + tags: $tags, + summaryMetrics: $summaryMetrics, + }) { + bucket { + id + name + displayName + description + config + sweepName + project { + id + name + entity { + id + name + } + } + historyLineCount + } + inserted + _Server_Settings_ + } + } + """ + self.server_settings_introspection() + + server_settings_string = ( + """ + serverSettings { + serverMessages{ + utfText + plainText + htmlText + messageType + messageLevel + } + } + """ + if self._server_settings_type + else "" + ) + + query_string = query_string.replace("_Server_Settings_", server_settings_string) + mutation = gql(query_string) + config_str = json.dumps(config) if config else None + if not description or description.isspace(): + description = None + + kwargs = {} + if num_retries is not None: + kwargs["num_retries"] = num_retries + + variable_values = { + "id": id, + "entity": entity or self.settings("entity"), + "name": name, + "project": project or util.auto_project_name(program_path), + "groupName": group, + "tags": tags, + "description": description, + "config": config_str, + "commit": commit, + "displayName": display_name, + "notes": notes, + "host": None + if self.settings().get("anonymous") in ["allow", "must"] + else host, + "debug": env.is_debug(env=self._environ), + "repo": repo, + "program": program_path, + "jobType": job_type, + "state": state, + "sweep": sweep_name, + "summaryMetrics": summary_metrics, + } + + # retry conflict errors for 2 minutes, default to no_auth_retry + check_retry_fn = util.make_check_retry_fn( + check_fn=util.check_retry_conflict_or_gone, + check_timedelta=datetime.timedelta(minutes=2), + fallback_retry_fn=util.no_retry_auth, + ) + + response = self.gql( + mutation, + variable_values=variable_values, + check_retry_fn=check_retry_fn, + **kwargs, + ) + + run_obj: Dict[str, Dict[str, Dict[str, str]]] = response["upsertBucket"][ + "bucket" + ] + project_obj: Dict[str, Dict[str, str]] = run_obj.get("project", {}) + if project_obj: + self.set_setting("project", project_obj["name"]) + entity_obj = project_obj.get("entity", {}) + if entity_obj: + self.set_setting("entity", entity_obj["name"]) + + server_messages = None + if self._server_settings_type: + server_messages = ( + response["upsertBucket"] + .get("serverSettings", {}) + .get("serverMessages", []) + ) + + return ( + response["upsertBucket"]["bucket"], + response["upsertBucket"]["inserted"], + server_messages, + ) + + @normalize_exceptions + def rewind_run( + self, + run_name: str, + metric_name: str, + metric_value: float, + program_path: Optional[str] = None, + entity: Optional[str] = None, + project: Optional[str] = None, + num_retries: Optional[int] = None, + ) -> dict: + """Rewinds a run to a previous state. + + Args: + run_name (str): The name of the run to rewind + metric_name (str): The name of the metric to rewind to + metric_value (float): The value of the metric to rewind to + program_path (str, optional): Path to the program + entity (str, optional): The entity to scope this project to + project (str, optional): The name of the project + num_retries (int, optional): Number of retries + + Returns: + A dict with the rewound run + + { + "id": "run_id", + "name": "run_name", + "displayName": "run_display_name", + "description": "run_description", + "config": "stringified_run_config_json", + "sweepName": "run_sweep_name", + "project": { + "id": "project_id", + "name": "project_name", + "entity": { + "id": "entity_id", + "name": "entity_name" + } + }, + "historyLineCount": 100, + } + """ + query_string = """ + mutation RewindRun($runName: String!, $entity: String, $project: String, $metricName: String!, $metricValue: Float!) { + rewindRun(input: {runName: $runName, entityName: $entity, projectName: $project, metricName: $metricName, metricValue: $metricValue}) { + rewoundRun { + id + name + displayName + description + config + sweepName + project { + id + name + entity { + id + name + } + } + historyLineCount + } + } + } + """ + + mutation = gql(query_string) + + kwargs = {} + if num_retries is not None: + kwargs["num_retries"] = num_retries + + variable_values = { + "runName": run_name, + "entity": entity or self.settings("entity"), + "project": project or util.auto_project_name(program_path), + "metricName": metric_name, + "metricValue": metric_value, + } + + # retry conflict errors for 2 minutes, default to no_auth_retry + check_retry_fn = util.make_check_retry_fn( + check_fn=util.check_retry_conflict_or_gone, + check_timedelta=datetime.timedelta(minutes=2), + fallback_retry_fn=util.no_retry_auth, + ) + + response = self.gql( + mutation, + variable_values=variable_values, + check_retry_fn=check_retry_fn, + **kwargs, + ) + + run_obj: Dict[str, Dict[str, Dict[str, str]]] = response.get( + "rewindRun", {} + ).get("rewoundRun", {}) + project_obj: Dict[str, Dict[str, str]] = run_obj.get("project", {}) + if project_obj: + self.set_setting("project", project_obj["name"]) + entity_obj = project_obj.get("entity", {}) + if entity_obj: + self.set_setting("entity", entity_obj["name"]) + + return run_obj + + @normalize_exceptions + def get_run_info( + self, + entity: str, + project: str, + name: str, + ) -> dict: + query = gql( + """ + query RunInfo($project: String!, $entity: String!, $name: String!) { + project(name: $project, entityName: $entity) { + run(name: $name) { + runInfo { + program + args + os + python + colab + executable + codeSaved + cpuCount + gpuCount + gpu + git { + remote + commit + } + } + } + } + } + """ + ) + variable_values = {"project": project, "entity": entity, "name": name} + res = self.gql(query, variable_values) + if res.get("project") is None: + raise CommError( + f"Error fetching run info for {entity}/{project}/{name}. Check that this project exists and you have access to this entity and project" + ) + elif res["project"].get("run") is None: + raise CommError( + f"Error fetching run info for {entity}/{project}/{name}. Check that this run id exists" + ) + run_info: dict = res["project"]["run"]["runInfo"] + return run_info + + @normalize_exceptions + def get_run_state(self, entity: str, project: str, name: str) -> str: + query = gql( + """ + query RunState( + $project: String!, + $entity: String!, + $name: String!) { + project(name: $project, entityName: $entity) { + run(name: $name) { + state + } + } + } + """ + ) + variable_values = { + "project": project, + "entity": entity, + "name": name, + } + res = self.gql(query, variable_values) + if res.get("project") is None or res["project"].get("run") is None: + raise CommError(f"Error fetching run state for {entity}/{project}/{name}.") + run_state: str = res["project"]["run"]["state"] + return run_state + + @normalize_exceptions + def create_run_files_introspection(self) -> bool: + _, _, mutations = self.server_info_introspection() + return "createRunFiles" in mutations + + @normalize_exceptions + def upload_urls( + self, + project: str, + files: Union[List[str], Dict[str, IO]], + run: Optional[str] = None, + entity: Optional[str] = None, + description: Optional[str] = None, + ) -> Tuple[str, List[str], Dict[str, Dict[str, Any]]]: + """Generate temporary resumable upload urls. + + Args: + project (str): The project to download + files (list or dict): The filenames to upload + run (str, optional): The run to upload to + entity (str, optional): The entity to scope this project to. + description (str, optional): description + + Returns: + (run_id, upload_headers, file_info) + run_id: id of run we uploaded files to + upload_headers: A list of headers to use when uploading files. + file_info: A dict of filenames and urls. + { + "run_id": "run_id", + "upload_headers": [""], + "file_info": [ + { "weights.h5": { "uploadUrl": "https://weights.url" } }, + { "model.json": { "uploadUrl": "https://model.json" } } + ] + } + """ + run_name = run or self.current_run_id + assert run_name, "run must be specified" + entity = entity or self.settings("entity") + assert entity, "entity must be specified" + + has_create_run_files_mutation = self.create_run_files_introspection() + if not has_create_run_files_mutation: + return self.legacy_upload_urls(project, files, run, entity, description) + + query = gql( + """ + mutation CreateRunFiles($entity: String!, $project: String!, $run: String!, $files: [String!]!) { + createRunFiles(input: {entityName: $entity, projectName: $project, runName: $run, files: $files}) { + runID + uploadHeaders + files { + name + uploadUrl + } + } + } + """ + ) + + query_result = self.gql( + query, + variable_values={ + "project": project, + "run": run_name, + "entity": entity, + "files": [file for file in files], + }, + ) + + result = query_result["createRunFiles"] + run_id = result["runID"] + if not run_id: + raise CommError( + f"Error uploading files to {entity}/{project}/{run_name}. Check that this project exists and you have access to this entity and project" + ) + file_name_urls = {file["name"]: file for file in result["files"]} + return run_id, result["uploadHeaders"], file_name_urls + + def legacy_upload_urls( + self, + project: str, + files: Union[List[str], Dict[str, IO]], + run: Optional[str] = None, + entity: Optional[str] = None, + description: Optional[str] = None, + ) -> Tuple[str, List[str], Dict[str, Dict[str, Any]]]: + """Generate temporary resumable upload urls. + + A new mutation createRunFiles was introduced after 0.15.4. + This function is used to support older versions. + """ + query = gql( + """ + query RunUploadUrls($name: String!, $files: [String]!, $entity: String, $run: String!, $description: String) { + model(name: $name, entityName: $entity) { + bucket(name: $run, desc: $description) { + id + files(names: $files) { + uploadHeaders + edges { + node { + name + url(upload: true) + updatedAt + } + } + } + } + } + } + """ + ) + run_id = run or self.current_run_id + assert run_id, "run must be specified" + entity = entity or self.settings("entity") + query_result = self.gql( + query, + variable_values={ + "name": project, + "run": run_id, + "entity": entity, + "files": [file for file in files], + "description": description, + }, + ) + + run_obj = query_result["model"]["bucket"] + if run_obj: + for file_node in run_obj["files"]["edges"]: + file = file_node["node"] + # we previously used "url" field but now use "uploadUrl" + # replace the "url" field with "uploadUrl for downstream compatibility + if "url" in file and "uploadUrl" not in file: + file["uploadUrl"] = file.pop("url") + + result = { + file["name"]: file for file in self._flatten_edges(run_obj["files"]) + } + return run_obj["id"], run_obj["files"]["uploadHeaders"], result + else: + raise CommError(f"Run does not exist {entity}/{project}/{run_id}.") + + @normalize_exceptions + def download_urls( + self, + project: str, + run: Optional[str] = None, + entity: Optional[str] = None, + ) -> Dict[str, Dict[str, str]]: + """Generate download urls. + + Args: + project (str): The project to download + run (str): The run to upload to + entity (str, optional): The entity to scope this project to. Defaults to wandb models + + Returns: + A dict of extensions and urls + + { + 'weights.h5': { "url": "https://weights.url", "updatedAt": '2013-04-26T22:22:23.832Z', 'md5': 'mZFLkyvTelC5g8XnyQrpOw==' }, + 'model.json': { "url": "https://model.url", "updatedAt": '2013-04-26T22:22:23.832Z', 'md5': 'mZFLkyvTelC5g8XnyQrpOw==' } + } + """ + query = gql( + """ + query RunDownloadUrls($name: String!, $entity: String, $run: String!) { + model(name: $name, entityName: $entity) { + bucket(name: $run) { + files { + edges { + node { + name + url + md5 + updatedAt + } + } + } + } + } + } + """ + ) + run = run or self.current_run_id + assert run, "run must be specified" + entity = entity or self.settings("entity") + query_result = self.gql( + query, + variable_values={ + "name": project, + "run": run, + "entity": entity, + }, + ) + if query_result["model"] is None: + raise CommError(f"Run does not exist {entity}/{project}/{run}.") + files = self._flatten_edges(query_result["model"]["bucket"]["files"]) + return {file["name"]: file for file in files if file} + + @normalize_exceptions + def download_url( + self, + project: str, + file_name: str, + run: Optional[str] = None, + entity: Optional[str] = None, + ) -> Optional[Dict[str, str]]: + """Generate download urls. + + Args: + project (str): The project to download + file_name (str): The name of the file to download + run (str): The run to upload to + entity (str, optional): The entity to scope this project to. Defaults to wandb models + + Returns: + A dict of extensions and urls + + { "url": "https://weights.url", "updatedAt": '2013-04-26T22:22:23.832Z', 'md5': 'mZFLkyvTelC5g8XnyQrpOw==' } + + """ + query = gql( + """ + query RunDownloadUrl($name: String!, $fileName: String!, $entity: String, $run: String!) { + model(name: $name, entityName: $entity) { + bucket(name: $run) { + files(names: [$fileName]) { + edges { + node { + name + url + md5 + updatedAt + } + } + } + } + } + } + """ + ) + run = run or self.current_run_id + assert run, "run must be specified" + query_result = self.gql( + query, + variable_values={ + "name": project, + "run": run, + "fileName": file_name, + "entity": entity or self.settings("entity"), + }, + ) + if query_result["model"]: + files = self._flatten_edges(query_result["model"]["bucket"]["files"]) + return files[0] if len(files) > 0 and files[0].get("updatedAt") else None + else: + return None + + @normalize_exceptions + def download_file(self, url: str) -> Tuple[int, requests.Response]: + """Initiate a streaming download. + + Args: + url (str): The url to download + + Returns: + A tuple of the content length and the streaming response + """ + check_httpclient_logger_handler() + + http_headers = _thread_local_api_settings.headers or {} + + auth = None + if self.access_token is not None: + http_headers["Authorization"] = f"Bearer {self.access_token}" + elif _thread_local_api_settings.cookies is None: + auth = ("api", self.api_key or "") + + response = requests.get( + url, + auth=auth, + cookies=_thread_local_api_settings.cookies or {}, + headers=http_headers, + stream=True, + ) + response.raise_for_status() + return int(response.headers.get("content-length", 0)), response + + @normalize_exceptions + def download_write_file( + self, + metadata: Dict[str, str], + out_dir: Optional[str] = None, + ) -> Tuple[str, Optional[requests.Response]]: + """Download a file from a run and write it to wandb/. + + Args: + metadata (obj): The metadata object for the file to download. Comes from Api.download_urls(). + out_dir (str, optional): The directory to write the file to. Defaults to wandb/ + + Returns: + A tuple of the file's local path and the streaming response. The streaming response is None if the file + already existed and was up-to-date. + """ + filename = metadata["name"] + path = os.path.join(out_dir or self.settings("wandb_dir"), filename) + if self.file_current(filename, B64MD5(metadata["md5"])): + return path, None + + size, response = self.download_file(metadata["url"]) + + with util.fsync_open(path, "wb") as file: + for data in response.iter_content(chunk_size=1024): + file.write(data) + + return path, response + + def upload_file_azure( + self, url: str, file: Any, extra_headers: Dict[str, str] + ) -> None: + """Upload a file to azure.""" + from azure.core.exceptions import AzureError # type: ignore + + # Configure the client without retries so our existing logic can handle them + client = self._azure_blob_module.BlobClient.from_blob_url( + url, retry_policy=self._azure_blob_module.LinearRetry(retry_total=0) + ) + try: + if extra_headers.get("Content-MD5") is not None: + md5: Optional[bytes] = base64.b64decode(extra_headers["Content-MD5"]) + else: + md5 = None + content_settings = self._azure_blob_module.ContentSettings( + content_md5=md5, + content_type=extra_headers.get("Content-Type"), + ) + client.upload_blob( + file, + max_concurrency=4, + length=len(file), + overwrite=True, + content_settings=content_settings, + ) + except AzureError as e: + if hasattr(e, "response"): + response = requests.models.Response() + response.status_code = e.response.status_code + response.headers = e.response.headers + raise requests.exceptions.RequestException(e.message, response=response) + else: + raise requests.exceptions.ConnectionError(e.message) + + def upload_multipart_file_chunk( + self, + url: str, + upload_chunk: bytes, + extra_headers: Optional[Dict[str, str]] = None, + ) -> Optional[requests.Response]: + """Upload a file chunk to S3 with failure resumption. + + Args: + url: The url to download + upload_chunk: The path to the file you want to upload + extra_headers: A dictionary of extra headers to send with the request + + Returns: + The `requests` library response object + """ + check_httpclient_logger_handler() + try: + if env.is_debug(env=self._environ): + logger.debug("upload_file: %s", url) + response = self._upload_file_session.put( + url, data=upload_chunk, headers=extra_headers + ) + if env.is_debug(env=self._environ): + logger.debug("upload_file: %s complete", url) + response.raise_for_status() + except requests.exceptions.RequestException as e: + logger.exception(f"upload_file exception for {url=}") + response_content = e.response.content if e.response is not None else "" + status_code = e.response.status_code if e.response is not None else 0 + # S3 reports retryable request timeouts out-of-band + is_aws_retryable = status_code == 400 and "RequestTimeout" in str( + response_content + ) + # Retry errors from cloud storage or local network issues + if ( + status_code in (308, 408, 409, 429, 500, 502, 503, 504) + or isinstance( + e, + (requests.exceptions.Timeout, requests.exceptions.ConnectionError), + ) + or is_aws_retryable + ): + _e = retry.TransientError(exc=e) + raise _e.with_traceback(sys.exc_info()[2]) + else: + wandb._sentry.reraise(e) + return response + + def upload_file( + self, + url: str, + file: IO[bytes], + callback: Optional["ProgressFn"] = None, + extra_headers: Optional[Dict[str, str]] = None, + ) -> Optional[requests.Response]: + """Upload a file to W&B with failure resumption. + + Args: + url: The url to download + file: The path to the file you want to upload + callback: A callback which is passed the number of + bytes uploaded since the last time it was called, used to report progress + extra_headers: A dictionary of extra headers to send with the request + + Returns: + The `requests` library response object + """ + check_httpclient_logger_handler() + extra_headers = extra_headers.copy() if extra_headers else {} + response: Optional[requests.Response] = None + progress = Progress(file, callback=callback) + try: + if "x-ms-blob-type" in extra_headers and self._azure_blob_module: + self.upload_file_azure(url, progress, extra_headers) + else: + if "x-ms-blob-type" in extra_headers: + wandb.termwarn( + "Azure uploads over 256MB require the azure SDK, install with pip install wandb[azure]", + repeat=False, + ) + if env.is_debug(env=self._environ): + logger.debug("upload_file: %s", url) + response = self._upload_file_session.put( + url, data=progress, headers=extra_headers + ) + if env.is_debug(env=self._environ): + logger.debug("upload_file: %s complete", url) + response.raise_for_status() + except requests.exceptions.RequestException as e: + logger.exception(f"upload_file exception for {url=}") + response_content = e.response.content if e.response is not None else "" + status_code = e.response.status_code if e.response is not None else 0 + # S3 reports retryable request timeouts out-of-band + is_aws_retryable = ( + "x-amz-meta-md5" in extra_headers + and status_code == 400 + and "RequestTimeout" in str(response_content) + ) + # We need to rewind the file for the next retry (the file passed in is `seek`'ed to 0) + progress.rewind() + # Retry errors from cloud storage or local network issues + if ( + status_code in (308, 408, 409, 429, 500, 502, 503, 504) + or isinstance( + e, + (requests.exceptions.Timeout, requests.exceptions.ConnectionError), + ) + or is_aws_retryable + ): + _e = retry.TransientError(exc=e) + raise _e.with_traceback(sys.exc_info()[2]) + else: + wandb._sentry.reraise(e) + + return response + + @normalize_exceptions + def register_agent( + self, + host: str, + sweep_id: Optional[str] = None, + project_name: Optional[str] = None, + entity: Optional[str] = None, + ) -> dict: + """Register a new agent. + + Args: + host (str): hostname + sweep_id (str): sweep id + project_name: (str): model that contains sweep + entity: (str): entity that contains sweep + """ + mutation = gql( + """ + mutation CreateAgent( + $host: String! + $projectName: String, + $entityName: String, + $sweep: String! + ) { + createAgent(input: { + host: $host, + projectName: $projectName, + entityName: $entityName, + sweep: $sweep, + }) { + agent { + id + } + } + } + """ + ) + if entity is None: + entity = self.settings("entity") + if project_name is None: + project_name = self.settings("project") + + response = self.gql( + mutation, + variable_values={ + "host": host, + "entityName": entity, + "projectName": project_name, + "sweep": sweep_id, + }, + check_retry_fn=util.no_retry_4xx, + ) + result: dict = response["createAgent"]["agent"] + return result + + def agent_heartbeat( + self, agent_id: str, metrics: dict, run_states: dict + ) -> List[Dict[str, Any]]: + """Notify server about agent state, receive commands. + + Args: + agent_id (str): agent_id + metrics (dict): system metrics + run_states (dict): run_id: state mapping + Returns: + List of commands to execute. + """ + mutation = gql( + """ + mutation Heartbeat( + $id: ID!, + $metrics: JSONString, + $runState: JSONString + ) { + agentHeartbeat(input: { + id: $id, + metrics: $metrics, + runState: $runState + }) { + agent { + id + } + commands + } + } + """ + ) + + if agent_id is None: + raise ValueError("Cannot call heartbeat with an unregistered agent.") + + try: + response = self.gql( + mutation, + variable_values={ + "id": agent_id, + "metrics": json.dumps(metrics), + "runState": json.dumps(run_states), + }, + timeout=60, + ) + except Exception: + logger.exception("Error communicating with W&B.") + return [] + else: + result: List[Dict[str, Any]] = json.loads( + response["agentHeartbeat"]["commands"] + ) + return result + + @staticmethod + def _validate_config_and_fill_distribution(config: dict) -> dict: + # verify that parameters are well specified. + # TODO(dag): deprecate this in favor of jsonschema validation once + # apiVersion 2 is released and local controller is integrated with + # wandb/client. + + # avoid modifying the original config dict in + # case it is reused outside the calling func + config = deepcopy(config) + + # explicitly cast to dict in case config was passed as a sweepconfig + # sweepconfig does not serialize cleanly to yaml and breaks graphql, + # but it is a subclass of dict, so this conversion is clean + config = dict(config) + + if "parameters" not in config: + # still shows an anaconda warning, but doesn't error + return config + + for parameter_name in config["parameters"]: + parameter = config["parameters"][parameter_name] + if "min" in parameter and "max" in parameter: + if "distribution" not in parameter: + if isinstance(parameter["min"], int) and isinstance( + parameter["max"], int + ): + parameter["distribution"] = "int_uniform" + elif isinstance(parameter["min"], float) and isinstance( + parameter["max"], float + ): + parameter["distribution"] = "uniform" + else: + raise ValueError( + f"Parameter {parameter_name} is ambiguous, please specify bounds as both floats (for a float_" + "uniform distribution) or ints (for an int_uniform distribution)." + ) + return config + + @normalize_exceptions + def upsert_sweep( + self, + config: dict, + controller: Optional[str] = None, + launch_scheduler: Optional[str] = None, + scheduler: Optional[str] = None, + obj_id: Optional[str] = None, + project: Optional[str] = None, + entity: Optional[str] = None, + state: Optional[str] = None, + prior_runs: Optional[List[str]] = None, + template_variable_values: Optional[Dict[str, Any]] = None, + ) -> Tuple[str, List[str]]: + """Upsert a sweep object. + + Args: + config (dict): sweep config (will be converted to yaml) + controller (str): controller to use + launch_scheduler (str): launch scheduler to use + scheduler (str): scheduler to use + obj_id (str): object id + project (str): project to use + entity (str): entity to use + state (str): state + prior_runs (list): IDs of existing runs to add to the sweep + template_variable_values (dict): template variable values + """ + project_query = """ + project { + id + name + entity { + id + name + } + } + """ + mutation_str = """ + mutation UpsertSweep( + $id: ID, + $config: String, + $description: String, + $entityName: String, + $projectName: String, + $controller: JSONString, + $scheduler: JSONString, + $state: String, + $priorRunsFilters: JSONString, + ) { + upsertSweep(input: { + id: $id, + config: $config, + description: $description, + entityName: $entityName, + projectName: $projectName, + controller: $controller, + scheduler: $scheduler, + state: $state, + priorRunsFilters: $priorRunsFilters, + }) { + sweep { + name + _PROJECT_QUERY_ + } + configValidationWarnings + } + } + """ + # TODO(jhr): we need protocol versioning to know schema is not supported + # for now we will just try both new and old query + mutation_5 = gql( + mutation_str.replace( + "$controller: JSONString,", + "$controller: JSONString,$launchScheduler: JSONString, $templateVariableValues: JSONString,", + ) + .replace( + "controller: $controller,", + "controller: $controller,launchScheduler: $launchScheduler,templateVariableValues: $templateVariableValues,", + ) + .replace("_PROJECT_QUERY_", project_query) + ) + # launchScheduler was introduced in core v0.14.0 + mutation_4 = gql( + mutation_str.replace( + "$controller: JSONString,", + "$controller: JSONString,$launchScheduler: JSONString,", + ) + .replace( + "controller: $controller,", + "controller: $controller,launchScheduler: $launchScheduler", + ) + .replace("_PROJECT_QUERY_", project_query) + ) + + # mutation 3 maps to backend that can support CLI version of at least 0.10.31 + mutation_3 = gql(mutation_str.replace("_PROJECT_QUERY_", project_query)) + mutation_2 = gql( + mutation_str.replace("_PROJECT_QUERY_", project_query).replace( + "configValidationWarnings", "" + ) + ) + mutation_1 = gql( + mutation_str.replace("_PROJECT_QUERY_", "").replace( + "configValidationWarnings", "" + ) + ) + + # TODO(dag): replace this with a query for protocol versioning + mutations = [mutation_5, mutation_4, mutation_3, mutation_2, mutation_1] + + config = self._validate_config_and_fill_distribution(config) + + # Silly, but attr-dicts like EasyDicts don't serialize correctly to yaml. + # This sanitizes them with a round trip pass through json to get a regular dict. + config_str = yaml.dump( + json.loads(json.dumps(config)), Dumper=util.NonOctalStringDumper + ) + filters = None + if prior_runs: + filters = json.dumps({"$or": [{"name": r} for r in prior_runs]}) + + err: Optional[Exception] = None + for mutation in mutations: + try: + variables = { + "id": obj_id, + "config": config_str, + "description": config.get("description"), + "entityName": entity or self.settings("entity"), + "projectName": project or self.settings("project"), + "controller": controller, + "launchScheduler": launch_scheduler, + "templateVariableValues": json.dumps(template_variable_values), + "scheduler": scheduler, + "priorRunsFilters": filters, + } + if state: + variables["state"] = state + + response = self.gql( + mutation, + variable_values=variables, + check_retry_fn=util.no_retry_4xx, + ) + except UsageError: + raise + except Exception as e: + # graphql schema exception is generic + err = e + continue + err = None + break + if err: + raise err + + sweep: Dict[str, Dict[str, Dict]] = response["upsertSweep"]["sweep"] + project_obj: Dict[str, Dict] = sweep.get("project", {}) + if project_obj: + self.set_setting("project", project_obj["name"]) + entity_obj: dict = project_obj.get("entity", {}) + if entity_obj: + self.set_setting("entity", entity_obj["name"]) + + warnings = response["upsertSweep"].get("configValidationWarnings", []) + return response["upsertSweep"]["sweep"]["name"], warnings + + @normalize_exceptions + def create_anonymous_api_key(self) -> str: + """Create a new API key belonging to a new anonymous user.""" + mutation = gql( + """ + mutation CreateAnonymousApiKey { + createAnonymousEntity(input: {}) { + apiKey { + name + } + } + } + """ + ) + + response = self.gql(mutation, variable_values={}) + key: str = str(response["createAnonymousEntity"]["apiKey"]["name"]) + return key + + @staticmethod + def file_current(fname: str, md5: B64MD5) -> bool: + """Checksum a file and compare the md5 with the known md5.""" + return os.path.isfile(fname) and md5_file_b64(fname) == md5 + + @normalize_exceptions + def pull( + self, project: str, run: Optional[str] = None, entity: Optional[str] = None + ) -> "List[requests.Response]": + """Download files from W&B. + + Args: + project (str): The project to download + run (str, optional): The run to upload to + entity (str, optional): The entity to scope this project to. Defaults to wandb models + + Returns: + The `requests` library response object + """ + project, run = self.parse_slug(project, run=run) + urls = self.download_urls(project, run, entity) + responses = [] + for filename in urls: + _, response = self.download_write_file(urls[filename]) + if response: + responses.append(response) + + return responses + + def get_project(self) -> str: + project: str = self.default_settings.get("project") or self.settings("project") + return project + + @normalize_exceptions + def push( + self, + files: Union[List[str], Dict[str, IO]], + run: Optional[str] = None, + entity: Optional[str] = None, + project: Optional[str] = None, + description: Optional[str] = None, + force: bool = True, + progress: Union[TextIO, Literal[False]] = False, + ) -> "List[Optional[requests.Response]]": + """Uploads multiple files to W&B. + + Args: + files (list or dict): The filenames to upload, when dict the values are open files + run (str, optional): The run to upload to + entity (str, optional): The entity to scope this project to. Defaults to wandb models + project (str, optional): The name of the project to upload to. Defaults to the one in settings. + description (str, optional): The description of the changes + force (bool, optional): Whether to prevent push if git has uncommitted changes + progress (callable, or stream): If callable, will be called with (chunk_bytes, + total_bytes) as argument. If TextIO, renders a progress bar to it. + + Returns: + A list of `requests.Response` objects + """ + if project is None: + project = self.get_project() + if project is None: + raise CommError("No project configured.") + if run is None: + run = self.current_run_id + + # TODO(adrian): we use a retriable version of self.upload_file() so + # will never retry self.upload_urls() here. Instead, maybe we should + # make push itself retriable. + _, upload_headers, result = self.upload_urls( + project, + files, + run, + entity, + ) + extra_headers = {} + for upload_header in upload_headers: + key, val = upload_header.split(":", 1) + extra_headers[key] = val + responses = [] + for file_name, file_info in result.items(): + file_url = file_info["uploadUrl"] + + # If the upload URL is relative, fill it in with the base URL, + # since it's a proxied file store like the on-prem VM. + if file_url.startswith("/"): + file_url = f"{self.api_url}{file_url}" + + try: + # To handle Windows paths + # TODO: this doesn't handle absolute paths... + normal_name = os.path.join(*file_name.split("/")) + open_file = ( + files[file_name] + if isinstance(files, dict) + else open(normal_name, "rb") + ) + except OSError: + print(f"{file_name} does not exist") # noqa: T201 + continue + if progress is False: + responses.append( + self.upload_file_retry( + file_info["uploadUrl"], open_file, extra_headers=extra_headers + ) + ) + else: + if callable(progress): + responses.append( # type: ignore + self.upload_file_retry( + file_url, open_file, progress, extra_headers=extra_headers + ) + ) + else: + length = os.fstat(open_file.fileno()).st_size + with click.progressbar( # type: ignore + file=progress, + length=length, + label=f"Uploading file: {file_name}", + fill_char=click.style("&", fg="green"), + ) as bar: + responses.append( + self.upload_file_retry( + file_url, + open_file, + lambda bites, _: bar.update(bites), + extra_headers=extra_headers, + ) + ) + open_file.close() + return responses + + def link_artifact( + self, + client_id: str, + server_id: str, + portfolio_name: str, + entity: str, + project: str, + aliases: Sequence[str], + organization: str, + ) -> Dict[str, Any]: + template = """ + mutation LinkArtifact( + $artifactPortfolioName: String!, + $entityName: String!, + $projectName: String!, + $aliases: [ArtifactAliasInput!], + ID_TYPE + ) { + linkArtifact(input: { + artifactPortfolioName: $artifactPortfolioName, + entityName: $entityName, + projectName: $projectName, + aliases: $aliases, + ID_VALUE + }) { + versionIndex + } + } + """ + + org_entity = "" + if is_artifact_registry_project(project): + try: + org_entity = self._resolve_org_entity_name( + entity=entity, organization=organization + ) + except ValueError as e: + wandb.termerror(str(e)) + raise + + def replace(a: str, b: str) -> None: + nonlocal template + template = template.replace(a, b) + + if server_id: + replace("ID_TYPE", "$artifactID: ID") + replace("ID_VALUE", "artifactID: $artifactID") + elif client_id: + replace("ID_TYPE", "$clientID: ID") + replace("ID_VALUE", "clientID: $clientID") + + variable_values = { + "clientID": client_id, + "artifactID": server_id, + "artifactPortfolioName": portfolio_name, + "entityName": org_entity or entity, + "projectName": project, + "aliases": [ + {"alias": alias, "artifactCollectionName": portfolio_name} + for alias in aliases + ], + } + + mutation = gql(template) + response = self.gql(mutation, variable_values=variable_values) + link_artifact: Dict[str, Any] = response["linkArtifact"] + return link_artifact + + def _resolve_org_entity_name(self, entity: str, organization: str = "") -> str: + # resolveOrgEntityName fetches the portfolio's org entity's name. + # + # The organization parameter may be empty, an org's display name, or an org entity name. + # + # If the server doesn't support fetching the org name of a portfolio, then this returns + # the organization parameter, or an error if it is empty. Otherwise, this returns the + # fetched value after validating that the given organization, if not empty, matches + # either the org's display or entity name. + + if not entity: + raise ValueError("Entity name is required to resolve org entity name.") + + org_fields = self.server_organization_type_introspection() + can_shorthand_org_entity = "orgEntity" in org_fields + if not organization and not can_shorthand_org_entity: + raise ValueError( + "Fetching Registry artifacts without inputting an organization " + "is unavailable for your server version. " + "Please upgrade your server to 0.50.0 or later." + ) + if not can_shorthand_org_entity: + # Server doesn't support fetching org entity to validate, + # assume org entity is correctly inputted + return organization + + orgs_from_entity = self._fetch_orgs_and_org_entities_from_entity(entity) + if organization: + return _match_org_with_fetched_org_entities(organization, orgs_from_entity) + + # If no input organization provided, error if entity belongs to multiple orgs because we + # cannot determine which one to use. + if len(orgs_from_entity) > 1: + raise ValueError( + f"Personal entity {entity!r} belongs to multiple organizations " + "and cannot be used without specifying the organization name. " + "Please specify the organization in the Registry path or use a team entity in the entity settings." + ) + return orgs_from_entity[0].entity_name + + def _fetch_orgs_and_org_entities_from_entity(self, entity: str) -> List[_OrgNames]: + """Fetches organization entity names and display names for a given entity. + + Args: + entity (str): Entity name to lookup. Can be either a personal or team entity. + + Returns: + List[_OrgNames]: List of _OrgNames tuples. (_OrgNames(entity_name, display_name)) + + Raises: + ValueError: If entity is not found, has no organizations, or other validation errors. + """ + query = gql( + """ + query FetchOrgEntityFromEntity($entityName: String!) { + entity(name: $entityName) { + organization { + name + orgEntity { + name + } + } + user { + organizations { + name + orgEntity { + name + } + } + } + } + } + """ + ) + response = self.gql( + query, + variable_values={ + "entityName": entity, + }, + ) + + # Parse organization from response + entity_resp = response["entity"]["organization"] + user_resp = response["entity"]["user"] + # Check for organization under team/org entity type + if entity_resp: + org_name = entity_resp.get("name") + org_entity_name = entity_resp.get("orgEntity") and entity_resp[ + "orgEntity" + ].get("name") + if not org_name or not org_entity_name: + raise ValueError( + f"Unable to find an organization under entity {entity!r}." + ) + return [_OrgNames(entity_name=org_entity_name, display_name=org_name)] + # Check for organization under personal entity type, where a user can belong to multiple orgs + elif user_resp: + orgs = user_resp.get("organizations", []) + org_entities_return = [ + _OrgNames( + entity_name=org["orgEntity"]["name"], display_name=org["name"] + ) + for org in orgs + if org.get("orgEntity") and org.get("name") + ] + if not org_entities_return: + raise ValueError( + f"Unable to resolve an organization associated with personal entity: {entity!r}. " + "This could be because its a personal entity that doesn't belong to any organizations. " + "Please specify the organization in the Registry path or use a team entity in the entity settings." + ) + return org_entities_return + else: + raise ValueError(f"Unable to find an organization under entity {entity!r}.") + + def _construct_use_artifact_query( + self, + artifact_id: str, + entity_name: Optional[str] = None, + project_name: Optional[str] = None, + run_name: Optional[str] = None, + use_as: Optional[str] = None, + artifact_entity_name: Optional[str] = None, + artifact_project_name: Optional[str] = None, + ) -> Tuple[Document, Dict[str, Any]]: + query_vars = [ + "$entityName: String!", + "$projectName: String!", + "$runName: String!", + "$artifactID: ID!", + ] + query_args = [ + "entityName: $entityName", + "projectName: $projectName", + "runName: $runName", + "artifactID: $artifactID", + ] + + artifact_types = self.server_use_artifact_input_introspection() + if "usedAs" in artifact_types and use_as: + query_vars.append("$usedAs: String") + query_args.append("usedAs: $usedAs") + + entity_name = entity_name or self.settings("entity") + project_name = project_name or self.settings("project") + run_name = run_name or self.current_run_id + + variable_values: Dict[str, Any] = { + "entityName": entity_name, + "projectName": project_name, + "runName": run_name, + "artifactID": artifact_id, + "usedAs": use_as, + } + + server_allows_entity_project_information = self._server_supports( + ServerFeature.USE_ARTIFACT_WITH_ENTITY_AND_PROJECT_INFORMATION + ) + if server_allows_entity_project_information: + query_vars.extend( + [ + "$artifactEntityName: String", + "$artifactProjectName: String", + ] + ) + query_args.extend( + [ + "artifactEntityName: $artifactEntityName", + "artifactProjectName: $artifactProjectName", + ] + ) + variable_values["artifactEntityName"] = artifact_entity_name + variable_values["artifactProjectName"] = artifact_project_name + + vars_str = ", ".join(query_vars) + args_str = ", ".join(query_args) + + query = gql( + f""" + mutation UseArtifact({vars_str}) {{ + useArtifact(input: {{{args_str}}}) {{ + artifact {{ + id + digest + description + state + createdAt + metadata + }} + }} + }} + """ + ) + return query, variable_values + + def use_artifact( + self, + artifact_id: str, + entity_name: Optional[str] = None, + project_name: Optional[str] = None, + run_name: Optional[str] = None, + artifact_entity_name: Optional[str] = None, + artifact_project_name: Optional[str] = None, + use_as: Optional[str] = None, + ) -> Optional[Dict[str, Any]]: + query, variable_values = self._construct_use_artifact_query( + artifact_id, + entity_name, + project_name, + run_name, + use_as, + artifact_entity_name, + artifact_project_name, + ) + response = self.gql(query, variable_values) + + if response["useArtifact"]["artifact"]: + artifact: Dict[str, Any] = response["useArtifact"]["artifact"] + return artifact + return None + + # Fetch fields available in backend of Organization type + def server_organization_type_introspection(self) -> List[str]: + query_string = """ + query ProbeServerOrganization { + OrganizationInfoType: __type(name:"Organization") { + fields { + name + } + } + } + """ + + if self.server_organization_type_fields_info is None: + query = gql(query_string) + res = self.gql(query) + input_fields = res.get("OrganizationInfoType", {}).get("fields", [{}]) + self.server_organization_type_fields_info = [ + field["name"] for field in input_fields if "name" in field + ] + + return self.server_organization_type_fields_info + + # Fetch input arguments for the "artifact" endpoint on the "Project" type + def server_project_type_introspection(self) -> bool: + if self.server_supports_enabling_artifact_usage_tracking is not None: + return self.server_supports_enabling_artifact_usage_tracking + + query_string = """ + query ProbeServerProjectInfo { + ProjectInfoType: __type(name:"Project") { + fields { + name + args { + name + } + } + } + } + """ + + query = gql(query_string) + res = self.gql(query) + input_fields = res.get("ProjectInfoType", {}).get("fields", [{}]) + artifact_args: List[Dict[str, str]] = next( + ( + field.get("args", []) + for field in input_fields + if field.get("name") == "artifact" + ), + [], + ) + self.server_supports_enabling_artifact_usage_tracking = any( + arg.get("name") == "enableTracking" for arg in artifact_args + ) + + return self.server_supports_enabling_artifact_usage_tracking + + def create_artifact_type( + self, + artifact_type_name: str, + entity_name: Optional[str] = None, + project_name: Optional[str] = None, + description: Optional[str] = None, + ) -> Optional[str]: + mutation = gql( + """ + mutation CreateArtifactType( + $entityName: String!, + $projectName: String!, + $artifactTypeName: String!, + $description: String + ) { + createArtifactType(input: { + entityName: $entityName, + projectName: $projectName, + name: $artifactTypeName, + description: $description + }) { + artifactType { + id + } + } + } + """ + ) + entity_name = entity_name or self.settings("entity") + project_name = project_name or self.settings("project") + response = self.gql( + mutation, + variable_values={ + "entityName": entity_name, + "projectName": project_name, + "artifactTypeName": artifact_type_name, + "description": description, + }, + ) + _id: Optional[str] = response["createArtifactType"]["artifactType"]["id"] + return _id + + def server_artifact_introspection(self) -> List[str]: + query_string = """ + query ProbeServerArtifact { + ArtifactInfoType: __type(name:"Artifact") { + fields { + name + } + } + } + """ + + if self.server_artifact_fields_info is None: + query = gql(query_string) + res = self.gql(query) + input_fields = res.get("ArtifactInfoType", {}).get("fields", [{}]) + self.server_artifact_fields_info = [ + field["name"] for field in input_fields if "name" in field + ] + + return self.server_artifact_fields_info + + def server_create_artifact_introspection(self) -> List[str]: + query_string = """ + query ProbeServerCreateArtifactInput { + CreateArtifactInputInfoType: __type(name:"CreateArtifactInput") { + inputFields{ + name + } + } + } + """ + + if self.server_create_artifact_input_info is None: + query = gql(query_string) + res = self.gql(query) + input_fields = res.get("CreateArtifactInputInfoType", {}).get( + "inputFields", [{}] + ) + self.server_create_artifact_input_info = [ + field["name"] for field in input_fields if "name" in field + ] + + return self.server_create_artifact_input_info + + def _get_create_artifact_mutation( + self, + fields: List, + history_step: Optional[int], + distributed_id: Optional[str], + ) -> str: + types = "" + values = "" + + if "historyStep" in fields and history_step not in [0, None]: + types += "$historyStep: Int64!," + values += "historyStep: $historyStep," + + if distributed_id: + types += "$distributedID: String," + values += "distributedID: $distributedID," + + if "clientID" in fields: + types += "$clientID: ID," + values += "clientID: $clientID," + + if "sequenceClientID" in fields: + types += "$sequenceClientID: ID," + values += "sequenceClientID: $sequenceClientID," + + if "enableDigestDeduplication" in fields: + values += "enableDigestDeduplication: true," + + if "ttlDurationSeconds" in fields: + types += "$ttlDurationSeconds: Int64," + values += "ttlDurationSeconds: $ttlDurationSeconds," + + if "tags" in fields: + types += "$tags: [TagInput!]," + values += "tags: $tags," + + query_template = """ + mutation CreateArtifact( + $artifactTypeName: String!, + $artifactCollectionNames: [String!], + $entityName: String!, + $projectName: String!, + $runName: String, + $description: String, + $digest: String!, + $aliases: [ArtifactAliasInput!], + $metadata: JSONString, + _CREATE_ARTIFACT_ADDITIONAL_TYPE_ + ) { + createArtifact(input: { + artifactTypeName: $artifactTypeName, + artifactCollectionNames: $artifactCollectionNames, + entityName: $entityName, + projectName: $projectName, + runName: $runName, + description: $description, + digest: $digest, + digestAlgorithm: MANIFEST_MD5, + aliases: $aliases, + metadata: $metadata, + _CREATE_ARTIFACT_ADDITIONAL_VALUE_ + }) { + artifact { + id + state + artifactSequence { + id + latestArtifact { + id + versionIndex + } + } + } + } + } + """ + + return query_template.replace( + "_CREATE_ARTIFACT_ADDITIONAL_TYPE_", types + ).replace("_CREATE_ARTIFACT_ADDITIONAL_VALUE_", values) + + def create_artifact( + self, + artifact_type_name: str, + artifact_collection_name: str, + digest: str, + client_id: Optional[str] = None, + sequence_client_id: Optional[str] = None, + entity_name: Optional[str] = None, + project_name: Optional[str] = None, + run_name: Optional[str] = None, + description: Optional[str] = None, + metadata: Optional[Dict] = None, + ttl_duration_seconds: Optional[int] = None, + aliases: Optional[List[Dict[str, str]]] = None, + tags: Optional[List[Dict[str, str]]] = None, + distributed_id: Optional[str] = None, + is_user_created: Optional[bool] = False, + history_step: Optional[int] = None, + ) -> Tuple[Dict, Dict]: + fields = self.server_create_artifact_introspection() + artifact_fields = self.server_artifact_introspection() + if ("ttlIsInherited" not in artifact_fields) and ttl_duration_seconds: + wandb.termwarn( + "Server not compatible with setting Artifact TTLs, please upgrade the server to use Artifact TTL" + ) + # ttlDurationSeconds is only usable if ttlIsInherited is also present + ttl_duration_seconds = None + if ("tags" not in artifact_fields) and tags: + wandb.termwarn( + "Server not compatible with Artifact tags. " + "To use Artifact tags, please upgrade the server to v0.85 or higher." + ) + + query_template = self._get_create_artifact_mutation( + fields, history_step, distributed_id + ) + + entity_name = entity_name or self.settings("entity") + project_name = project_name or self.settings("project") + if not is_user_created: + run_name = run_name or self.current_run_id + + mutation = gql(query_template) + response = self.gql( + mutation, + variable_values={ + "entityName": entity_name, + "projectName": project_name, + "runName": run_name, + "artifactTypeName": artifact_type_name, + "artifactCollectionNames": [artifact_collection_name], + "clientID": client_id, + "sequenceClientID": sequence_client_id, + "digest": digest, + "description": description, + "aliases": list(aliases or []), + "tags": list(tags or []), + "metadata": json.dumps(util.make_safe_for_json(metadata)) + if metadata + else None, + "ttlDurationSeconds": ttl_duration_seconds, + "distributedID": distributed_id, + "historyStep": history_step, + }, + ) + av = response["createArtifact"]["artifact"] + latest = response["createArtifact"]["artifact"]["artifactSequence"].get( + "latestArtifact" + ) + return av, latest + + def commit_artifact(self, artifact_id: str) -> "_Response": + mutation = gql( + """ + mutation CommitArtifact( + $artifactID: ID!, + ) { + commitArtifact(input: { + artifactID: $artifactID, + }) { + artifact { + id + digest + } + } + } + """ + ) + + response: _Response = self.gql( + mutation, + variable_values={"artifactID": artifact_id}, + timeout=60, + ) + return response + + def complete_multipart_upload_artifact( + self, + artifact_id: str, + storage_path: str, + completed_parts: List[Dict[str, Any]], + upload_id: Optional[str], + complete_multipart_action: str = "Complete", + ) -> Optional[str]: + mutation = gql( + """ + mutation CompleteMultipartUploadArtifact( + $completeMultipartAction: CompleteMultipartAction!, + $completedParts: [UploadPartsInput!]!, + $artifactID: ID! + $storagePath: String! + $uploadID: String! + ) { + completeMultipartUploadArtifact( + input: { + completeMultipartAction: $completeMultipartAction, + completedParts: $completedParts, + artifactID: $artifactID, + storagePath: $storagePath + uploadID: $uploadID + } + ) { + digest + } + } + """ + ) + response = self.gql( + mutation, + variable_values={ + "completeMultipartAction": complete_multipart_action, + "artifactID": artifact_id, + "storagePath": storage_path, + "completedParts": completed_parts, + "uploadID": upload_id, + }, + ) + digest: Optional[str] = response["completeMultipartUploadArtifact"]["digest"] + return digest + + def create_artifact_manifest( + self, + name: str, + digest: str, + artifact_id: Optional[str], + base_artifact_id: Optional[str] = None, + entity: Optional[str] = None, + project: Optional[str] = None, + run: Optional[str] = None, + include_upload: bool = True, + type: str = "FULL", + ) -> Tuple[str, Dict[str, Any]]: + mutation = gql( + """ + mutation CreateArtifactManifest( + $name: String!, + $digest: String!, + $artifactID: ID!, + $baseArtifactID: ID, + $entityName: String!, + $projectName: String!, + $runName: String!, + $includeUpload: Boolean!, + {} + ) {{ + createArtifactManifest(input: {{ + name: $name, + digest: $digest, + artifactID: $artifactID, + baseArtifactID: $baseArtifactID, + entityName: $entityName, + projectName: $projectName, + runName: $runName, + {} + }}) {{ + artifactManifest {{ + id + file {{ + id + name + displayName + uploadUrl @include(if: $includeUpload) + uploadHeaders @include(if: $includeUpload) + }} + }} + }} + }} + """.format( + "$type: ArtifactManifestType = FULL" if type != "FULL" else "", + "type: $type" if type != "FULL" else "", + ) + ) + + entity_name = entity or self.settings("entity") + project_name = project or self.settings("project") + run_name = run or self.current_run_id + + response = self.gql( + mutation, + variable_values={ + "name": name, + "digest": digest, + "artifactID": artifact_id, + "baseArtifactID": base_artifact_id, + "entityName": entity_name, + "projectName": project_name, + "runName": run_name, + "includeUpload": include_upload, + "type": type, + }, + ) + return ( + response["createArtifactManifest"]["artifactManifest"]["id"], + response["createArtifactManifest"]["artifactManifest"]["file"], + ) + + def update_artifact_manifest( + self, + artifact_manifest_id: str, + base_artifact_id: Optional[str] = None, + digest: Optional[str] = None, + include_upload: Optional[bool] = True, + ) -> Tuple[str, Dict[str, Any]]: + mutation = gql( + """ + mutation UpdateArtifactManifest( + $artifactManifestID: ID!, + $digest: String, + $baseArtifactID: ID, + $includeUpload: Boolean!, + ) { + updateArtifactManifest(input: { + artifactManifestID: $artifactManifestID, + digest: $digest, + baseArtifactID: $baseArtifactID, + }) { + artifactManifest { + id + file { + id + name + displayName + uploadUrl @include(if: $includeUpload) + uploadHeaders @include(if: $includeUpload) + } + } + } + } + """ + ) + + response = self.gql( + mutation, + variable_values={ + "artifactManifestID": artifact_manifest_id, + "digest": digest, + "baseArtifactID": base_artifact_id, + "includeUpload": include_upload, + }, + ) + + return ( + response["updateArtifactManifest"]["artifactManifest"]["id"], + response["updateArtifactManifest"]["artifactManifest"]["file"], + ) + + def update_artifact_metadata( + self, artifact_id: str, metadata: Dict[str, Any] + ) -> Dict[str, Any]: + """Set the metadata of the given artifact version.""" + mutation = gql( + """ + mutation UpdateArtifact( + $artifactID: ID!, + $metadata: JSONString, + ) { + updateArtifact(input: { + artifactID: $artifactID, + metadata: $metadata, + }) { + artifact { + id + } + } + } + """ + ) + response = self.gql( + mutation, + variable_values={ + "artifactID": artifact_id, + "metadata": json.dumps(metadata), + }, + ) + return response["updateArtifact"]["artifact"] + + def _resolve_client_id( + self, + client_id: str, + ) -> Optional[str]: + if client_id in self._client_id_mapping: + return self._client_id_mapping[client_id] + + query = gql( + """ + query ClientIDMapping($clientID: ID!) { + clientIDMapping(clientID: $clientID) { + serverID + } + } + """ + ) + response = self.gql( + query, + variable_values={ + "clientID": client_id, + }, + ) + server_id = None + if response is not None: + client_id_mapping = response.get("clientIDMapping") + if client_id_mapping is not None: + server_id = client_id_mapping.get("serverID") + if server_id is not None: + self._client_id_mapping[client_id] = server_id + return server_id + + def server_create_artifact_file_spec_input_introspection(self) -> List: + query_string = """ + query ProbeServerCreateArtifactFileSpecInput { + CreateArtifactFileSpecInputInfoType: __type(name:"CreateArtifactFileSpecInput") { + inputFields{ + name + } + } + } + """ + + query = gql(query_string) + res = self.gql(query) + create_artifact_file_spec_input_info = [ + field.get("name", "") + for field in res.get("CreateArtifactFileSpecInputInfoType", {}).get( + "inputFields", [{}] + ) + ] + return create_artifact_file_spec_input_info + + @normalize_exceptions + def create_artifact_files( + self, artifact_files: Iterable["CreateArtifactFileSpecInput"] + ) -> Mapping[str, "CreateArtifactFilesResponseFile"]: + query_template = """ + mutation CreateArtifactFiles( + $storageLayout: ArtifactStorageLayout! + $artifactFiles: [CreateArtifactFileSpecInput!]! + ) { + createArtifactFiles(input: { + artifactFiles: $artifactFiles, + storageLayout: $storageLayout, + }) { + files { + edges { + node { + id + name + displayName + uploadUrl + uploadHeaders + _MULTIPART_UPLOAD_FIELDS_ + artifact { + id + } + } + } + } + } + } + """ + multipart_upload_url_query = """ + storagePath + uploadMultipartUrls { + uploadID + uploadUrlParts { + partNumber + uploadUrl + } + } + """ + + # TODO: we should use constants here from interface/artifacts.py + # but probably don't want the dependency. We're going to remove + # this setting in a future release, so I'm just hard-coding the strings. + storage_layout = "V2" + if env.get_use_v1_artifacts(): + storage_layout = "V1" + + create_artifact_file_spec_input_fields = ( + self.server_create_artifact_file_spec_input_introspection() + ) + if "uploadPartsInput" in create_artifact_file_spec_input_fields: + query_template = query_template.replace( + "_MULTIPART_UPLOAD_FIELDS_", multipart_upload_url_query + ) + else: + query_template = query_template.replace("_MULTIPART_UPLOAD_FIELDS_", "") + + mutation = gql(query_template) + response = self.gql( + mutation, + variable_values={ + "storageLayout": storage_layout, + "artifactFiles": [af for af in artifact_files], + }, + ) + + result = {} + for edge in response["createArtifactFiles"]["files"]["edges"]: + node = edge["node"] + result[node["displayName"]] = node + return result + + @normalize_exceptions + def notify_scriptable_run_alert( + self, + title: str, + text: str, + level: Optional[str] = None, + wait_duration: Optional["Number"] = None, + ) -> bool: + mutation = gql( + """ + mutation NotifyScriptableRunAlert( + $entityName: String!, + $projectName: String!, + $runName: String!, + $title: String!, + $text: String!, + $severity: AlertSeverity = INFO, + $waitDuration: Duration + ) { + notifyScriptableRunAlert(input: { + entityName: $entityName, + projectName: $projectName, + runName: $runName, + title: $title, + text: $text, + severity: $severity, + waitDuration: $waitDuration + }) { + success + } + } + """ + ) + + response = self.gql( + mutation, + variable_values={ + "entityName": self.settings("entity"), + "projectName": self.settings("project"), + "runName": self.current_run_id, + "title": title, + "text": text, + "severity": level, + "waitDuration": wait_duration, + }, + ) + success: bool = response["notifyScriptableRunAlert"]["success"] + return success + + def get_sweep_state( + self, sweep: str, entity: Optional[str] = None, project: Optional[str] = None + ) -> "SweepState": + state: SweepState = self.sweep( + sweep=sweep, entity=entity, project=project, specs="{}" + )["state"] + return state + + def set_sweep_state( + self, + sweep: str, + state: "SweepState", + entity: Optional[str] = None, + project: Optional[str] = None, + ) -> None: + assert state in ("RUNNING", "PAUSED", "CANCELED", "FINISHED") + s = self.sweep(sweep=sweep, entity=entity, project=project, specs="{}") + curr_state = s["state"].upper() + if state == "PAUSED" and curr_state not in ("PAUSED", "RUNNING"): + raise Exception(f"Cannot pause {curr_state.lower()} sweep.") + elif state != "RUNNING" and curr_state not in ("RUNNING", "PAUSED", "PENDING"): + raise Exception(f"Sweep already {curr_state.lower()}.") + sweep_id = s["id"] + mutation = gql( + """ + mutation UpsertSweep( + $id: ID, + $state: String, + $entityName: String, + $projectName: String + ) { + upsertSweep(input: { + id: $id, + state: $state, + entityName: $entityName, + projectName: $projectName + }){ + sweep { + name + } + } + } + """ + ) + self.gql( + mutation, + variable_values={ + "id": sweep_id, + "state": state, + "entityName": entity or self.settings("entity"), + "projectName": project or self.settings("project"), + }, + ) + + def stop_sweep( + self, + sweep: str, + entity: Optional[str] = None, + project: Optional[str] = None, + ) -> None: + """Finish the sweep to stop running new runs and let currently running runs finish.""" + self.set_sweep_state( + sweep=sweep, state="FINISHED", entity=entity, project=project + ) + + def cancel_sweep( + self, + sweep: str, + entity: Optional[str] = None, + project: Optional[str] = None, + ) -> None: + """Cancel the sweep to kill all running runs and stop running new runs.""" + self.set_sweep_state( + sweep=sweep, state="CANCELED", entity=entity, project=project + ) + + def pause_sweep( + self, + sweep: str, + entity: Optional[str] = None, + project: Optional[str] = None, + ) -> None: + """Pause the sweep to temporarily stop running new runs.""" + self.set_sweep_state( + sweep=sweep, state="PAUSED", entity=entity, project=project + ) + + def resume_sweep( + self, + sweep: str, + entity: Optional[str] = None, + project: Optional[str] = None, + ) -> None: + """Resume the sweep to continue running new runs.""" + self.set_sweep_state( + sweep=sweep, state="RUNNING", entity=entity, project=project + ) + + def _status_request(self, url: str, length: int) -> requests.Response: + """Ask google how much we've uploaded.""" + check_httpclient_logger_handler() + return requests.put( + url=url, + headers={"Content-Length": "0", "Content-Range": f"bytes */{length}"}, + ) + + def _flatten_edges(self, response: "_Response") -> List[Dict]: + """Return an array from the nested graphql relay structure.""" + return [node["node"] for node in response["edges"]] + + @normalize_exceptions + def stop_run( + self, + run_id: str, + ) -> bool: + mutation = gql( + """ + mutation stopRun($id: ID!) { + stopRun(input: { + id: $id + }) { + clientMutationId + success + } + } + """ + ) + + response = self.gql( + mutation, + variable_values={ + "id": run_id, + }, + ) + + success: bool = response["stopRun"].get("success") + + return success + + @normalize_exceptions + def create_custom_chart( + self, + entity: str, + name: str, + display_name: str, + spec_type: str, + access: str, + spec: Union[str, Mapping[str, Any]], + ) -> Optional[Dict[str, Any]]: + if not isinstance(spec, str): + spec = json.dumps(spec) + + mutation = gql( + """ + mutation CreateCustomChart( + $entity: String! + $name: String! + $displayName: String! + $type: String! + $access: String! + $spec: JSONString! + ) { + createCustomChart( + input: { + entity: $entity + name: $name + displayName: $displayName + type: $type + access: $access + spec: $spec + } + ) { + chart { id } + } + } + """ + ) + + variable_values = { + "entity": entity, + "name": name, + "displayName": display_name, + "type": spec_type, + "access": access, + "spec": spec, + } + + result: Optional[Dict[str, Any]] = self.gql(mutation, variable_values)[ + "createCustomChart" + ] + return result diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/job_builder.py b/lib/python3.12/site-packages/wandb/sdk/internal/job_builder.py new file mode 100644 index 0000000000000000000000000000000000000000..2bf231961dff9199f5c52efbc54644ba2a8ffc55 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/internal/job_builder.py @@ -0,0 +1,633 @@ +"""job builder.""" + +import json +import logging +import os +import re +import sys +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + List, + Literal, + Optional, + Tuple, + TypedDict, + Union, +) + +import wandb +from wandb.sdk.artifacts._internal_artifact import InternalArtifact +from wandb.sdk.artifacts.artifact import Artifact +from wandb.sdk.data_types._dtypes import TypeRegistry +from wandb.sdk.internal.internal_api import Api +from wandb.sdk.lib.filenames import DIFF_FNAME, METADATA_FNAME, REQUIREMENTS_FNAME +from wandb.util import make_artifact_name_safe + +from .settings_static import SettingsStatic + +_logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from wandb.proto.wandb_internal_pb2 import ArtifactRecord + +FROZEN_REQUIREMENTS_FNAME = "requirements.frozen.txt" +JOB_FNAME = "wandb-job.json" +JOB_ARTIFACT_TYPE = "job" + +LOG_LEVEL = Literal["log", "warn", "error"] + + +class Version: + def __init__(self, major: int, minor: int, patch: int): + self._major = major + self._minor = minor + self._patch = patch + + def __repr__(self) -> str: + return f"{self._major}.{self._minor}.{self._patch}" + + def __lt__(self, other: "Version") -> bool: + if self._major < other._major: + return True + elif self._major == other._major: + if self._minor < other._minor: + return True + elif self._minor == other._minor: + if self._patch < other._patch: + return True + return False + + def __eq__(self, other: object) -> bool: + if not isinstance(other, Version): + return NotImplemented + return ( + self._major == other._major + and self._minor == other._minor + and self._patch == other._patch + ) + + +# Minimum supported wandb version for keys in the source dict of wandb-job.json +SOURCE_KEYS_MIN_SUPPORTED_VERSION = { + "dockerfile": Version(0, 17, 0), + "build_context": Version(0, 17, 0), +} + + +class GitInfo(TypedDict): + remote: str + commit: str + + +class GitSourceDict(TypedDict): + git: GitInfo + entrypoint: List[str] + notebook: bool + build_context: Optional[str] + dockerfile: Optional[str] + + +class ArtifactSourceDict(TypedDict): + artifact: str + entrypoint: List[str] + notebook: bool + build_context: Optional[str] + dockerfile: Optional[str] + + +class ImageSourceDict(TypedDict): + image: str + + +class JobSourceDict(TypedDict, total=False): + _version: str + source_type: str + source: Union[GitSourceDict, ArtifactSourceDict, ImageSourceDict] + input_types: Dict[str, Any] + output_types: Dict[str, Any] + runtime: Optional[str] + + +class ArtifactInfoForJob(TypedDict): + id: str + name: str + + +def get_min_supported_for_source_dict( + source: Union[GitSourceDict, ArtifactSourceDict, ImageSourceDict], +) -> Optional[Version]: + """Get the minimum supported wandb version the source dict of wandb-job.json.""" + min_seen = None + for key in source: + new_ver = SOURCE_KEYS_MIN_SUPPORTED_VERSION.get(key) + if new_ver: + if min_seen is None or new_ver < min_seen: + min_seen = new_ver + return min_seen + + +class JobBuilder: + _settings: SettingsStatic + _metadatafile_path: Optional[str] + _requirements_path: Optional[str] + _config: Optional[Dict[str, Any]] + _summary: Optional[Dict[str, Any]] + _logged_code_artifact: Optional[ArtifactInfoForJob] + _disable: bool + _partial_source_id: Optional[str] # Partial job source artifact id. + _aliases: List[str] + _job_seq_id: Optional[str] + _job_version_alias: Optional[str] + _is_notebook_run: bool + _verbose: bool + + def __init__(self, settings: SettingsStatic, verbose: bool = False): + self._settings = settings + self._metadatafile_path = None + self._requirements_path = None + self._config = None + self._summary = None + self._logged_code_artifact = None + self._job_seq_id = None + self._job_version_alias = None + self._disable = settings.disable_job_creation or settings.x_disable_machine_info + self._partial_source_id = None + self._aliases = [] + self._source_type: Optional[Literal["repo", "artifact", "image"]] = ( + settings.job_source # type: ignore[assignment] + ) + self._is_notebook_run = self._get_is_notebook_run() + self._verbose = verbose + self._partial = False + + def set_config(self, config: Dict[str, Any]) -> None: + self._config = config + + def set_summary(self, summary: Dict[str, Any]) -> None: + self._summary = summary + + @property + def disable(self) -> bool: + return self._disable + + @disable.setter + def disable(self, val: bool) -> None: + self._disable = val + + @property + def input_types(self) -> Dict[str, Any]: + return TypeRegistry.type_of(self._config).to_json() + + @property + def output_types(self) -> Dict[str, Any]: + return TypeRegistry.type_of(self._summary).to_json() + + def set_partial_source_id(self, source_id: str) -> None: + self._partial_source_id = source_id + + def _handle_server_artifact( + self, res: Optional[Dict], artifact: "ArtifactRecord" + ) -> None: + if artifact.type == "job" and res is not None: + try: + if res["artifactSequence"]["latestArtifact"] is None: + self._job_version_alias = "v0" + elif res["artifactSequence"]["latestArtifact"]["id"] == res["id"]: + self._job_version_alias = ( + f"v{res['artifactSequence']['latestArtifact']['versionIndex']}" + ) + else: + self._job_version_alias = f"v{res['artifactSequence']['latestArtifact']['versionIndex'] + 1}" + self._job_seq_id = res["artifactSequence"]["id"] + except KeyError as e: + _logger.info(f"Malformed response from ArtifactSaver.save {e}") + if artifact.type == "code" and res is not None: + self._logged_code_artifact = ArtifactInfoForJob( + { + "id": res["id"], + "name": artifact.name, + } + ) + + def _build_repo_job_source( + self, + program_relpath: str, + metadata: Dict[str, Any], + ) -> Tuple[Optional[GitSourceDict], Optional[str]]: + git_info: Dict[str, str] = metadata.get("git", {}) + remote = git_info.get("remote") + commit = git_info.get("commit") + root = metadata.get("root") + assert remote is not None + assert commit is not None + if self._is_notebook_run: + if not os.path.exists( + os.path.join(os.getcwd(), os.path.basename(program_relpath)) + ): + return None, None + + if root is None or self._settings.x_jupyter_root is None: + _logger.info("target path does not exist, exiting") + return None, None + assert self._settings.x_jupyter_root is not None + # git notebooks set the root to the git root, + # jupyter_root contains the path where the jupyter notebook was started + # program_relpath contains the path from jupyter_root to the file + # full program path here is actually the relpath from the program to the git root + full_program_path = os.path.join( + os.path.relpath(str(self._settings.x_jupyter_root), root), + program_relpath, + ) + full_program_path = os.path.normpath(full_program_path) + # if the notebook server is started above the git repo need to clear all the ..s + if full_program_path.startswith(".."): + split_path = full_program_path.split("/") + count_dots = 0 + for p in split_path: + if p == "..": + count_dots += 1 + full_program_path = "/".join(split_path[2 * count_dots :]) + else: + full_program_path = program_relpath + + entrypoint = self._get_entrypoint(full_program_path, metadata) + # TODO: update executable to a method that supports pex + source: GitSourceDict = { + "git": {"remote": remote, "commit": commit}, + "entrypoint": entrypoint, + "notebook": self._is_notebook_run, + "build_context": metadata.get("build_context"), + "dockerfile": metadata.get("dockerfile"), + } + name = self._make_job_name(f"{remote}_{program_relpath}") + + return source, name + + def _log_if_verbose(self, message: str, level: LOG_LEVEL) -> None: + log_func: Optional[Union[Callable[[Any], None], Callable[[Any], None]]] = None + if level == "log": + _logger.info(message) + log_func = wandb.termlog + elif level == "warn": + _logger.warning(message) + log_func = wandb.termwarn + elif level == "error": + _logger.error(message) + log_func = wandb.termerror + + if self._verbose and log_func is not None: + log_func(message) + + def _build_artifact_job_source( + self, + program_relpath: str, + metadata: Dict[str, Any], + ) -> Tuple[Optional[ArtifactSourceDict], Optional[str]]: + assert isinstance(self._logged_code_artifact, dict) + # TODO: should we just always exit early if the path doesn't exist? + if self._is_notebook_run and not self._is_colab_run(): + full_program_relpath = os.path.relpath(program_relpath, os.getcwd()) + # if the resolved path doesn't exist, then we shouldn't make a job because it will fail + if not os.path.exists(full_program_relpath): + # when users call log code in a notebook the code artifact starts + # at the directory the notebook is in instead of the jupyter core + if not os.path.exists(os.path.basename(program_relpath)): + _logger.info("target path does not exist, exiting") + self._log_if_verbose( + "No program path found when generating artifact job source for a non-colab notebook run. See https://docs.wandb.ai/guides/launch/create-job", + "warn", + ) + return None, None + full_program_relpath = os.path.basename(program_relpath) + else: + full_program_relpath = program_relpath + + entrypoint = self._get_entrypoint(full_program_relpath, metadata) + # TODO: update executable to a method that supports pex + source: ArtifactSourceDict = { + "entrypoint": entrypoint, + "notebook": self._is_notebook_run, + "artifact": f"wandb-artifact://_id/{self._logged_code_artifact['id']}", + "build_context": metadata.get("build_context"), + "dockerfile": metadata.get("dockerfile"), + } + artifact_basename, *_ = self._logged_code_artifact["name"].split(":") + name = self._make_job_name(artifact_basename) + + return source, name + + def _build_image_job_source( + self, metadata: Dict[str, Any] + ) -> Tuple[ImageSourceDict, str]: + image_name = metadata.get("docker") + assert isinstance(image_name, str) + + raw_image_name = image_name + if ":" in image_name: + tag = image_name.split(":")[-1] + + # if tag looks properly formatted, assume its a tag + # regex: alphanumeric and "_" "-" "." + if re.fullmatch(r"([a-zA-Z0-9_\-\.]+)", tag): + raw_image_name = raw_image_name.replace(f":{tag}", "") + self._aliases += [tag] + + source: ImageSourceDict = { + "image": image_name, + } + name = self._make_job_name(raw_image_name) + + return source, name + + def _make_job_name(self, input_str: str) -> str: + """Use job name from settings if provided, else use programmatic name.""" + if self._settings.job_name: + return self._settings.job_name + + return make_artifact_name_safe(f"job-{input_str}") + + def _get_entrypoint( + self, + program_relpath: str, + metadata: Dict[str, Any], + ) -> List[str]: + # if building a partial job from CLI, overwrite entrypoint and notebook + # should already be in metadata from create_job + if self._partial: + if metadata.get("entrypoint"): + entrypoint: List[str] = metadata["entrypoint"] + return entrypoint + # job is being built from a run + entrypoint = [os.path.basename(sys.executable), program_relpath] + + return entrypoint + + def _get_is_notebook_run(self) -> bool: + return hasattr(self._settings, "_jupyter") and bool(self._settings._jupyter) + + def _is_colab_run(self) -> bool: + return hasattr(self._settings, "_colab") and bool(self._settings._colab) + + def _build_job_source( + self, + source_type: str, + program_relpath: Optional[str], + metadata: Dict[str, Any], + ) -> Tuple[ + Union[GitSourceDict, ArtifactSourceDict, ImageSourceDict, None], + Optional[str], + ]: + """Construct a job source dict and name from the current run. + + Args: + source_type (str): The type of source to build the job from. One of + "repo", "artifact", or "image". + """ + source: Union[ + GitSourceDict, + ArtifactSourceDict, + ImageSourceDict, + None, + ] = None + + if source_type == "repo": + source, name = self._build_repo_job_source( + program_relpath or "", + metadata, + ) + elif source_type == "artifact": + source, name = self._build_artifact_job_source( + program_relpath or "", + metadata, + ) + elif source_type == "image" and self._has_image_job_ingredients(metadata): + source, name = self._build_image_job_source(metadata) + else: + source = None + + if source is None: + if source_type: + self._log_if_verbose( + f"Source type is set to '{source_type}' but some required information is missing " + "from the environment. A job will not be created from this run. See " + "https://docs.wandb.ai/guides/launch/create-job", + "warn", + ) + return None, None + + return source, name + + def build( + self, + api: Api, + build_context: Optional[str] = None, + dockerfile: Optional[str] = None, + base_image: Optional[str] = None, + ) -> Optional[Artifact]: + """Build a job artifact from the current run. + + Args: + api (Api): The API object to use to create the job artifact. + build_context (Optional[str]): Path within the job source code to + the image build context. Saved as part of the job for future + builds. + dockerfile (Optional[str]): Path within the build context the + Dockerfile. Saved as part of the job for future builds. + base_image (Optional[str]): The base image used to run the job code. + + Returns: + Optional[Artifact]: The job artifact if it was successfully built, + otherwise None. + """ + _logger.info("Attempting to build job artifact") + + # If a partial job was used, write the input/output types to the metadata + # rather than building a new job version. + if self._partial_source_id is not None: + new_metadata = { + "input_types": {"@wandb.config": self.input_types}, + "output_types": self.output_types, + } + api.update_artifact_metadata( + self._partial_source_id, + new_metadata, + ) + return None + + if not os.path.exists( + os.path.join(self._settings.files_dir, REQUIREMENTS_FNAME) + ): + self._log_if_verbose( + "No requirements.txt found, not creating job artifact. See https://docs.wandb.ai/guides/launch/create-job", + "warn", + ) + return None + metadata = self._handle_metadata_file() + if metadata is None: + self._log_if_verbose( + f"Ensure read and write access to run files dir: {self._settings.files_dir}, control this via the WANDB_DIR env var. See https://docs.wandb.ai/guides/track/environment-variables", + "warn", + ) + return None + + runtime: Optional[str] = metadata.get("python") + # can't build a job without a python version + if runtime is None: + self._log_if_verbose( + "No python version found in metadata, not creating job artifact. " + "See https://docs.wandb.ai/guides/launch/create-job", + "warn", + ) + return None + + input_types = TypeRegistry.type_of(self._config).to_json() + output_types = TypeRegistry.type_of(self._summary).to_json() + + name: Optional[str] = None + source_info: Optional[JobSourceDict] = None + + # configure job from environment + source_type = self._get_source_type(metadata) + if not source_type: + # if source_type is None, then we don't have enough information to build a job + # if the user intended to create a job, warn. + if ( + self._settings.job_name + or self._settings.job_source + or self._source_type + ): + self._log_if_verbose( + "No source type found, not creating job artifact", "warn" + ) + return None + + program_relpath = self._get_program_relpath(source_type, metadata) + if not self._partial and source_type != "image" and not program_relpath: + self._log_if_verbose( + "No program path found, not creating job artifact. " + "See https://docs.wandb.ai/guides/launch/create-job", + "warn", + ) + return None + + source, name = self._build_job_source( + source_type, + program_relpath, + metadata, + ) + if source is None: + return None + + if build_context: + source["build_context"] = build_context # type: ignore[typeddict-item] + if dockerfile: + source["dockerfile"] = dockerfile # type: ignore[typeddict-item] + if base_image: + source["base_image"] = base_image # type: ignore[typeddict-item] + + # Pop any keys that are initialized to None. The current TypedDict + # system for source dicts requires all keys to be present, but we + # don't want to include keys that are None in the final dict. + for key in list(source.keys()): + if source[key] is None: # type: ignore[literal-required] + source.pop(key) # type: ignore[literal-require,misc] + + source_info = { + "_version": str(get_min_supported_for_source_dict(source) or "v0"), + "source_type": source_type, + "source": source, + "input_types": input_types, + "output_types": output_types, + "runtime": runtime, + } + + assert source_info is not None + assert name is not None + + artifact = InternalArtifact(name, JOB_ARTIFACT_TYPE) + + _logger.info("adding wandb-job metadata file") + with artifact.new_file("wandb-job.json") as f: + f.write(json.dumps(source_info, indent=4)) + + artifact.add_file( + os.path.join(self._settings.files_dir, REQUIREMENTS_FNAME), + name=FROZEN_REQUIREMENTS_FNAME, + ) + + if source_type == "repo": + # add diff + if os.path.exists(os.path.join(self._settings.files_dir, DIFF_FNAME)): + artifact.add_file( + os.path.join(self._settings.files_dir, DIFF_FNAME), + name=DIFF_FNAME, + ) + + return artifact + + def _get_source_type(self, metadata: Dict[str, Any]) -> Optional[str]: + if self._source_type: + return self._source_type + + if self._has_git_job_ingredients(metadata): + _logger.info("is repo sourced job") + return "repo" + + if self._has_artifact_job_ingredients(): + _logger.info("is artifact sourced job") + return "artifact" + + if self._has_image_job_ingredients(metadata): + _logger.info("is image sourced job") + return "image" + + _logger.info("no source found") + return None + + def _get_program_relpath( + self, source_type: str, metadata: Dict[str, Any] + ) -> Optional[str]: + if self._is_notebook_run: + _logger.info("run is notebook based run") + program = metadata.get("program") + + if not program: + self._log_if_verbose( + "Notebook 'program' path not found in metadata. See https://docs.wandb.ai/guides/launch/create-job", + "warn", + ) + + return program + + if source_type == "artifact" or self._settings.job_source == "artifact": + # if the job is set to be an artifact, use relpath guaranteed + # to be correct. 'codePath' uses the root path when in git repo + # fallback to codePath if strictly local relpath not present + return metadata.get("codePathLocal") or metadata.get("codePath") + + return metadata.get("codePath") + + def _handle_metadata_file( + self, + ) -> Optional[Dict]: + if os.path.exists(os.path.join(self._settings.files_dir, METADATA_FNAME)): + with open(os.path.join(self._settings.files_dir, METADATA_FNAME)) as f: + metadata: Dict = json.load(f) + return metadata + + return None + + def _has_git_job_ingredients(self, metadata: Dict[str, Any]) -> bool: + git_info: Dict[str, str] = metadata.get("git", {}) + if self._is_notebook_run and metadata.get("root") is None: + return False + return git_info.get("remote") is not None and git_info.get("commit") is not None + + def _has_artifact_job_ingredients(self) -> bool: + return self._logged_code_artifact is not None + + def _has_image_job_ingredients(self, metadata: Dict[str, Any]) -> bool: + return metadata.get("docker") is not None diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/profiler.py b/lib/python3.12/site-packages/wandb/sdk/internal/profiler.py new file mode 100644 index 0000000000000000000000000000000000000000..9973e42984580a3fca6dc6706728a58ed22d4277 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/internal/profiler.py @@ -0,0 +1,79 @@ +"""Integration with pytorch profiler.""" + +import os + +import wandb +from wandb.errors import Error, UsageError +from wandb.sdk.lib import telemetry + +PYTORCH_MODULE = "torch" +PYTORCH_PROFILER_MODULE = "torch.profiler" + + +def torch_trace_handler(): + """Create a trace handler for traces generated by the profiler. + + Provide as an argument to `torch.profiler.profile`: + ```python + torch.profiler.profile(..., on_trace_ready=wandb.profiler.torch_trace_handler()) + ``` + + Calling this function ensures that profiler charts & tables can be viewed in + your run dashboard on wandb.ai. + + Please note that `wandb.init()` must be called before this function is + invoked, and the reinit setting must not be set to "create_new". The PyTorch + (torch) version must also be at least 1.9, in order to ensure stability of + their Profiler API. + + Args: + None + + Returns: + None + + Raises: + UsageError if wandb.init() hasn't been called before profiling. + Error if torch version is less than 1.9.0. + + Examples: + ```python + run = wandb.init() + run.config.id = "profile_code" + + with torch.profiler.profile( + schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1), + on_trace_ready=wandb.profiler.torch_trace_handler(), + record_shapes=True, + with_stack=True, + ) as prof: + for i, batch in enumerate(dataloader): + if step >= 5: + break + train(batch) + prof.step() + ``` + """ + from packaging.version import parse + + torch = wandb.util.get_module(PYTORCH_MODULE, required=True) + torch_profiler = wandb.util.get_module(PYTORCH_PROFILER_MODULE, required=True) + + if parse(torch.__version__) < parse("1.9.0"): + raise Error( + f"torch version must be at least 1.9 in order to use the PyTorch Profiler API.\ + \nVersion of torch currently installed: {torch.__version__}" + ) + + try: + logdir = os.path.join(wandb.run.dir, "pytorch_traces") # type: ignore + os.mkdir(logdir) + except AttributeError: + raise UsageError( + "Please call `wandb.init()` before `wandb.profiler.torch_trace_handler()`" + ) from None + + with telemetry.context() as tel: + tel.feature.torch_profiler_trace = True + + return torch_profiler.tensorboard_trace_handler(logdir) diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/progress.py b/lib/python3.12/site-packages/wandb/sdk/internal/progress.py new file mode 100644 index 0000000000000000000000000000000000000000..ef145aefe2510f86010ee5377bd0a7d9ecf715b2 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/internal/progress.py @@ -0,0 +1,77 @@ +"""progress.""" + +import os +from typing import IO, TYPE_CHECKING, Optional + +from wandb.errors import CommError + +if TYPE_CHECKING: + from typing import Protocol + + class ProgressFn(Protocol): + def __call__(self, new_bytes: int, total_bytes: int) -> None: + pass + + +class Progress: + """A helper class for displaying progress.""" + + ITER_BYTES = 1024 * 1024 + + def __init__( + self, file: IO[bytes], callback: Optional["ProgressFn"] = None + ) -> None: + self.file = file + if callback is None: + + def callback_(new_bytes: int, total_bytes: int) -> None: + pass + + callback = callback_ + + self.callback: ProgressFn = callback + self.bytes_read = 0 + self.len = os.fstat(file.fileno()).st_size + + def read(self, size=-1): + """Read bytes and call the callback.""" + bites = self.file.read(size) + self.bytes_read += len(bites) + if not bites and self.bytes_read < self.len: + # Files shrinking during uploads causes request timeouts. Maybe + # we could avoid those by updating the self.len in real-time, but + # files getting truncated while uploading seems like something + # that shouldn't really be happening anyway. + raise CommError( + f"File {self.file.name} size shrank from {self.len} to {self.bytes_read} while it was being uploaded." + ) + # Growing files are also likely to be bad, but our code didn't break + # on those in the past, so it's riskier to make that an error now. + self.callback(len(bites), self.bytes_read) + return bites + + def rewind(self) -> None: + self.callback(-self.bytes_read, 0) + self.bytes_read = 0 + self.file.seek(0) + + def __getattr__(self, name): + """Fallback to the file object for attrs not defined here.""" + if hasattr(self.file, name): + return getattr(self.file, name) + else: + raise AttributeError + + def __iter__(self): + return self + + def __next__(self): + bites = self.read(self.ITER_BYTES) + if len(bites) == 0: + raise StopIteration + return bites + + def __len__(self): + return self.len + + next = __next__ diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/run.py b/lib/python3.12/site-packages/wandb/sdk/internal/run.py new file mode 100644 index 0000000000000000000000000000000000000000..885729a8b0c27eb91f9b0ff006576186be491235 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/internal/run.py @@ -0,0 +1,27 @@ +# +"""InternalRun - Internal-only run object. + +Semi-stubbed run for internal process use. + +""" + +from typing_extensions import override + +from wandb.sdk import wandb_run + + +class InternalRun(wandb_run.Run): + def __init__(self, run_obj, settings, datatypes_cb): + super().__init__(settings=settings) + self._run_obj = run_obj + self._datatypes_cb = datatypes_cb + + @override + def _set_backend(self, backend): + # This type of run object can't have a backend + # or do any writes. + pass + + @override + def _publish_file(self, fname: str) -> None: + self._datatypes_cb(fname) diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/sample.py b/lib/python3.12/site-packages/wandb/sdk/internal/sample.py new file mode 100644 index 0000000000000000000000000000000000000000..7bffe7562e4b28614bad3e4c1cbf5be5423532fe --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/internal/sample.py @@ -0,0 +1,70 @@ +"""sample.""" + +import math + + +class UniformSampleAccumulator: + def __init__(self, min_samples=None): + self._samples = min_samples or 64 + # force power of 2 samples + self._samples = 2 ** int(math.ceil(math.log(self._samples, 2))) + # target oversample by factor of 2 + self._samples2 = self._samples * 2 + # max size of each buffer + self._max = self._samples2 // 2 + self._shift = 0 + self._mask = (1 << self._shift) - 1 + self._buckets = int(math.log(self._samples2, 2)) + self._buckets_bits = int(math.log(self._buckets, 2)) + self._buckets_mask = (1 << self._buckets_bits + 1) - 1 + self._buckets_index = 0 + self._bucket = [] + self._index = [0] * self._buckets + self._count = 0 + self._log2 = [0] + + # pre-allocate buckets + for _ in range(self._buckets): + self._bucket.append([0] * self._max) + # compute integer log2 + self._log2 += [int(math.log(i, 2)) for i in range(1, 2**self._buckets + 1)] + + def _show(self): + print("=" * 20) # noqa: T201 + for b in range(self._buckets): + b = (b + self._buckets_index) % self._buckets + vals = [self._bucket[b][i] for i in range(self._index[b])] + print(f"{b}: {vals}") # noqa: T201 + + def add(self, val): + self._count += 1 + cnt = self._count + if cnt & self._mask: + return + b = cnt >> self._shift + b = self._log2[b] # b = int(math.log(b, 2)) + if b >= self._buckets: + self._index[self._buckets_index] = 0 + self._buckets_index = (self._buckets_index + 1) % self._buckets + self._shift += 1 + self._mask = (self._mask << 1) | 1 + b += self._buckets - 1 + b = (b + self._buckets_index) % self._buckets + self._bucket[b][self._index[b]] = val + self._index[b] += 1 + + def get(self): + full = [] + sampled = [] + # self._show() + for b in range(self._buckets): + max_num = 2**b + b = (b + self._buckets_index) % self._buckets + modb = self._index[b] // max_num + for i in range(self._index[b]): + if not modb or i % modb == 0: + sampled.append(self._bucket[b][i]) + full.append(self._bucket[b][i]) + if len(sampled) < self._samples: + return tuple(full) + return tuple(sampled) diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/sender.py b/lib/python3.12/site-packages/wandb/sdk/internal/sender.py new file mode 100644 index 0000000000000000000000000000000000000000..e87606d39e814175acd494e2014ed78c245a0a1b --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/internal/sender.py @@ -0,0 +1,1692 @@ +"""sender.""" + +import contextlib +import glob +import gzip +import json +import logging +import os +import queue +import threading +import time +import traceback +from collections import defaultdict +from datetime import datetime +from queue import Queue +from typing import ( + TYPE_CHECKING, + Any, + Dict, + Generator, + List, + Literal, + Optional, + Tuple, + Type, + Union, +) + +import requests + +import wandb +from wandb import util +from wandb.errors import CommError, UsageError +from wandb.errors.util import ProtobufErrorHandler +from wandb.filesync.dir_watcher import DirWatcher +from wandb.proto import wandb_internal_pb2 +from wandb.sdk.artifacts.artifact_saver import ArtifactSaver +from wandb.sdk.interface import interface +from wandb.sdk.interface.interface_queue import InterfaceQueue +from wandb.sdk.internal import ( + context, + datastore, + file_stream, + internal_api, + sender_config, +) +from wandb.sdk.internal.file_pusher import FilePusher +from wandb.sdk.internal.job_builder import JobBuilder +from wandb.sdk.internal.settings_static import SettingsStatic +from wandb.sdk.lib import ( + config_util, + filenames, + filesystem, + proto_util, + redirect, + retry, + telemetry, +) +from wandb.sdk.lib.proto_util import message_to_dict + +if TYPE_CHECKING: + from wandb.proto.wandb_internal_pb2 import ( + ArtifactManifest, + ArtifactManifestEntry, + ArtifactRecord, + EnvironmentRecord, + HttpResponse, + LocalInfo, + Record, + Result, + RunExitResult, + RunRecord, + SummaryRecord, + ) + + StreamLiterals = Literal["stdout", "stderr"] + + +logger = logging.getLogger(__name__) + + +_OUTPUT_MIN_CALLBACK_INTERVAL = 2 # seconds + + +def _framework_priority() -> Generator[Tuple[str, str], None, None]: + yield from [ + ("lightgbm", "lightgbm"), + ("catboost", "catboost"), + ("xgboost", "xgboost"), + ("transformers_huggingface", "huggingface"), # backwards compatibility + ("transformers", "huggingface"), + ("pytorch_ignite", "ignite"), # backwards compatibility + ("ignite", "ignite"), + ("pytorch_lightning", "lightning"), + ("fastai", "fastai"), + ("torch", "torch"), + ("keras", "keras"), + ("tensorflow", "tensorflow"), + ("sklearn", "sklearn"), + ] + + +def _manifest_json_from_proto(manifest: "ArtifactManifest") -> Dict: + if manifest.version == 1: + if manifest.manifest_file_path: + contents = {} + with gzip.open(manifest.manifest_file_path, "rt") as f: + for line in f: + entry_json = json.loads(line) + path = entry_json.pop("path") + contents[path] = entry_json + else: + contents = { + content.path: _manifest_entry_from_proto(content) + for content in manifest.contents + } + else: + raise ValueError(f"unknown artifact manifest version: {manifest.version}") + + return { + "version": manifest.version, + "storagePolicy": manifest.storage_policy, + "storagePolicyConfig": { + config.key: json.loads(config.value_json) + for config in manifest.storage_policy_config + }, + "contents": contents, + } + + +def _manifest_entry_from_proto(entry: "ArtifactManifestEntry") -> Dict: + birth_artifact_id = entry.birth_artifact_id if entry.birth_artifact_id else None + return { + "digest": entry.digest, + "birthArtifactID": birth_artifact_id, + "ref": entry.ref if entry.ref else None, + "size": entry.size if entry.size is not None else None, + "local_path": entry.local_path if entry.local_path else None, + "skip_cache": entry.skip_cache, + "extra": {extra.key: json.loads(extra.value_json) for extra in entry.extra}, + } + + +class ResumeState: + resumed: bool + step: int + history: int + events: int + output: int + runtime: float + wandb_runtime: Optional[int] + summary: Optional[Dict[str, Any]] + config: Optional[Dict[str, Any]] + tags: Optional[List[str]] + + def __init__(self) -> None: + self.resumed = False + self.step = 0 + self.history = 0 + self.events = 0 + self.output = 0 + self.runtime = 0 + # wandb_runtime is the canonical runtime (stored in summary._wandb.runtime) + self.wandb_runtime = None + self.summary = None + self.config = None + self.tags = None + + def __str__(self) -> str: + obj = ",".join(map(lambda it: f"{it[0]}={it[1]}", vars(self).items())) + return f"ResumeState({obj})" + + +class _OutputRawStream: + _stopped: threading.Event + _queue: queue.Queue + _emulator: redirect.TerminalEmulator + _writer_thr: threading.Thread + _reader_thr: threading.Thread + + def __init__(self, stream: str, sm: "SendManager"): + self._stopped = threading.Event() + self._queue = queue.Queue() + self._emulator = redirect.TerminalEmulator() + self._writer_thr = threading.Thread( + target=sm._output_raw_writer_thread, + kwargs=dict(stream=stream), + daemon=True, + name=f"OutRawWr-{stream}", + ) + self._reader_thr = threading.Thread( + target=sm._output_raw_reader_thread, + kwargs=dict(stream=stream), + daemon=True, + name=f"OutRawRd-{stream}", + ) + + def start(self) -> None: + self._writer_thr.start() + self._reader_thr.start() + + +class SendManager: + UPDATE_CONFIG_TIME: int = 30 + UPDATE_STATUS_TIME: int = 5 + + _settings: SettingsStatic + _record_q: "Queue[Record]" + _result_q: "Queue[Result]" + _interface: InterfaceQueue + _api_settings: Dict[str, str] + _partial_output: Dict[str, str] + _context_keeper: context.ContextKeeper + + _telemetry_obj: telemetry.TelemetryRecord + _environment_obj: "EnvironmentRecord" + _fs: Optional["file_stream.FileStreamApi"] + _run: Optional["RunRecord"] + _entity: Optional[str] + _project: Optional[str] + _dir_watcher: Optional["DirWatcher"] + _pusher: Optional["FilePusher"] + _record_exit: Optional["Record"] + _exit_result: Optional["RunExitResult"] + _resume_state: ResumeState + _rewind_response: Optional[Dict[str, Any]] + _cached_server_info: Dict[str, Any] + _cached_viewer: Dict[str, Any] + _server_messages: List[Dict[str, Any]] + _ds: Optional[datastore.DataStore] + _output_raw_streams: Dict["StreamLiterals", _OutputRawStream] + _output_raw_file: Optional[filesystem.CRDedupedFile] + _send_record_num: int + _send_end_offset: int + _debounce_config_time: float + _debounce_status_time: float + + def __init__( + self, + settings: SettingsStatic, + record_q: "Queue[Record]", + result_q: "Queue[Result]", + interface: InterfaceQueue, + context_keeper: context.ContextKeeper, + ) -> None: + self._settings = settings + self._record_q = record_q + self._result_q = result_q + self._interface = interface + self._context_keeper = context_keeper + + self._ds = None + self._send_record_num = 0 + self._send_end_offset = 0 + + self._fs = None + self._pusher = None + self._dir_watcher = None + + # State updated by login + self._entity = None + self._flags = None + + # State updated by wandb.init + self._run = None + self._project = None + + # keep track of config from key/val updates + self._consolidated_config = sender_config.ConfigState() + + self._start_time: int = 0 + self._telemetry_obj = telemetry.TelemetryRecord() + self._environment_obj = wandb_internal_pb2.EnvironmentRecord() + self._config_metric_pbdict_list: List[Dict[int, Any]] = [] + self._metadata_summary: Dict[str, Any] = defaultdict() + self._cached_summary: Dict[str, Any] = dict() + self._config_metric_index_dict: Dict[str, int] = {} + self._config_metric_dict: Dict[str, wandb_internal_pb2.MetricRecord] = {} + self._consolidated_summary: Dict[str, Any] = dict() + + self._cached_server_info = dict() + self._cached_viewer = dict() + self._server_messages = [] + + # State updated by resuming + self._resume_state = ResumeState() + self._rewind_response = None + + # State added when run_exit is initiated and complete + self._record_exit = None + self._exit_result = None + + self._api = internal_api.Api( + default_settings=settings, retry_callback=self.retry_callback + ) + self._api_settings = dict() + + # queue filled by retry_callback + self._retry_q: Queue[HttpResponse] = queue.Queue() + + # do we need to debounce? + self._config_needs_debounce: bool = False + + # TODO(jhr): do something better, why do we need to send full lines? + self._partial_output = dict() + + self._exit_code = 0 + + # internal vars for handing raw console output + self._output_raw_streams = dict() + self._output_raw_file = None + + # job builder + self._job_builder = JobBuilder(settings) + + time_now = time.monotonic() + self._debounce_config_time = time_now + self._debounce_status_time = time_now + + @classmethod + def setup( + cls, + root_dir: str, + resume: Union[None, bool, str], + ) -> "SendManager": + """Set up a standalone SendManager. + + Exclusively used in `sync.py`. + """ + files_dir = os.path.join(root_dir, "files") + settings = wandb.Settings( + x_files_dir=files_dir, + root_dir=root_dir, + # _start_time=0, + resume=resume, + # ignore_globs=(), + x_sync=True, + disable_job_creation=False, + x_file_stream_timeout_seconds=0, + ) + record_q: Queue[Record] = queue.Queue() + result_q: Queue[Result] = queue.Queue() + publish_interface = InterfaceQueue(record_q=record_q) + context_keeper = context.ContextKeeper() + return SendManager( + settings=SettingsStatic(settings.to_proto()), + record_q=record_q, + result_q=result_q, + interface=publish_interface, + context_keeper=context_keeper, + ) + + def __len__(self) -> int: + return self._record_q.qsize() + + def __enter__(self) -> "SendManager": + return self + + def __exit__( + self, + exc_type: Optional[Type[BaseException]], + exc_value: Optional[BaseException], + exc_traceback: Optional[traceback.TracebackException], + ) -> Literal[False]: + while self: + data = next(self) + self.send(data) + self.finish() + return False + + def retry_callback(self, status: int, response_text: str) -> None: + response = wandb_internal_pb2.HttpResponse() + response.http_status_code = status + response.http_response_text = response_text + self._retry_q.put(response) + + def send(self, record: "Record") -> None: + self._update_record_num(record.num) + self._update_end_offset(record.control.end_offset) + + record_type = record.WhichOneof("record_type") + assert record_type + handler_str = "send_" + record_type + send_handler = getattr(self, handler_str, None) + # Don't log output to reduce log noise + if record_type not in {"output", "request", "output_raw"}: + logger.debug(f"send: {record_type}") + assert send_handler, f"unknown send handler: {handler_str}" + + context_id = context.context_id_from_record(record) + api_context = self._context_keeper.get(context_id) + try: + self._api.set_local_context(api_context) + send_handler(record) + except retry.RetryCancelledError: + logger.debug(f"Record cancelled: {record_type}") + self._context_keeper.release(context_id) + finally: + self._api.clear_local_context() + + def send_preempting(self, _: "Record") -> None: + if self._fs: + self._fs.enqueue_preempting() + + def send_request_sender_mark(self, _: "Record") -> None: + self._maybe_report_status(always=True) + + def send_request(self, record: "Record") -> None: + request_type = record.request.WhichOneof("request_type") + assert request_type + handler_str = "send_request_" + request_type + send_handler = getattr(self, handler_str, None) + if request_type != "network_status": + logger.debug(f"send_request: {request_type}") + assert send_handler, f"unknown handle: {handler_str}" + send_handler(record) + + def _respond_result(self, result: "Result") -> None: + context_id = context.context_id_from_result(result) + self._context_keeper.release(context_id) + self._result_q.put(result) + + def _flatten(self, dictionary: Dict) -> None: + if isinstance(dictionary, dict): + for k, v in list(dictionary.items()): + if isinstance(v, dict): + self._flatten(v) + dictionary.pop(k) + for k2, v2 in v.items(): + dictionary[k + "." + k2] = v2 + + def _update_record_num(self, record_num: int) -> None: + if not record_num: + return + # Currently how we handle offline mode and syncing is not + # compatible with this assertion due to how the exit record + # is (mis)handled: + # - using "always_send" in offline mode to trigger defer + # state machine + # - skipping the exit record in `wandb sync` mode so that + # it is always executed as the last record + if not self._settings._offline and not self._settings.x_sync: + assert record_num == self._send_record_num + 1 + self._send_record_num = record_num + + def _update_end_offset(self, end_offset: int) -> None: + if not end_offset: + return + self._send_end_offset = end_offset + + def send_request_sender_read(self, record: "Record") -> None: + if self._ds is None: + self._ds = datastore.DataStore() + self._ds.open_for_scan(self._settings.sync_file) + + # TODO(cancel_paused): implement cancel_set logic + # The idea is that there is an active request to cancel a + # message that is being read from the transaction log below + + start_offset = record.request.sender_read.start_offset + final_offset = record.request.sender_read.final_offset + self._ds.seek(start_offset) + + current_end_offset = 0 + while current_end_offset < final_offset: + data = self._ds.scan_data() + assert data + current_end_offset = self._ds.get_offset() + + send_record = wandb_internal_pb2.Record() + send_record.ParseFromString(data) + self._update_end_offset(current_end_offset) + self.send(send_record) + + # make sure we perform deferred operations + self.debounce() + + # make sure that we always update writer for every sended read request + self._maybe_report_status(always=True) + + def send_request_stop_status(self, record: "Record") -> None: + result = proto_util._result_from_record(record) + status_resp = result.response.stop_status_response + status_resp.run_should_stop = False + if self._entity and self._project and self._run and self._run.run_id: + try: + status_resp.run_should_stop = self._api.check_stop_requested( + self._project, self._entity, self._run.run_id + ) + except Exception as e: + logger.warning("Failed to check stop requested status: %s", e) + self._respond_result(result) + + def _maybe_update_config(self, always: bool = False) -> None: + time_now = time.monotonic() + if ( + not always + and time_now < self._debounce_config_time + self.UPDATE_CONFIG_TIME + ): + return + if self._config_needs_debounce: + self._debounce_config() + self._debounce_config_time = time_now + + def _maybe_report_status(self, always: bool = False) -> None: + time_now = time.monotonic() + if ( + not always + and time_now < self._debounce_status_time + self.UPDATE_STATUS_TIME + ): + return + self._debounce_status_time = time_now + + status_report = wandb_internal_pb2.StatusReportRequest( + record_num=self._send_record_num, + sent_offset=self._send_end_offset, + ) + status_time = time.time() + status_report.sync_time.FromMicroseconds(int(status_time * 1e6)) + record = self._interface._make_request(status_report=status_report) + self._interface._publish(record) + + def debounce(self, final: bool = False) -> None: + self._maybe_report_status(always=final) + self._maybe_update_config(always=final) + + def _debounce_config(self) -> None: + config_value_dict = self._config_backend_dict() + # TODO(jhr): check result of upsert_run? + if self._run: + self._api.upsert_run( + name=self._run.run_id, + config=config_value_dict, + **self._api_settings, # type: ignore + ) + self._config_save(config_value_dict) + self._config_needs_debounce = False + + def send_request_network_status(self, record: "Record") -> None: + result = proto_util._result_from_record(record) + status_resp = result.response.network_status_response + while True: + try: + status_resp.network_responses.append(self._retry_q.get_nowait()) + except queue.Empty: + break + except Exception as e: + logger.warning(f"Error emptying retry queue: {e}") + self._respond_result(result) + + def send_exit(self, record: "Record") -> None: + # track where the exit came from + self._record_exit = record + + run_exit = record.exit + self._exit_code = run_exit.exit_code + logger.info("handling exit code: %s", run_exit.exit_code) + runtime = run_exit.runtime + logger.info("handling runtime: %s", run_exit.runtime) + self._metadata_summary["runtime"] = runtime + self._update_summary() + + # We need to give the request queue a chance to empty between states + # so use handle_request_defer as a state machine. + logger.info("send defer") + self._interface.publish_defer() + + def send_final(self, record: "Record") -> None: + pass + + def _flush_run(self) -> None: + pass + + def send_request_status_report(self, record: "Record") -> None: + # todo? this is just a noop to please wandb sync + pass + + def send_request_defer(self, record: "Record") -> None: # noqa: C901 + defer = record.request.defer + state = defer.state + logger.info(f"handle sender defer: {state}") + + def transition_state() -> None: + state = defer.state + 1 + logger.info(f"send defer: {state}") + self._interface.publish_defer(state) + + done = False + if state == defer.BEGIN: + transition_state() + elif state == defer.FLUSH_RUN: + self._flush_run() + transition_state() + elif state == defer.FLUSH_STATS: + # NOTE: this is handled in handler.py:handle_request_defer() + transition_state() + elif state == defer.FLUSH_PARTIAL_HISTORY: + # NOTE: this is handled in handler.py:handle_request_defer() + transition_state() + elif state == defer.FLUSH_TB: + # NOTE: this is handled in handler.py:handle_request_defer() + transition_state() + elif state == defer.FLUSH_SUM: + # NOTE: this is handled in handler.py:handle_request_defer() + transition_state() + elif state == defer.FLUSH_DEBOUNCER: + self.debounce(final=True) + transition_state() + elif state == defer.FLUSH_OUTPUT: + self._output_raw_finish() + transition_state() + elif state == defer.FLUSH_JOB: + self._flush_job() + transition_state() + elif state == defer.FLUSH_DIR: + if self._dir_watcher: + self._dir_watcher.finish() + self._dir_watcher = None + transition_state() + elif state == defer.FLUSH_FP: + if self._pusher: + # FilePusher generates some events for FileStreamApi, so we + # need to wait for pusher to finish before going to the next + # state to ensure that filestream gets all the events that we + # want before telling it to finish up + self._pusher.finish(transition_state) + else: + transition_state() + elif state == defer.JOIN_FP: + if self._pusher: + self._pusher.join() + transition_state() + elif state == defer.FLUSH_FS: + if self._fs: + # TODO(jhr): now is a good time to output pending output lines + self._fs.finish(self._exit_code) + self._fs = None + transition_state() + elif state == defer.FLUSH_FINAL: + self._interface.publish_final() + self._interface.publish_footer() + transition_state() + elif state == defer.END: + done = True + else: + raise AssertionError("unknown state") + + if not done: + return + + exit_result = wandb_internal_pb2.RunExitResult() + + # mark exit done in case we are polling on exit + self._exit_result = exit_result + + # Report response to mailbox + if self._record_exit and self._record_exit.control.mailbox_slot: + result = proto_util._result_from_record(self._record_exit) + result.exit_result.CopyFrom(exit_result) + self._respond_result(result) + + def send_request_poll_exit(self, record: "Record") -> None: + if not record.control.req_resp and not record.control.mailbox_slot: + return + + result = proto_util._result_from_record(record) + + if self._pusher: + _alive, status = self._pusher.get_status() + file_counts = self._pusher.file_counts_by_category() + resp = result.response.poll_exit_response + resp.pusher_stats.uploaded_bytes = status.uploaded_bytes + resp.pusher_stats.total_bytes = status.total_bytes + resp.pusher_stats.deduped_bytes = status.deduped_bytes + resp.file_counts.wandb_count = file_counts.wandb + resp.file_counts.media_count = file_counts.media + resp.file_counts.artifact_count = file_counts.artifact + resp.file_counts.other_count = file_counts.other + + if self._exit_result: + result.response.poll_exit_response.done = True + result.response.poll_exit_response.exit_result.CopyFrom(self._exit_result) + + self._respond_result(result) + + def _setup_resume( + self, run: "RunRecord" + ) -> Optional["wandb_internal_pb2.ErrorInfo"]: + """Queries the backend for a run; fail if the settings are incompatible.""" + if not self._settings.resume: + return None + + # TODO: This causes a race, we need to make the upsert atomically + # only create or update depending on the resume config + # we use the runs entity if set, otherwise fallback to users entity + # todo: ensure entity is not None as self._entity is Optional[str] + entity = run.entity or self._entity + logger.info( + "checking resume status for %s/%s/%s", entity, run.project, run.run_id + ) + resume_status = self._api.run_resume_status( + entity=entity, # type: ignore + project_name=run.project, + name=run.run_id, + ) + # No resume status = run does not exist; No t key in wandbConfig = run exists but hasn't been inited + if not resume_status or '"t":' not in resume_status.get("wandbConfig", ""): + if self._settings.resume == "must": + error = wandb_internal_pb2.ErrorInfo() + error.code = wandb_internal_pb2.ErrorInfo.ErrorCode.USAGE + error.message = ( + "You provided an invalid value for the `resume` argument." + f" The value 'must' is not a valid option for resuming a run ({run.run_id}) that has not been initialized." + " Please check your inputs and try again with a valid run ID." + " If you are trying to start a new run, please omit the `resume` argument or use `resume='allow'`." + ) + return error + return None + + # + # handle cases where we have resume_status + # + if self._settings.resume == "never": + error = wandb_internal_pb2.ErrorInfo() + error.code = wandb_internal_pb2.ErrorInfo.ErrorCode.USAGE + error.message = ( + "You provided an invalid value for the `resume` argument." + f" The value 'never' is not a valid option for resuming a run ({run.run_id}) that already exists." + " Please check your inputs and try again with a valid value for the `resume` argument." + ) + return error + + history = {} + events = {} + config = {} + summary = {} + try: + events_rt = 0 + history_rt = 0 + history = json.loads(resume_status["historyTail"]) + if history: + history = json.loads(history[-1]) + history_rt = history.get("_runtime", 0) + events = json.loads(resume_status["eventsTail"]) + if events: + events = json.loads(events[-1]) + events_rt = events.get("_runtime", 0) + config = json.loads(resume_status["config"] or "{}") + summary = json.loads(resume_status["summaryMetrics"] or "{}") + new_runtime = summary.get("_wandb", {}).get("runtime", None) + if new_runtime is not None: + self._resume_state.wandb_runtime = new_runtime + tags = resume_status.get("tags") or [] + + except (IndexError, ValueError): + logger.exception("unable to load resume tails") + if self._settings.resume == "must": + error = wandb_internal_pb2.ErrorInfo() + error.code = wandb_internal_pb2.ErrorInfo.ErrorCode.USAGE + error.message = f"resume='must' but could not resume ({run.run_id}) " + return error + + # TODO: Do we need to restore config / summary? + # System metrics runtime is usually greater than history + self._resume_state.runtime = max(events_rt, history_rt) + last_step = history.get("_step", 0) + history_line_count = resume_status["historyLineCount"] + self._resume_state.step = last_step + 1 if history_line_count > 0 else last_step + self._resume_state.history = history_line_count + self._resume_state.events = resume_status["eventsLineCount"] + self._resume_state.output = resume_status["logLineCount"] + self._resume_state.config = config + self._resume_state.summary = summary + self._resume_state.tags = tags + self._resume_state.resumed = True + logger.info(f"configured resuming with: {self._resume_state}") + return None + + def _telemetry_get_framework(self) -> str: + """Get telemetry data for internal config structure.""" + # detect framework by checking what is loaded + imports: telemetry.TelemetryImports + if self._telemetry_obj.HasField("imports_finish"): + imports = self._telemetry_obj.imports_finish + elif self._telemetry_obj.HasField("imports_init"): + imports = self._telemetry_obj.imports_init + else: + return "" + framework = next( + (n for f, n in _framework_priority() if getattr(imports, f, False)), "" + ) + return framework + + def _config_backend_dict(self) -> sender_config.BackendConfigDict: + config = self._consolidated_config or sender_config.ConfigState() + return config.to_backend_dict( + telemetry_record=self._telemetry_obj, + framework=self._telemetry_get_framework(), + start_time_millis=self._start_time, + metric_pbdicts=self._config_metric_pbdict_list, + environment_record=self._environment_obj, + ) + + def _config_save( + self, + config_value_dict: sender_config.BackendConfigDict, + ) -> None: + config_path = os.path.join(self._settings.files_dir, "config.yaml") + config_util.save_config_file_from_dict(config_path, config_value_dict) + + def _sync_spell(self) -> None: + """Sync this run with spell.""" + if not self._run: + return + try: + env = os.environ + self._interface.publish_config( + key=("_wandb", "spell_url"), val=env.get("SPELL_RUN_URL") + ) + url = f"{self._api.app_url}/{self._run.entity}/{self._run.project}/runs/{self._run.run_id}" + requests.put( + env.get("SPELL_API_URL", "https://api.spell.run") + "/wandb_url", + json={"access_token": env.get("WANDB_ACCESS_TOKEN"), "url": url}, + timeout=2, + ) + except requests.RequestException: + pass + # TODO: do something if sync spell is not successful? + + def _setup_fork(self, server_run: dict): + assert self._run + assert self._run.branch_point + first_step = int(self._run.branch_point.value) + 1 + self._resume_state.step = first_step + self._resume_state.history = server_run.get("historyLineCount", 0) + self._run.forked = True + self._run.starting_step = first_step + + def _load_rewind_state(self, run: "RunRecord"): + assert run.branch_point + self._rewind_response = self._api.rewind_run( + run_name=run.run_id, + entity=run.entity or None, + project=run.project or None, + metric_name=run.branch_point.metric, + metric_value=run.branch_point.value, + program_path=self._settings.program or None, + ) + self._resume_state.history = self._rewind_response.get("historyLineCount", 0) + self._resume_state.config = json.loads( + self._rewind_response.get("config", "{}") + ) + + def _install_rewind_state(self): + assert self._run + assert self._run.branch_point + assert self._rewind_response + + first_step = int(self._run.branch_point.value) + 1 + self._resume_state.step = first_step + + # We set the fork flag here because rewind uses the forking + # infrastructure under the hood. Setting `forked` here + # ensures that run._step is properly set in the user process. + self._run.forked = True + self._run.starting_step = first_step + + def _handle_error( + self, + record: "Record", + error: "wandb_internal_pb2.ErrorInfo", + run: "RunRecord", + ) -> None: + if record.control.req_resp or record.control.mailbox_slot: + result = proto_util._result_from_record(record) + result.run_result.run.CopyFrom(run) + result.run_result.error.CopyFrom(error) + self._respond_result(result) + else: + logger.error("Got error in async mode: %s", error.message) + + def send_run(self, record: "Record", file_dir: Optional[str] = None) -> None: + run = record.run + error = None + is_wandb_init = self._run is None + + # save start time of a run + self._start_time = int(run.start_time.ToMicroseconds() // 1e6) + + # update telemetry + if run.telemetry: + self._telemetry_obj.MergeFrom(run.telemetry) + if self._settings.x_sync: + self._telemetry_obj.feature.sync = True + + # build config dict + config_value_dict: Optional[sender_config.BackendConfigDict] = None + if run.config: + self._consolidated_config.update_from_proto(run.config) + config_value_dict = self._config_backend_dict() + self._config_save(config_value_dict) + + do_rewind = run.branch_point.run == run.run_id + do_fork = not do_rewind and run.branch_point.run != "" + do_resume = bool(self._settings.resume) + + num_resume_options_set = sum([do_fork, do_rewind, do_resume]) + if num_resume_options_set > 1: + error = wandb_internal_pb2.ErrorInfo() + error.code = wandb_internal_pb2.ErrorInfo.ErrorCode.USAGE + error.message = ( + "Multiple resume options specified. " + "Please specify only one of `fork_from`, `resume`, or `resume_from`." + ) + self._handle_error(record, error, run) + + if is_wandb_init: + # Ensure we have a project to query for status + if run.project == "": + run.project = util.auto_project_name(self._settings.program) + # Only check resume status on `wandb.init` + + if do_resume: + error = self._setup_resume(run) + + elif do_rewind: + error = self._load_rewind_state(run) + + if error is not None: + self._handle_error(record, error, run) + return + + # Save the resumed config + if self._resume_state.config is not None: + self._consolidated_config.merge_resumed_config( + config_util.dict_strip_value_dict(self._resume_state.config) + ) + + config_value_dict = self._config_backend_dict() + self._config_save(config_value_dict) + + # handle empty config + # TODO(jhr): consolidate the 4 ways config is built: + # (passed config, empty config, resume config, send_config) + if not config_value_dict: + config_value_dict = self._config_backend_dict() + self._config_save(config_value_dict) + + try: + server_run = self._init_run(run, config_value_dict) + except (CommError, UsageError) as e: + logger.error(e, exc_info=True) + error = ProtobufErrorHandler.from_exception(e) + self._handle_error(record, error, run) + return + + assert self._run # self._run is configured in _init_run() + + if do_fork: + error = self._setup_fork(server_run) + + if error is not None: + self._handle_error(record, error, run) + return + + if record.control.req_resp or record.control.mailbox_slot: + result = proto_util._result_from_record(record) + # TODO: we could do self._interface.publish_defer(resp) to notify + # the handler not to actually perform server updates for this uuid + # because the user process will send a summary update when we resume + result.run_result.run.CopyFrom(self._run) + self._respond_result(result) + + # Only spin up our threads on the first run message + if is_wandb_init: + self._start_run_threads(file_dir) + else: + logger.info("updated run: %s", self._run.run_id) + + def _update_resume_state(self, is_rewinding: bool, inserted: bool): + assert self._run + if self._resume_state.resumed: + self._run.resumed = True + if self._resume_state.wandb_runtime is not None: + self._run.runtime = self._resume_state.wandb_runtime + elif is_rewinding: + # because is_rewinding is mutually exclusive with self._resume_state.resumed, + # this block will always execute if is_rewinding is set + self._install_rewind_state() + else: + # If the user is not resuming, and we didn't insert on upsert_run then + # it is likely that we are overwriting the run which we might want to + # prevent in the future. This could be a false signal since an upsert_run + # message which gets retried in the network could also show up as not + # inserted. + if not inserted: + # no need to flush this, it will get updated eventually + self._telemetry_obj.feature.maybe_run_overwrite = True + + def _init_run( + self, + run: "RunRecord", + config_dict: Optional[sender_config.BackendConfigDict], + ) -> dict: + # We subtract the previous runs runtime when resuming + start_time = ( + run.start_time.ToMicroseconds() / 1e6 + ) - self._resume_state.runtime + # TODO: we don't check inserted currently, ultimately we should make + # the upsert know the resume state and fail transactionally + + if self._resume_state and self._resume_state.tags and not run.tags: + run.tags.extend(self._resume_state.tags) + + is_rewinding = bool(self._settings.resume_from) + if is_rewinding: + assert self._rewind_response + server_run = self._rewind_response + server_messages = None + inserted = True + else: + server_run, inserted, server_messages = self._api.upsert_run( + name=run.run_id, + entity=run.entity or None, + project=run.project or None, + group=run.run_group or None, + job_type=run.job_type or None, + display_name=run.display_name or None, + notes=run.notes or None, + tags=run.tags[:] or None, + config=config_dict or None, + sweep_name=run.sweep_id or None, + host=run.host or None, + program_path=self._settings.program or None, + repo=run.git.remote_url or None, + commit=run.git.commit or None, + ) + + # TODO: we don't want to create jobs in sweeps, since the + # executable doesn't appear to be consistent + if run.sweep_id: + self._job_builder.disable = True + + self._server_messages = server_messages or [] + self._run = run + + if self._resume_state.resumed and is_rewinding: + # this should not ever be possible to hit, since we check for + # resumption above and raise an error if resumption is specified + # twice. + raise ValueError( + "Cannot attempt to rewind and resume a run - only one of " + "`resume` or `resume_from` can be specified." + ) + + self._update_resume_state(is_rewinding, inserted) + self._run.starting_step = self._resume_state.step + self._run.start_time.FromMicroseconds(int(start_time * 1e6)) + self._run.config.CopyFrom(self._interface._make_config(config_dict)) + if self._resume_state.summary is not None: + self._run.summary.CopyFrom( + self._interface._make_summary_from_dict(self._resume_state.summary) + ) + storage_id = server_run.get("id") + if storage_id: + self._run.storage_id = storage_id + id = server_run.get("name") + if id: + self._api.set_current_run_id(id) + display_name = server_run.get("displayName") + if display_name: + self._run.display_name = display_name + project = server_run.get("project") + # TODO: remove self._api.set_settings, and make self._project a property? + if project: + project_name = project.get("name") + if project_name: + self._run.project = project_name + self._project = project_name + self._api_settings["project"] = project_name + self._api.set_setting("project", project_name) + entity = project.get("entity") + if entity: + entity_name = entity.get("name") + if entity_name: + self._run.entity = entity_name + self._entity = entity_name + self._api_settings["entity"] = entity_name + self._api.set_setting("entity", entity_name) + sweep_id = server_run.get("sweepName") + if sweep_id: + self._run.sweep_id = sweep_id + if os.getenv("SPELL_RUN_URL"): + self._sync_spell() + return server_run + + def _start_run_threads(self, file_dir: Optional[str] = None) -> None: + assert self._run # self._run is configured by caller + self._fs = file_stream.FileStreamApi( + self._api, + self._run.run_id, + self._run.start_time.ToMicroseconds() / 1e6, + timeout=self._settings.x_file_stream_timeout_seconds or 0, + settings=self._api_settings, + ) + # Ensure the streaming polices have the proper offsets + self._fs.set_file_policy("wandb-summary.json", file_stream.SummaryFilePolicy()) + self._fs.set_file_policy( + "wandb-history.jsonl", + file_stream.JsonlFilePolicy(start_chunk_id=self._resume_state.history), + ) + self._fs.set_file_policy( + "wandb-events.jsonl", + file_stream.JsonlFilePolicy(start_chunk_id=self._resume_state.events), + ) + self._fs.set_file_policy( + "output.log", + file_stream.CRDedupeFilePolicy(start_chunk_id=self._resume_state.output), + ) + + # hack to merge run_settings and self._settings object together + # so that fields like entity or project are available to be attached to Sentry events. + run_settings = message_to_dict(self._run) + _settings = dict(self._settings) + _settings.update(run_settings) + wandb._sentry.configure_scope(tags=_settings, process_context="internal") + + self._fs.start() + self._pusher = FilePusher(self._api, self._fs, settings=self._settings) + self._dir_watcher = DirWatcher(self._settings, self._pusher, file_dir) + logger.info( + "run started: %s with start time %s", + self._run.run_id, + self._run.start_time.ToMicroseconds() / 1e6, + ) + + def _save_history(self, history_dict: Dict[str, Any]) -> None: + if self._fs: + self._fs.push(filenames.HISTORY_FNAME, json.dumps(history_dict)) + + def send_history(self, record: "Record") -> None: + history = record.history + history_dict = proto_util.dict_from_proto_list(history.item) + self._save_history(history_dict) + + def _update_summary_record(self, summary: "SummaryRecord") -> None: + summary_dict = proto_util.dict_from_proto_list(summary.update) + self._cached_summary = summary_dict + self._update_summary() + + def send_summary(self, record: "Record") -> None: + self._update_summary_record(record.summary) + + def send_request_summary_record(self, record: "Record") -> None: + self._update_summary_record(record.request.summary_record.summary) + + def _update_summary(self) -> None: + summary_dict = self._cached_summary.copy() + summary_dict.pop("_wandb", None) + if self._metadata_summary: + summary_dict["_wandb"] = self._metadata_summary + # merge with consolidated summary + self._consolidated_summary.update(summary_dict) + json_summary = json.dumps(self._consolidated_summary) + if self._fs: + self._fs.push(filenames.SUMMARY_FNAME, json_summary) + # TODO(jhr): we should only write this at the end of the script + summary_path = os.path.join(self._settings.files_dir, filenames.SUMMARY_FNAME) + with open(summary_path, "w") as f: + f.write(json_summary) + self._save_file(interface.GlobStr(filenames.SUMMARY_FNAME)) + + def send_stats(self, record: "Record") -> None: + stats = record.stats + if stats.stats_type != wandb_internal_pb2.StatsRecord.StatsType.SYSTEM: + return + if not self._fs: + return + if not self._run: + return + now_us = stats.timestamp.ToMicroseconds() + start_us = self._run.start_time.ToMicroseconds() + d = dict() + for item in stats.item: + try: + d[item.key] = json.loads(item.value_json) + except json.JSONDecodeError: + logger.exception("error decoding stats json: %s", item.value_json) + row: Dict[str, Any] = dict(system=d) + self._flatten(row) + row["_wandb"] = True + row["_timestamp"] = now_us / 1e6 + row["_runtime"] = (now_us - start_us) / 1e6 + self._fs.push(filenames.EVENTS_FNAME, json.dumps(row)) + # TODO(jhr): check fs.push results? + + def _output_raw_finish(self) -> None: + for stream, output_raw in self._output_raw_streams.items(): + output_raw._stopped.set() + + # shut down threads + output_raw._writer_thr.join(timeout=5) + if output_raw._writer_thr.is_alive(): + logger.info("processing output...") + output_raw._writer_thr.join() + output_raw._reader_thr.join() + + # flush output buffers and files + self._output_raw_flush(stream) + self._output_raw_streams = {} + if self._output_raw_file: + self._output_raw_file.close() + self._output_raw_file = None + + def _output_raw_writer_thread(self, stream: "StreamLiterals") -> None: + while True: + output_raw = self._output_raw_streams[stream] + if output_raw._queue.empty(): + if output_raw._stopped.is_set(): + return + time.sleep(0.5) + continue + data = [] + while not output_raw._queue.empty(): + data.append(output_raw._queue.get()) + if output_raw._stopped.is_set() and sum(map(len, data)) > 100000: + logger.warning("Terminal output too large. Logging without processing.") + self._output_raw_flush(stream) + for line in data: + self._output_raw_flush(stream, line) + # TODO: lets mark that this happened in telemetry + return + try: + output_raw._emulator.write("".join(data)) + except Exception as e: + logger.warning(f"problem writing to output_raw emulator: {e}") + + def _output_raw_reader_thread(self, stream: "StreamLiterals") -> None: + output_raw = self._output_raw_streams[stream] + while not (output_raw._stopped.is_set() and output_raw._queue.empty()): + self._output_raw_flush(stream) + time.sleep(_OUTPUT_MIN_CALLBACK_INTERVAL) + + def _output_raw_flush( + self, stream: "StreamLiterals", data: Optional[str] = None + ) -> None: + if data is None: + output_raw = self._output_raw_streams[stream] + try: + data = output_raw._emulator.read() + except Exception as e: + logger.warning(f"problem reading from output_raw emulator: {e}") + if data: + self._send_output_line(stream, data) + if self._output_raw_file: + self._output_raw_file.write(data.encode("utf-8")) + + def send_request_python_packages(self, record: "Record") -> None: + import os + + from wandb.sdk.lib.filenames import REQUIREMENTS_FNAME + + installed_packages_list = sorted( + f"{r.name}=={r.version}" for r in record.request.python_packages.package + ) + with open(os.path.join(self._settings.files_dir, REQUIREMENTS_FNAME), "w") as f: + f.write("\n".join(installed_packages_list)) + + def send_output(self, record: "Record") -> None: + if not self._fs: + return + out = record.output + stream: StreamLiterals = "stdout" + if out.output_type == wandb_internal_pb2.OutputRecord.OutputType.STDERR: + stream = "stderr" + line = out.line + self._send_output_line(stream, line) + + def send_output_raw(self, record: "Record") -> None: + if not self._fs: + return + out = record.output_raw + stream: StreamLiterals = "stdout" + if out.output_type == wandb_internal_pb2.OutputRawRecord.OutputType.STDERR: + stream = "stderr" + line = out.line + + output_raw = self._output_raw_streams.get(stream) + if not output_raw: + output_raw = _OutputRawStream(stream=stream, sm=self) + self._output_raw_streams[stream] = output_raw + + # open the console output file shared between both streams + if not self._output_raw_file: + output_log_path = os.path.join( + self._settings.files_dir, filenames.OUTPUT_FNAME + ) + output_raw_file = None + try: + output_raw_file = filesystem.CRDedupedFile( + open(output_log_path, "wb") + ) + except OSError as e: + logger.warning(f"could not open output_raw_file: {e}") + if output_raw_file: + self._output_raw_file = output_raw_file + output_raw.start() + + output_raw._queue.put(line) + + def _send_output_line(self, stream: "StreamLiterals", line: str) -> None: + """Combined writer for raw and non raw output lines. + + This is combined because they are both post emulator. + """ + prepend = "" + if stream == "stderr": + prepend = "ERROR " + if not line.endswith("\n"): + self._partial_output.setdefault(stream, "") + if line.startswith("\r"): + # TODO: maybe we shouldn't just drop this, what if there was some \ns in the partial + # that should probably be the check instead of not line.endswith(\n") + # logger.info(f"Dropping data {self._partial_output[stream]}") + self._partial_output[stream] = "" + self._partial_output[stream] += line + # TODO(jhr): how do we make sure this gets flushed? + # we might need this for other stuff like telemetry + else: + # TODO(jhr): use time from timestamp proto + # TODO(jhr): do we need to make sure we write full lines? + # seems to be some issues with line breaks + cur_time = time.time() + timestamp = datetime.utcfromtimestamp(cur_time).isoformat() + " " + prev_str = self._partial_output.get(stream, "") + line = f"{prepend}{timestamp}{prev_str}{line}" + if self._fs: + self._fs.push(filenames.OUTPUT_FNAME, line) + self._partial_output[stream] = "" + + def _update_config(self) -> None: + self._config_needs_debounce = True + + def send_config(self, record: "Record") -> None: + self._consolidated_config.update_from_proto(record.config) + self._update_config() + + def send_metric(self, record: "Record") -> None: + metric = record.metric + if metric.glob_name: + logger.warning("Seen metric with glob (shouldn't happen)") + return + + # merge or overwrite + old_metric = self._config_metric_dict.get( + metric.name, wandb_internal_pb2.MetricRecord() + ) + if metric._control.overwrite: + old_metric.CopyFrom(metric) + else: + old_metric.MergeFrom(metric) + self._config_metric_dict[metric.name] = old_metric + metric = old_metric + + # convert step_metric to index + if metric.step_metric: + find_step_idx = self._config_metric_index_dict.get(metric.step_metric) + if find_step_idx is not None: + # make a copy of this metric as we will be modifying it + rec = wandb_internal_pb2.Record() + rec.metric.CopyFrom(metric) + metric = rec.metric + + metric.ClearField("step_metric") + metric.step_metric_index = find_step_idx + 1 + + md: Dict[int, Any] = proto_util.proto_encode_to_dict(metric) + find_idx = self._config_metric_index_dict.get(metric.name) + if find_idx is not None: + self._config_metric_pbdict_list[find_idx] = md + else: + next_idx = len(self._config_metric_pbdict_list) + self._config_metric_pbdict_list.append(md) + self._config_metric_index_dict[metric.name] = next_idx + self._debounce_config() + + def _update_telemetry_record(self, telemetry: telemetry.TelemetryRecord) -> None: + self._telemetry_obj.MergeFrom(telemetry) + self._debounce_config() + + def send_telemetry(self, record: "Record") -> None: + self._update_telemetry_record(record.telemetry) + + def send_request_telemetry_record(self, record: "Record") -> None: + self._update_telemetry_record(record.request.telemetry_record.telemetry) + + def _save_file( + self, fname: interface.GlobStr, policy: "interface.PolicyName" = "end" + ) -> None: + logger.info("saving file %s with policy %s", fname, policy) + if self._dir_watcher: + self._dir_watcher.update_policy(fname, policy) + + def send_files(self, record: "Record") -> None: + files = record.files + for k in files.files: + # TODO(jhr): fix paths with directories + self._save_file( + interface.GlobStr(glob.escape(k.path)), + interface.file_enum_to_policy(k.policy), + ) + + def send_header(self, record: "Record") -> None: + pass + + def send_footer(self, record: "Record") -> None: + pass + + def send_tbrecord(self, record: "Record") -> None: + # tbrecord watching threads are handled by handler.py + pass + + def _update_environment_record(self, environment: "EnvironmentRecord") -> None: + self._environment_obj.MergeFrom(environment) + self._debounce_config() + + def send_environment(self, record: "Record") -> None: + """Inject environment info into config and upload as a JSON file.""" + self._update_environment_record(record.environment) + + environment_json = json.dumps(proto_util.message_to_dict(self._environment_obj)) + + with open( + os.path.join(self._settings.files_dir, filenames.METADATA_FNAME), "w" + ) as f: + f.write(environment_json) + + self._save_file(interface.GlobStr(filenames.METADATA_FNAME), policy="now") + + def send_request_link_artifact(self, record: "Record") -> None: + if not (record.control.req_resp or record.control.mailbox_slot): + raise ValueError( + f"Expected either `req_resp` or `mailbox_slot`, got: {record.control!r}" + ) + result = proto_util._result_from_record(record) + link = record.request.link_artifact + client_id = link.client_id + server_id = link.server_id + portfolio_name = link.portfolio_name + entity = link.portfolio_entity + project = link.portfolio_project + aliases = link.portfolio_aliases + organization = link.portfolio_organization + logger.debug( + f"link_artifact params - client_id={client_id}, server_id={server_id}, " + f"portfolio_name={portfolio_name}, entity={entity}, project={project}, " + f"organization={organization}" + ) + if (client_id or server_id) and portfolio_name and entity and project: + try: + response = self._api.link_artifact( + client_id, + server_id, + portfolio_name, + entity, + project, + aliases, + organization, + ) + result.response.link_artifact_response.version_index = response[ + "versionIndex" + ] + except Exception as e: + org_or_entity = organization or entity + result.response.link_artifact_response.error_message = ( + f"error linking artifact to " + f'"{org_or_entity}/{project}/{portfolio_name}"; error: {e}' + ) + logger.warning("Failed to link artifact to portfolio: %s", e) + self._respond_result(result) + + def send_use_artifact(self, record: "Record") -> None: + """Pretend to send a used artifact. + + This function doesn't actually send anything, it is just used internally. + """ + use = record.use_artifact + + if use.type == "job" and not use.partial.job_name: + self._job_builder.disable = True + elif use.partial.job_name: + # job is partial, let job builder rebuild job, set job source dict + self._job_builder.set_partial_source_id(use.id) + + def send_request_log_artifact(self, record: "Record") -> None: + result = proto_util._result_from_record(record) + artifact = record.request.log_artifact.artifact + history_step = record.request.log_artifact.history_step + + try: + res = self._send_artifact(artifact, history_step) + assert res, "Unable to send artifact" + result.response.log_artifact_response.artifact_id = res["id"] + logger.info(f"logged artifact {artifact.name} - {res}") + except Exception as e: + result.response.log_artifact_response.error_message = ( + f'error logging artifact "{artifact.type}/{artifact.name}": {e}' + ) + + self._respond_result(result) + + def send_artifact(self, record: "Record") -> None: + artifact = record.artifact + try: + res = self._send_artifact(artifact) + logger.info(f"sent artifact {artifact.name} - {res}") + except Exception: + logger.exception( + f'send_artifact: failed for artifact "{artifact.type}/{artifact.name}"' + ) + + def _send_artifact( + self, artifact: "ArtifactRecord", history_step: Optional[int] = None + ) -> Optional[Dict]: + from packaging.version import parse + + assert self._pusher + saver = ArtifactSaver( + api=self._api, + digest=artifact.digest, + manifest_json=_manifest_json_from_proto(artifact.manifest), + file_pusher=self._pusher, + is_user_created=artifact.user_created, + ) + + if artifact.distributed_id: + max_cli_version = self._max_cli_version() + if max_cli_version is None or parse(max_cli_version) < parse("0.10.16"): + logger.warning( + "This W&B Server doesn't support distributed artifacts, " + "have your administrator install wandb/local >= 0.9.37" + ) + return None + + metadata = json.loads(artifact.metadata) if artifact.metadata else None + res = saver.save( + entity=artifact.entity, + project=artifact.project, + type=artifact.type, + name=artifact.name, + client_id=artifact.client_id, + sequence_client_id=artifact.sequence_client_id, + metadata=metadata, + ttl_duration_seconds=artifact.ttl_duration_seconds or None, + description=artifact.description or None, + aliases=artifact.aliases, + tags=artifact.tags, + use_after_commit=artifact.use_after_commit, + distributed_id=artifact.distributed_id, + finalize=artifact.finalize, + incremental=artifact.incremental_beta1, + history_step=history_step, + base_id=artifact.base_id or None, + ) + + self._job_builder._handle_server_artifact(res, artifact) + + if artifact.manifest.manifest_file_path: + with contextlib.suppress(FileNotFoundError): + os.remove(artifact.manifest.manifest_file_path) + return res + + def send_alert(self, record: "Record") -> None: + from packaging.version import parse + + alert = record.alert + max_cli_version = self._max_cli_version() + if max_cli_version is None or parse(max_cli_version) < parse("0.10.9"): + logger.warning( + "This W&B server doesn't support alerts, " + "have your administrator install wandb/local >= 0.9.31" + ) + else: + try: + self._api.notify_scriptable_run_alert( + title=alert.title, + text=alert.text, + level=alert.level, + wait_duration=alert.wait_duration, + ) + except Exception: + logger.exception(f"send_alert: failed for alert {alert.title!r}") + + def finish(self) -> None: + logger.info("shutting down sender") + # if self._tb_watcher: + # self._tb_watcher.finish() + self._output_raw_finish() + if self._dir_watcher: + self._dir_watcher.finish() + self._dir_watcher = None + if self._pusher: + self._pusher.finish() + self._pusher.join() + self._pusher = None + if self._fs: + self._fs.finish(self._exit_code) + self._fs = None + wandb._sentry.end_session() + + def _max_cli_version(self) -> Optional[str]: + server_info = self.get_server_info() + max_cli_version = server_info.get("cliVersionInfo", {}).get( + "max_cli_version", None + ) + if not isinstance(max_cli_version, str): + return None + return max_cli_version + + def get_viewer_server_info(self) -> None: + if self._cached_server_info and self._cached_viewer: + return + self._cached_viewer, self._cached_server_info = self._api.viewer_server_info() + + def get_viewer_info(self) -> Dict[str, Any]: + if not self._cached_viewer: + self.get_viewer_server_info() + return self._cached_viewer + + def get_server_info(self) -> Dict[str, Any]: + if not self._cached_server_info: + self.get_viewer_server_info() + return self._cached_server_info + + def get_local_info(self) -> "LocalInfo": + """Queries the server to get the local version information. + + First, we perform an introspection, if it returns empty we deduce that the + docker image is out-of-date. Otherwise, we use the returned values to deduce the + state of the local server. + """ + local_info = wandb_internal_pb2.LocalInfo() + if self._settings._offline: + local_info.out_of_date = False + return local_info + + latest_local_version = "latest" + + # Assuming the query is successful if the result is empty it indicates that + # the backend is out of date since it doesn't have the desired field + server_info = self.get_server_info() + latest_local_version_info = server_info.get("latestLocalVersionInfo", {}) + if latest_local_version_info is None: + local_info.out_of_date = False + else: + local_info.out_of_date = latest_local_version_info.get("outOfDate", True) + local_info.version = latest_local_version_info.get( + "latestVersionString", latest_local_version + ) + return local_info + + def _flush_job(self) -> None: + if self._job_builder.disable or self._settings._offline: + return + self._job_builder.set_config(self._consolidated_config.non_internal_config()) + summary_dict = self._cached_summary.copy() + summary_dict.pop("_wandb", None) + self._job_builder.set_summary(summary_dict) + + artifact = self._job_builder.build(api=self._api) + if artifact is not None and self._run is not None: + proto_artifact = self._interface._make_artifact(artifact) + proto_artifact.run_id = self._run.run_id + proto_artifact.project = self._run.project + proto_artifact.entity = self._run.entity + # TODO: this should be removed when the latest tag is handled + # by the backend (WB-12116) + proto_artifact.aliases.append("latest") + # add docker image tag + for alias in self._job_builder._aliases: + proto_artifact.aliases.append(alias) + + proto_artifact.user_created = True + proto_artifact.use_after_commit = True + proto_artifact.finalize = True + + self._interface._publish_artifact(proto_artifact) + + def __next__(self) -> "Record": + return self._record_q.get(block=True) + + next = __next__ diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/sender_config.py b/lib/python3.12/site-packages/wandb/sdk/internal/sender_config.py new file mode 100644 index 0000000000000000000000000000000000000000..0cc478958173637fefba1c620977af1d52faad0d --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/internal/sender_config.py @@ -0,0 +1,203 @@ +import json +from typing import Any, Dict, NewType, Optional, Sequence + +from wandb.proto import wandb_internal_pb2 +from wandb.sdk.lib import proto_util, telemetry + +BackendConfigDict = NewType("BackendConfigDict", Dict[str, Any]) +"""Run config dictionary in the format used by the backend.""" + +_WANDB_INTERNAL_KEY = "_wandb" + + +class ConfigState: + """The configuration of a run.""" + + def __init__(self, tree: Optional[Dict[str, Any]] = None) -> None: + self._tree: Dict[str, Any] = tree or {} + """A tree with string-valued nodes and JSON leaves. + + Leaves are Python objects that are valid JSON values: + + * Primitives like strings and numbers + * Dictionaries from strings to JSON objects + * Lists of JSON objects + """ + + def non_internal_config(self) -> Dict[str, Any]: + """Returns the config settings minus "_wandb".""" + return {k: v for k, v in self._tree.items() if k != _WANDB_INTERNAL_KEY} + + def update_from_proto( + self, + config_record: wandb_internal_pb2.ConfigRecord, + ) -> None: + """Applies update and remove commands.""" + for config_item in config_record.update: + self._update_at_path( + _key_path(config_item), + json.loads(config_item.value_json), + ) + + for config_item in config_record.remove: + self._delete_at_path(_key_path(config_item)) + + def merge_resumed_config(self, old_config_tree: Dict[str, Any]) -> None: + """Merges the config from a run that's being resumed.""" + # Add any top-level keys that aren't already set. + self._add_unset_keys_from_subtree(old_config_tree, []) + + # When resuming a run, we want to ensure the some of the old configs keys + # are maintained. So we have this logic here to add back + # any keys that were in the old config but not in the new config + for key in ["viz", "visualize", "mask/class_labels"]: + self._add_unset_keys_from_subtree( + old_config_tree, + [_WANDB_INTERNAL_KEY, key], + ) + + def _add_unset_keys_from_subtree( + self, + old_config_tree: Dict[str, Any], + path: Sequence[str], + ) -> None: + """Uses the given subtree for keys that aren't already set.""" + old_subtree = _subtree(old_config_tree, path, create=False) + if not old_subtree: + return + + new_subtree = _subtree(self._tree, path, create=True) + assert new_subtree is not None + + for key, value in old_subtree.items(): + if key not in new_subtree: + new_subtree[key] = value + + def to_backend_dict( + self, + telemetry_record: telemetry.TelemetryRecord, + framework: Optional[str], + start_time_millis: int, + metric_pbdicts: Sequence[Dict[int, Any]], + environment_record: wandb_internal_pb2.EnvironmentRecord, + ) -> BackendConfigDict: + """Returns a dictionary representation expected by the backend. + + The backend expects the configuration in a specific format, and the + config is also used to store additional metadata about the run. + + Args: + telemetry_record: Telemetry information to insert. + framework: The detected framework used in the run (e.g. TensorFlow). + start_time_millis: The run's start time in Unix milliseconds. + metric_pbdicts: List of dict representations of metric protobuffers. + """ + backend_dict = self._tree.copy() + wandb_internal = backend_dict.setdefault(_WANDB_INTERNAL_KEY, {}) + + ################################################### + # Telemetry information + ################################################### + py_version = telemetry_record.python_version + if py_version: + wandb_internal["python_version"] = py_version + + cli_version = telemetry_record.cli_version + if cli_version: + wandb_internal["cli_version"] = cli_version + + if framework: + wandb_internal["framework"] = framework + + huggingface_version = telemetry_record.huggingface_version + if huggingface_version: + wandb_internal["huggingface_version"] = huggingface_version + + wandb_internal["is_jupyter_run"] = telemetry_record.env.jupyter + wandb_internal["is_kaggle_kernel"] = telemetry_record.env.kaggle + wandb_internal["start_time"] = start_time_millis + + # The full telemetry record. + wandb_internal["t"] = proto_util.proto_encode_to_dict(telemetry_record) + + ################################################### + # Metrics + ################################################### + if metric_pbdicts: + wandb_internal["m"] = metric_pbdicts + + ################################################### + # Environment + ################################################### + writer_id = environment_record.writer_id + if writer_id: + environment_dict = proto_util.message_to_dict(environment_record) + wandb_internal["e"] = {writer_id: environment_dict} + + return BackendConfigDict( + { + key: { + # Configurations can be stored in a hand-written YAML file, + # and users can add descriptions to their hyperparameters + # there. However, we don't support a way to set descriptions + # via code, so this is always None. + "desc": None, + "value": value, + } + for key, value in self._tree.items() + } + ) + + def _update_at_path( + self, + key_path: Sequence[str], + value: Any, + ) -> None: + """Sets the value at the path in the config tree.""" + subtree = _subtree(self._tree, key_path[:-1], create=True) + assert subtree is not None + + subtree[key_path[-1]] = value + + def _delete_at_path( + self, + key_path: Sequence[str], + ) -> None: + """Removes the subtree at the path in the config tree.""" + subtree = _subtree(self._tree, key_path[:-1], create=False) + if subtree: + del subtree[key_path[-1]] + + +def _key_path(config_item: wandb_internal_pb2.ConfigItem) -> Sequence[str]: + """Returns the key path referenced by the config item.""" + if config_item.nested_key: + return config_item.nested_key + elif config_item.key: + return [config_item.key] + else: + raise AssertionError( + "Invalid ConfigItem: either key or nested_key must be set", + ) + + +def _subtree( + tree: Dict[str, Any], + key_path: Sequence[str], + *, + create: bool = False, +) -> Optional[Dict[str, Any]]: + """Returns a subtree at the given path.""" + for key in key_path: + subtree = tree.get(key) + + if not subtree: + if create: + subtree = {} + tree[key] = subtree + else: + return None + + tree = subtree + + return tree diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/settings_static.py b/lib/python3.12/site-packages/wandb/sdk/internal/settings_static.py new file mode 100644 index 0000000000000000000000000000000000000000..f3dfd6b25315f300e6461c74021a77b0fad207ad --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/internal/settings_static.py @@ -0,0 +1,121 @@ +from __future__ import annotations + +from typing import Any, Iterable + +from wandb.proto import wandb_settings_pb2 +from wandb.sdk.lib import RunMoment +from wandb.sdk.wandb_settings import Settings + + +class SettingsStatic(Settings): + """A readonly object that wraps a protobuf Settings message. + + Implements the mapping protocol, so you can access settings as + attributes or items. + """ + + def __init__(self, proto: wandb_settings_pb2.Settings) -> None: + data = self._proto_to_dict(proto) + super().__init__(**data) + + def _proto_to_dict(self, proto: wandb_settings_pb2.Settings) -> dict: + data = {} + + exclude_fields = { + "model_config", + "model_fields", + "model_fields_set", + "__fields__", + "__model_fields_set", + "__pydantic_self__", + "__pydantic_initialised__", + } + + fields = ( + Settings.model_fields + if hasattr(Settings, "model_fields") + else Settings.__fields__ + ) # type: ignore [attr-defined] + + fields = {k: v for k, v in fields.items() if k not in exclude_fields} # type: ignore [union-attr] + + forks_specified: list[str] = [] + for key in fields: + # Skip Python-only keys that do not exist on the proto. + if key in ("reinit",): + continue + + value: Any = None + + field_info = fields[key] + annotation = str(field_info.annotation) + + if key == "_stats_open_metrics_filters": + # todo: it's an underscored field, refactor into + # something more elegant? + # I'm really about this. It's ugly, but it works. + # Do not try to repeat this at home. + value_type = getattr(proto, key).WhichOneof("value") + if value_type == "sequence": + value = list(getattr(proto, key).sequence.value) + elif value_type == "mapping": + unpacked_mapping = {} + for outer_key, outer_value in getattr( + proto, key + ).mapping.value.items(): + unpacked_inner = {} + for inner_key, inner_value in outer_value.value.items(): + unpacked_inner[inner_key] = inner_value + unpacked_mapping[outer_key] = unpacked_inner + value = unpacked_mapping + elif key == "fork_from" or key == "resume_from": + value = getattr(proto, key) + if value.run: + value = RunMoment( + run=value.run, value=value.value, metric=value.metric + ) + forks_specified.append(key) + else: + value = None + else: + if proto.HasField(key): # type: ignore [arg-type] + value = getattr(proto, key).value + # Convert to list if the field is a sequence + if any(t in annotation for t in ("tuple", "Sequence", "list")): + value = list(value) + else: + value = None + + if value is not None: + data[key] = value + + if len(forks_specified) > 1: + raise ValueError( + "Only one of fork_from or resume_from can be specified, not both" + ) + + return data + + def __setattr__(self, name: str, value: object) -> None: + raise AttributeError("Error: SettingsStatic is a readonly object") + + def __setitem__(self, key: str, val: object) -> None: + raise AttributeError("Error: SettingsStatic is a readonly object") + + def keys(self) -> Iterable[str]: + return self.__dict__.keys() + + def __getitem__(self, key: str) -> Any: + return self.__dict__[key] + + def __getattr__(self, name: str) -> Any: + try: + return self.__dict__[name] + except KeyError: + raise AttributeError(f"SettingsStatic has no attribute {name}") + + def __str__(self) -> str: + return str(self.__dict__) + + def __contains__(self, key: str) -> bool: + return key in self.__dict__ diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/tb_watcher.py b/lib/python3.12/site-packages/wandb/sdk/internal/tb_watcher.py new file mode 100644 index 0000000000000000000000000000000000000000..67cd1b5e6aac85cdd7a93bf1d393fa55c09aefcc --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/internal/tb_watcher.py @@ -0,0 +1,519 @@ +"""tensorboard watcher.""" + +import glob +import logging +import os +import queue +import socket +import sys +import threading +import time +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +import wandb +from wandb import util +from wandb.plot import CustomChart +from wandb.sdk.interface.interface import GlobStr +from wandb.sdk.lib import filesystem + +from . import run as internal_run + +if TYPE_CHECKING: + from queue import PriorityQueue + + from tensorboard.backend.event_processing.event_file_loader import EventFileLoader + from tensorboard.compat.proto.event_pb2 import ProtoEvent + + from wandb.proto.wandb_internal_pb2 import RunRecord + from wandb.sdk.interface.interface import FilesDict + + from ..interface.interface_queue import InterfaceQueue + from .settings_static import SettingsStatic + + HistoryDict = Dict[str, Any] + +# Give some time for tensorboard data to be flushed +SHUTDOWN_DELAY = 5 +ERROR_DELAY = 5 +REMOTE_FILE_TOKEN = "://" +logger = logging.getLogger(__name__) + + +def _link_and_save_file( + path: str, base_path: str, interface: "InterfaceQueue", settings: "SettingsStatic" +) -> None: + # TODO(jhr): should this logic be merged with Run.save() + files_dir = settings.files_dir + file_name = os.path.relpath(path, base_path) + abs_path = os.path.abspath(path) + wandb_path = os.path.join(files_dir, file_name) + filesystem.mkdir_exists_ok(os.path.dirname(wandb_path)) + # We overwrite existing symlinks because namespaces can change in Tensorboard + if os.path.islink(wandb_path) and abs_path != os.readlink(wandb_path): + os.remove(wandb_path) + os.symlink(abs_path, wandb_path) + elif not os.path.exists(wandb_path): + os.symlink(abs_path, wandb_path) + # TODO(jhr): need to figure out policy, live/throttled? + interface.publish_files(dict(files=[(GlobStr(glob.escape(file_name)), "live")])) + + +def is_tfevents_file_created_by( + path: str, hostname: Optional[str], start_time: Optional[float] +) -> bool: + """Check if a path is a tfevents file. + + Optionally checks that it was created by [hostname] after [start_time]. + + tensorboard tfevents filename format: + https://github.com/tensorflow/tensorboard/blob/f3f26b46981da5bd46a5bb93fcf02d9eb7608bc1/tensorboard/summary/writer/event_file_writer.py#L81 + tensorflow tfevents filename format: + https://github.com/tensorflow/tensorflow/blob/8f597046dc30c14b5413813d02c0e0aed399c177/tensorflow/core/util/events_writer.cc#L68 + """ + if not path: + raise ValueError("Path must be a nonempty string") + basename = os.path.basename(path) + if basename.endswith((".profile-empty", ".sagemaker-uploaded")): + return False + fname_components = basename.split(".") + try: + tfevents_idx = fname_components.index("tfevents") + except ValueError: + return False + # check the hostname, which may have dots + if hostname is not None: + for i, part in enumerate(hostname.split(".")): + try: + fname_component_part = fname_components[tfevents_idx + 2 + i] + except IndexError: + return False + if part != fname_component_part: + return False + if start_time is not None: + try: + created_time = int(fname_components[tfevents_idx + 1]) + except (ValueError, IndexError): + return False + # Ensure that the file is newer then our start time, and that it was + # created from the same hostname. + # TODO: we should also check the PID (also contained in the tfevents + # filename). Can we assume that our parent pid is the user process + # that wrote these files? + if created_time < int(start_time): + return False + return True + + +class TBWatcher: + _logdirs: "Dict[str, TBDirWatcher]" + _watcher_queue: "PriorityQueue" + + def __init__( + self, + settings: "SettingsStatic", + run_proto: "RunRecord", + interface: "InterfaceQueue", + force: bool = False, + ) -> None: + self._logdirs = {} + self._consumer: Optional[TBEventConsumer] = None + self._settings = settings + self._interface = interface + self._run_proto = run_proto + self._force = force + # TODO(jhr): do we need locking in this queue? + self._watcher_queue = queue.PriorityQueue() + wandb.tensorboard.reset_state() # type: ignore + + def _calculate_namespace(self, logdir: str, rootdir: str) -> Optional[str]: + namespace: Optional[str] + dirs = list(self._logdirs) + [logdir] + + if os.path.isfile(logdir): + filename = os.path.basename(logdir) + else: + filename = "" + + if rootdir == "": + rootdir = util.to_forward_slash_path( + os.path.dirname(os.path.commonprefix(dirs)) + ) + # Tensorboard loads all tfevents files in a directory and prepends + # their values with the path. Passing namespace to log allows us + # to nest the values in wandb + # Note that we strip '/' instead of os.sep, because elsewhere we've + # converted paths to forward slash. + namespace = logdir.replace(filename, "").replace(rootdir, "").strip("/") + + # TODO: revisit this heuristic, it exists because we don't know the + # root log directory until more than one tfevents file is written to + if len(dirs) == 1 and namespace not in ["train", "validation"]: + namespace = None + else: + namespace = logdir.replace(filename, "").replace(rootdir, "").strip("/") + + return namespace + + def add(self, logdir: str, save: bool, root_dir: str) -> None: + logdir = util.to_forward_slash_path(logdir) + root_dir = util.to_forward_slash_path(root_dir) + if logdir in self._logdirs: + return + namespace = self._calculate_namespace(logdir, root_dir) + # TODO(jhr): implement the deferred tbdirwatcher to find namespace + + if not self._consumer: + self._consumer = TBEventConsumer( + self, self._watcher_queue, self._run_proto, self._settings + ) + self._consumer.start() + + tbdir_watcher = TBDirWatcher( + self, logdir, save, namespace, self._watcher_queue, self._force + ) + self._logdirs[logdir] = tbdir_watcher + tbdir_watcher.start() + + def finish(self) -> None: + for tbdirwatcher in self._logdirs.values(): + tbdirwatcher.shutdown() + for tbdirwatcher in self._logdirs.values(): + tbdirwatcher.finish() + if self._consumer: + self._consumer.finish() + + +class TBDirWatcher: + def __init__( + self, + tbwatcher: "TBWatcher", + logdir: str, + save: bool, + namespace: Optional[str], + queue: "PriorityQueue", + force: bool = False, + ) -> None: + self.directory_watcher = util.get_module( + "tensorboard.backend.event_processing.directory_watcher", + required="Please install tensorboard package", + ) + # self.event_file_loader = util.get_module( + # "tensorboard.backend.event_processing.event_file_loader", + # required="Please install tensorboard package", + # ) + self.tf_compat = util.get_module( + "tensorboard.compat", required="Please install tensorboard package" + ) + self._tbwatcher = tbwatcher + self._generator = self.directory_watcher.DirectoryWatcher( + logdir, self._loader(save, namespace), self._is_our_tfevents_file + ) + self._thread = threading.Thread(target=self._thread_except_body) + self._first_event_timestamp = None + self._shutdown = threading.Event() + self._queue = queue + self._file_version = None + self._namespace = namespace + self._logdir = logdir + self._hostname = socket.gethostname() + self._force = force + self._process_events_lock = threading.Lock() + + def start(self) -> None: + self._thread.start() + + def _is_our_tfevents_file(self, path: str) -> bool: + """Check if a path has been modified since launch and contains tfevents.""" + if not path: + raise ValueError("Path must be a nonempty string") + path = self.tf_compat.tf.compat.as_str_any(path) + if self._force: + return is_tfevents_file_created_by(path, None, None) + else: + return is_tfevents_file_created_by( + path, self._hostname, self._tbwatcher._settings.x_start_time + ) + + def _loader( + self, save: bool = True, namespace: Optional[str] = None + ) -> "EventFileLoader": + """Incredibly hacky class generator to optionally save / prefix tfevent files.""" + _loader_interface = self._tbwatcher._interface + _loader_settings = self._tbwatcher._settings + try: + from tensorboard.backend.event_processing import event_file_loader + except ImportError: + raise Exception("Please install tensorboard package") + + class EventFileLoader(event_file_loader.EventFileLoader): + def __init__(self, file_path: str) -> None: + super().__init__(file_path) + if save: + if REMOTE_FILE_TOKEN in file_path: + logger.warning( + "Not persisting remote tfevent file: %s", file_path + ) + else: + # TODO: save plugins? + logdir = os.path.dirname(file_path) + parts = list(os.path.split(logdir)) + if namespace and parts[-1] == namespace: + parts.pop() + logdir = os.path.join(*parts) + _link_and_save_file( + path=file_path, + base_path=logdir, + interface=_loader_interface, + settings=_loader_settings, + ) + + return EventFileLoader + + def _process_events(self, shutdown_call: bool = False) -> None: + try: + with self._process_events_lock: + for event in self._generator.Load(): + self.process_event(event) + except ( + self.directory_watcher.DirectoryDeletedError, + StopIteration, + RuntimeError, + OSError, + ) as e: + # When listing s3 the directory may not yet exist, or could be empty + logger.debug("Encountered tensorboard directory watcher error: %s", e) + if not self._shutdown.is_set() and not shutdown_call: + time.sleep(ERROR_DELAY) + + def _thread_except_body(self) -> None: + try: + self._thread_body() + except Exception: + logger.exception("generic exception in TBDirWatcher thread") + raise + + def _thread_body(self) -> None: + """Check for new events every second.""" + shutdown_time: Optional[float] = None + while True: + self._process_events() + if self._shutdown.is_set(): + now = time.time() + if not shutdown_time: + shutdown_time = now + SHUTDOWN_DELAY + elif now > shutdown_time: + break + time.sleep(1) + + def process_event(self, event: "ProtoEvent") -> None: + # print("\nEVENT:::", self._logdir, self._namespace, event, "\n") + if self._first_event_timestamp is None: + self._first_event_timestamp = event.wall_time + + if event.HasField("file_version"): + self._file_version = event.file_version + + if event.HasField("summary"): + self._queue.put(Event(event, self._namespace)) + + def shutdown(self) -> None: + self._process_events(shutdown_call=True) + self._shutdown.set() + + def finish(self) -> None: + self.shutdown() + self._thread.join() + + +class Event: + """An event wrapper to enable priority queueing.""" + + def __init__(self, event: "ProtoEvent", namespace: Optional[str]): + self.event = event + self.namespace = namespace + self.created_at = time.time() + + def __lt__(self, other: "Event") -> bool: + if self.event.wall_time < other.event.wall_time: + return True + return False + + +class TBEventConsumer: + """Consume tfevents from a priority queue. + + There should always only be one of these per run_manager. We wait for 10 seconds of + queued events to reduce the chance of multiple tfevent files triggering out of order + steps. + """ + + def __init__( + self, + tbwatcher: TBWatcher, + queue: "PriorityQueue", + run_proto: "RunRecord", + settings: "SettingsStatic", + delay: int = 10, + ) -> None: + self._tbwatcher = tbwatcher + self._queue = queue + self._thread = threading.Thread(target=self._thread_except_body) + self._shutdown = threading.Event() + self.tb_history = TBHistory() + self._delay = delay + + # This is a bit of a hack to get file saving to work as it does in the user + # process. Since we don't have a real run object, we have to define the + # datatypes callback ourselves. + def datatypes_cb(fname: GlobStr) -> None: + files: FilesDict = dict(files=[(fname, "now")]) + self._tbwatcher._interface.publish_files(files) + + # this is only used for logging artifacts + self._internal_run = internal_run.InternalRun(run_proto, settings, datatypes_cb) + self._internal_run._set_internal_run_interface(self._tbwatcher._interface) + + def start(self) -> None: + self._start_time = time.time() + self._thread.start() + + def finish(self) -> None: + self._delay = 0 + self._shutdown.set() + self._thread.join() + while not self._queue.empty(): + event = self._queue.get(True, 1) + if event: + self._handle_event(event, history=self.tb_history) + items = self.tb_history._get_and_reset() + for item in items: + self._save_row( + item, + ) + + def _thread_except_body(self) -> None: + try: + self._thread_body() + except Exception: + logger.exception("generic exception in TBEventConsumer thread") + raise + + def _thread_body(self) -> None: + while True: + try: + event = self._queue.get(True, 1) + # Wait self._delay seconds from consumer start before logging events + if ( + time.time() < self._start_time + self._delay + and not self._shutdown.is_set() + ): + self._queue.put(event) + time.sleep(0.1) + continue + except queue.Empty: + event = None + if self._shutdown.is_set(): + break + if event: + self._handle_event(event, history=self.tb_history) + items = self.tb_history._get_and_reset() + for item in items: + self._save_row( + item, + ) + # flush uncommitted data + self.tb_history._flush() + items = self.tb_history._get_and_reset() + for item in items: + self._save_row(item) + + def _handle_event( + self, event: "ProtoEvent", history: Optional["TBHistory"] = None + ) -> None: + wandb.tensorboard._log( # type: ignore + event.event, + step=event.event.step, + namespace=event.namespace, + history=history, + ) + + def _save_row(self, row: "HistoryDict") -> None: + chart_keys = set() + for k, v in row.items(): + if isinstance(v, CustomChart): + chart_keys.add(k) + v.set_key(k) + self._tbwatcher._interface.publish_config( + key=v.spec.config_key, + val=v.spec.config_value, + ) + + for k in chart_keys: + chart = row.pop(k) + if isinstance(chart, CustomChart): + row[chart.spec.table_key] = chart.table + + self._tbwatcher._interface.publish_history( + self._internal_run, + row, + publish_step=False, + ) + + +class TBHistory: + _data: "HistoryDict" + _added: "List[HistoryDict]" + + def __init__(self) -> None: + self._step = 0 + self._step_size = 0 + self._data = dict() + self._added = [] + + def _flush(self) -> None: + if not self._data: + return + # A single tensorboard step may have too much data + # we just drop the largest keys in the step if it does. + # TODO: we could flush the data across multiple steps + if self._step_size > util.MAX_LINE_BYTES: + metrics = [(k, sys.getsizeof(v)) for k, v in self._data.items()] + metrics.sort(key=lambda t: t[1], reverse=True) + bad = 0 + dropped_keys = [] + for k, v in metrics: + # TODO: (cvp) Added a buffer of 100KiB, this feels rather brittle. + if self._step_size - bad < util.MAX_LINE_BYTES - 100000: + break + else: + bad += v + dropped_keys.append(k) + del self._data[k] + wandb.termwarn( + f"Step {self._step} exceeds max data limit, dropping {len(dropped_keys)} of the largest keys:" + ) + print("\t" + ("\n\t".join(dropped_keys))) # noqa: T201 + self._data["_step"] = self._step + self._added.append(self._data) + self._step += 1 + self._step_size = 0 + + def add(self, d: "HistoryDict") -> None: + self._flush() + self._data = dict() + self._data.update(self._track_history_dict(d)) + + def _track_history_dict(self, d: "HistoryDict") -> "HistoryDict": + e = {} + for k in d.keys(): + e[k] = d[k] + self._step_size += sys.getsizeof(e[k]) + return e + + def _row_update(self, d: "HistoryDict") -> None: + self._data.update(self._track_history_dict(d)) + + def _get_and_reset(self) -> "List[HistoryDict]": + added = self._added[:] + self._added = [] + return added diff --git a/lib/python3.12/site-packages/wandb/sdk/internal/thread_local_settings.py b/lib/python3.12/site-packages/wandb/sdk/internal/thread_local_settings.py new file mode 100644 index 0000000000000000000000000000000000000000..2ee0e74cc44151617cb90e7eb45996e0d1a3859d --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/internal/thread_local_settings.py @@ -0,0 +1,18 @@ +import threading +from typing import Dict, Optional + + +# Context variable for setting API settings (api keys, etc.) for internal and public apis thread-locally +# TODO: move this into actual settings +class _ThreadLocalApiSettings(threading.local): + api_key: Optional[str] + cookies: Optional[Dict] + headers: Optional[Dict] + + def __init__(self) -> None: + self.api_key = None + self.cookies = None + self.headers = None + + +_thread_local_api_settings: _ThreadLocalApiSettings = _ThreadLocalApiSettings() diff --git a/lib/python3.12/site-packages/wandb/sdk/projects/_generated/upsert_registry_project.py b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/upsert_registry_project.py new file mode 100644 index 0000000000000000000000000000000000000000..384a480ab317f6d5ec2f66556b85743cb9587ecc --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/projects/_generated/upsert_registry_project.py @@ -0,0 +1,27 @@ +# Generated by ariadne-codegen +# Source: tools/graphql_codegen/projects/ + +from __future__ import annotations + +from typing import Optional + +from pydantic import Field + +from wandb._pydantic import GQLBase + +from .fragments import RegistryFragment + + +class UpsertRegistryProject(GQLBase): + upsert_model: Optional[UpsertRegistryProjectUpsertModel] = Field( + alias="upsertModel" + ) + + +class UpsertRegistryProjectUpsertModel(GQLBase): + project: Optional[RegistryFragment] + inserted: Optional[bool] + + +UpsertRegistryProject.model_rebuild() +UpsertRegistryProjectUpsertModel.model_rebuild() diff --git a/lib/python3.12/site-packages/wandb/sdk/wandb_alerts.py b/lib/python3.12/site-packages/wandb/sdk/wandb_alerts.py new file mode 100644 index 0000000000000000000000000000000000000000..072547ab219e907efede3399d660817926074490 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/wandb_alerts.py @@ -0,0 +1,12 @@ +# +from enum import Enum + +""" +Call run.alert() to generate an email or Slack notification programmatically. +""" + + +class AlertLevel(Enum): + INFO = "INFO" + WARN = "WARN" + ERROR = "ERROR" diff --git a/lib/python3.12/site-packages/wandb/sdk/wandb_config.py b/lib/python3.12/site-packages/wandb/sdk/wandb_config.py new file mode 100644 index 0000000000000000000000000000000000000000..35a1693bc55cd674a724d1081756e89c20adeedf --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/wandb_config.py @@ -0,0 +1,323 @@ +"""config.""" + +import logging +from typing import Optional + +import wandb +from wandb.util import ( + _is_artifact_representation, + check_dict_contains_nested_artifact, + json_friendly_val, +) + +from . import wandb_helper +from .lib import config_util + +logger = logging.getLogger("wandb") + + +# TODO(jhr): consider a callback for persisting changes? +# if this is done right we might make sure this is pickle-able +# we might be able to do this on other objects like Run? +class Config: + """Config object. + + Config objects are intended to hold all of the hyperparameters associated + with a wandb run and are saved with the run object when `wandb.init` is + called. + + We recommend setting the config once when initializing your run by passing + the `config` parameter to `init`: + + ``` + wandb.init(config=my_config_dict) + ``` + + You can create a file called `config-defaults.yaml`, and it will + automatically be loaded as each run's config. You can also pass the name + of the file as the `config` parameter to `init`: + + ``` + wandb.init(config="my_config.yaml") + ``` + + See https://docs.wandb.com/guides/track/config#file-based-configs. + + Examples: + Basic usage + ``` + with wandb.init(config={"epochs": 4}) as run: + for x in range(run.config.epochs): + # train + ``` + + Nested values + ``` + with wandb.init(config={"train": {"epochs": 4}}) as run: + for x in range(run.config["train"]["epochs"]): + # train + ``` + + Using absl flags + ``` + flags.DEFINE_string("model", None, "model to run") # name, default, help + with wandb.init() as run: + run.config.update(flags.FLAGS) # adds all absl flags to config + ``` + + Argparse flags + ```python + with wandb.init(config={"epochs": 4}) as run: + parser = argparse.ArgumentParser() + parser.add_argument( + "-b", + "--batch-size", + type=int, + default=8, + metavar="N", + help="input batch size for training (default: 8)", + ) + args = parser.parse_args() + run.config.update(args) + ``` + + Using TensorFlow flags (deprecated in tensorflow v2) + ```python + flags = tf.app.flags + flags.DEFINE_string("data_dir", "/tmp/data") + flags.DEFINE_integer("batch_size", 128, "Batch size.") + + with wandb.init() as run: + run.config.update(flags.FLAGS) + ``` + """ + + def __init__(self): + object.__setattr__(self, "_items", dict()) + object.__setattr__(self, "_locked", dict()) + object.__setattr__(self, "_users", dict()) + object.__setattr__(self, "_users_inv", dict()) + object.__setattr__(self, "_users_cnt", 0) + object.__setattr__(self, "_callback", None) + object.__setattr__(self, "_settings", None) + object.__setattr__(self, "_artifact_callback", None) + + self._load_defaults() + + def _set_callback(self, cb): + object.__setattr__(self, "_callback", cb) + + def _set_artifact_callback(self, cb): + object.__setattr__(self, "_artifact_callback", cb) + + def _set_settings(self, settings): + object.__setattr__(self, "_settings", settings) + + def __repr__(self): + return str(dict(self)) + + def keys(self): + return [k for k in self._items.keys() if not k.startswith("_")] + + def _as_dict(self): + return self._items + + def as_dict(self): + # TODO: add telemetry, deprecate, then remove + return dict(self) + + def __getitem__(self, key): + return self._items[key] + + def __iter__(self): + return iter(self._items) + + def _check_locked(self, key, ignore_locked=False) -> bool: + locked = self._locked.get(key) + if locked is not None: + locked_user = self._users_inv[locked] + if not ignore_locked: + wandb.termwarn( + f"Config item '{key}' was locked by '{locked_user}' (ignored update)." + ) + return True + return False + + def __setitem__(self, key, val): + if self._check_locked(key): + return + with wandb.sdk.lib.telemetry.context() as tel: + tel.feature.set_config_item = True + self._raise_value_error_on_nested_artifact(val, nested=True) + key, val = self._sanitize(key, val) + self._items[key] = val + logger.info("config set %s = %s - %s", key, val, self._callback) + if self._callback: + self._callback(key=key, val=val) + + def items(self): + return [(k, v) for k, v in self._items.items() if not k.startswith("_")] + + __setattr__ = __setitem__ + + def __getattr__(self, key): + try: + return self.__getitem__(key) + except KeyError as ke: + raise AttributeError( + f"{self.__class__!r} object has no attribute {key!r}" + ) from ke + + def __contains__(self, key): + return key in self._items + + def _update(self, d, allow_val_change=None, ignore_locked=None): + parsed_dict = wandb_helper.parse_config(d) + locked_keys = set() + for key in list(parsed_dict): + if self._check_locked(key, ignore_locked=ignore_locked): + locked_keys.add(key) + sanitized = self._sanitize_dict( + parsed_dict, allow_val_change, ignore_keys=locked_keys + ) + self._items.update(sanitized) + return sanitized + + def update(self, d, allow_val_change=None): + sanitized = self._update(d, allow_val_change) + if self._callback: + self._callback(data=sanitized) + + def get(self, *args): + return self._items.get(*args) + + def persist(self): + """Call the callback if it's set.""" + if self._callback: + self._callback(data=self._as_dict()) + + def setdefaults(self, d): + d = wandb_helper.parse_config(d) + # strip out keys already configured + d = {k: v for k, v in d.items() if k not in self._items} + d = self._sanitize_dict(d) + self._items.update(d) + if self._callback: + self._callback(data=d) + + def _get_user_id(self, user) -> int: + if user not in self._users: + self._users[user] = self._users_cnt + self._users_inv[self._users_cnt] = user + object.__setattr__(self, "_users_cnt", self._users_cnt + 1) + + return self._users[user] + + def update_locked(self, d, user=None, _allow_val_change=None): + """Shallow-update config with `d` and lock config updates on d's keys.""" + num = self._get_user_id(user) + + for k, v in d.items(): + k, v = self._sanitize(k, v, allow_val_change=_allow_val_change) + self._locked[k] = num + self._items[k] = v + + if self._callback: + self._callback(data=d) + + def merge_locked(self, d, user=None, _allow_val_change=None): + """Recursively merge-update config with `d` and lock config updates on d's keys.""" + num = self._get_user_id(user) + callback_d = {} + + for k, v in d.items(): + k, v = self._sanitize(k, v, allow_val_change=_allow_val_change) + self._locked[k] = num + + if ( + k in self._items + and isinstance(self._items[k], dict) + and isinstance(v, dict) + ): + self._items[k] = config_util.merge_dicts(self._items[k], v) + else: + self._items[k] = v + + callback_d[k] = self._items[k] + + if self._callback: + self._callback(data=callback_d) + + def _load_defaults(self): + conf_dict = config_util.dict_from_config_file("config-defaults.yaml") + if conf_dict is not None: + self.update(conf_dict) + + def _sanitize_dict( + self, + config_dict, + allow_val_change=None, + ignore_keys: Optional[set] = None, + ): + sanitized = {} + self._raise_value_error_on_nested_artifact(config_dict) + for k, v in config_dict.items(): + if ignore_keys and k in ignore_keys: + continue + k, v = self._sanitize(k, v, allow_val_change) + sanitized[k] = v + return sanitized + + def _sanitize(self, key, val, allow_val_change=None): + # TODO: enable WBValues in the config in the future + # refuse all WBValues which is all Media and Histograms + if isinstance(val, wandb.sdk.data_types.base_types.wb_value.WBValue): + raise TypeError("WBValue objects cannot be added to the run config") + # Let jupyter change config freely by default + if self._settings and self._settings._jupyter and allow_val_change is None: + allow_val_change = True + # We always normalize keys by stripping '-' + key = key.strip("-") + if _is_artifact_representation(val): + val = self._artifact_callback(key, val) + # if the user inserts an artifact into the config + if not isinstance(val, wandb.Artifact): + val = json_friendly_val(val) + if not allow_val_change: + if key in self._items and val != self._items[key]: + raise config_util.ConfigError( + f'Attempted to change value of key "{key}" ' + f"from {self._items[key]} to {val}\n" + "If you really want to do this, pass" + " allow_val_change=True to config.update()" + ) + return key, val + + def _raise_value_error_on_nested_artifact(self, v, nested=False): + # we can't swap nested artifacts because their root key can be locked by other values + # best if we don't allow nested artifacts until we can lock nested keys in the config + if isinstance(v, dict) and check_dict_contains_nested_artifact(v, nested): + raise ValueError( + "Instances of wandb.Artifact can only be top level keys in" + " a run's config" + ) + + +class ConfigStatic: + def __init__(self, config): + object.__setattr__(self, "__dict__", dict(config)) + + def __setattr__(self, name, value): + raise AttributeError("Error: run.config_static is a readonly object") + + def __setitem__(self, key, val): + raise AttributeError("Error: run.config_static is a readonly object") + + def keys(self): + return self.__dict__.keys() + + def __getitem__(self, key): + return self.__dict__[key] + + def __str__(self): + return str(self.__dict__) diff --git a/lib/python3.12/site-packages/wandb/sdk/wandb_helper.py b/lib/python3.12/site-packages/wandb/sdk/wandb_helper.py new file mode 100644 index 0000000000000000000000000000000000000000..5c5e5250f71b2b96db0e1300653ba1f94d62d682 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/wandb_helper.py @@ -0,0 +1,54 @@ +import inspect +import types + +from wandb.errors import UsageError + +from .lib import config_util + + +def parse_config(params, exclude=None, include=None): + if exclude and include: + raise UsageError("Expected at most only one of exclude or include") + if isinstance(params, str): + params = config_util.dict_from_config_file(params, must_exist=True) + params = _to_dict(params) + if include: + params = {key: value for key, value in params.items() if key in include} + if exclude: + params = {key: value for key, value in params.items() if key not in exclude} + return params + + +def _to_dict(params): + if isinstance(params, dict): + return params + + # Handle some cases where params is not a dictionary + # by trying to convert it into a dictionary + meta = inspect.getmodule(params) + if meta: + is_tf_flags_module = ( + isinstance(params, types.ModuleType) + and meta.__name__ == "tensorflow.python.platform.flags" + ) + if is_tf_flags_module or meta.__name__ == "absl.flags": + params = params.FLAGS + meta = inspect.getmodule(params) + + # newer tensorflow flags (post 1.4) uses absl.flags + if meta and meta.__name__ == "absl.flags._flagvalues": + params = {name: params[name].value for name in dir(params)} + elif not hasattr(params, "__dict__"): + raise TypeError("config must be a dict or have a __dict__ attribute.") + elif "__flags" in vars(params): + # for older tensorflow flags (pre 1.4) + if not "__parsed" not in vars(params): + params._parse_flags() + params = vars(params)["__flags"] + else: + # params is a Namespace object (argparse) + # or something else + params = vars(params) + + # assume argparse Namespace + return params diff --git a/lib/python3.12/site-packages/wandb/sdk/wandb_init.py b/lib/python3.12/site-packages/wandb/sdk/wandb_init.py new file mode 100644 index 0000000000000000000000000000000000000000..08753c5702f0d4a51a510305cf70b27565a59770 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/wandb_init.py @@ -0,0 +1,1620 @@ +"""Defines wandb.init() and associated classes and methods. + +`wandb.init()` indicates the beginning of a new run. In an ML training pipeline, +you could add `wandb.init()` to the beginning of your training script as well as +your evaluation script, and each step would be tracked as a run in W&B. + +For more on using `wandb.init()`, including code snippets, check out our +[guide and FAQs](https://docs.wandb.ai/guides/track/launch). +""" + +from __future__ import annotations + +import contextlib +import copy +import dataclasses +import json +import logging +import os +import pathlib +import platform +import sys +import tempfile +import time +from typing import TYPE_CHECKING, Iterable, Iterator, Sequence + +from typing_extensions import Any, Literal, Protocol, Self + +import wandb +import wandb.env +from wandb import env, trigger +from wandb.errors import CommError, Error, UsageError +from wandb.errors.links import url_registry +from wandb.errors.util import ProtobufErrorHandler +from wandb.integration import sagemaker +from wandb.proto.wandb_deprecated import Deprecated +from wandb.sdk.lib import ipython as wb_ipython +from wandb.sdk.lib import progress, runid, wb_logging +from wandb.sdk.lib.paths import StrPath +from wandb.util import _is_artifact_representation + +from . import wandb_login, wandb_setup +from .backend.backend import Backend +from .lib import SummaryDisabled, filesystem, module, paths, printer, telemetry +from .lib.deprecate import deprecate +from .mailbox import wait_with_progress +from .wandb_helper import parse_config +from .wandb_run import Run, TeardownHook, TeardownStage +from .wandb_settings import Settings + +if TYPE_CHECKING: + import wandb.jupyter + + +def _huggingface_version() -> str | None: + if "transformers" in sys.modules: + trans = wandb.util.get_module("transformers") + if hasattr(trans, "__version__"): + return str(trans.__version__) + return None + + +def _handle_launch_config(settings: Settings) -> dict[str, Any]: + launch_run_config: dict[str, Any] = {} + if not settings.launch: + return launch_run_config + if os.environ.get("WANDB_CONFIG") is not None: + try: + launch_run_config = json.loads(os.environ.get("WANDB_CONFIG", "{}")) + except (ValueError, SyntaxError): + wandb.termwarn("Malformed WANDB_CONFIG, using original config") + elif settings.launch_config_path and os.path.exists(settings.launch_config_path): + with open(settings.launch_config_path) as fp: + launch_config = json.loads(fp.read()) + launch_run_config = launch_config.get("overrides", {}).get("run_config") + else: + i = 0 + chunks = [] + while True: + key = f"WANDB_CONFIG_{i}" + if key in os.environ: + chunks.append(os.environ[key]) + i += 1 + else: + break + if len(chunks) > 0: + config_string = "".join(chunks) + try: + launch_run_config = json.loads(config_string) + except (ValueError, SyntaxError): + wandb.termwarn("Malformed WANDB_CONFIG, using original config") + + return launch_run_config + + +@dataclasses.dataclass(frozen=True) +class _ConfigParts: + base_no_artifacts: dict[str, Any] + """The run config passed to `init()` minus any artifact-valued keys.""" + + sweep_no_artifacts: dict[str, Any] + """The config loaded as part of a sweep minus any artifact-valued keys.""" + + launch_no_artifacts: dict[str, Any] + """The config loaded as part of Launch minus any artifact-valued keys.""" + + artifacts: dict[str, Any] + """Artifact keys removed from config dictionaries. + + Due to implementation details of how a Run is constructed, + artifacts must be inserted into its config after initialization. + """ + + +class _PrinterCallback(Protocol): + """A callback for displaying messages after a printer is configured. + + This is used for a few messages that may be generated before run settings + are computed, which are necessary for creating a printer. + """ + + def __call__(self, run_printer: printer.Printer) -> None: + """Display information through the given printer.""" + + +def _noop_printer_callback() -> _PrinterCallback: + """A printer callback that does not print anything.""" + return lambda _: None + + +def _concat_printer_callbacks( + cbs: Iterable[_PrinterCallback], +) -> _PrinterCallback: + """Returns a printer callback that runs the given callbacks in order.""" + + def do_callbacks(run_printer: printer.Printer) -> None: + for cb in cbs: + cb(run_printer) + + return do_callbacks + + +class _WandbInit: + def __init__( + self, + wl: wandb_setup._WandbSetup, + telemetry: telemetry.TelemetryRecord, + ) -> None: + self._wl = wl + + self._telemetry = telemetry + """Telemetry gathered before creating a run. + + After the run is created, `telemetry.context()` is used instead. + """ + + self.kwargs = None + self.run: Run | None = None + self.backend: Backend | None = None + + self._teardown_hooks: list[TeardownHook] = [] + self.notebook: wandb.jupyter.Notebook | None = None + + self.deprecated_features_used: dict[str, str] = dict() + + @property + def _logger(self) -> wandb_setup.Logger: + return self._wl._get_logger() + + def maybe_login(self, init_settings: Settings) -> None: + """Log in if we are not creating an offline or disabled run. + + This may change the W&B singleton settings. + + Args: + init_settings: Settings passed to `wandb.init()` or set via + keyword arguments. + """ + # Allow settings passed to init() to override inferred values. + # + # Calling login() may change settings on the singleton, + # so these may not be the final run settings. + run_settings = self._wl.settings.model_copy() + run_settings.update_from_settings(init_settings) + + # NOTE: _noop or _offline can become true after _login(). + # _noop happens if _login hits a timeout. + # _offline can be selected by the user at the login prompt. + if run_settings._noop or run_settings._offline: + return + + wandb_login._login( + anonymous=run_settings.anonymous, + host=run_settings.base_url, + force=run_settings.force, + _disable_warning=True, + _silent=run_settings.quiet or run_settings.silent, + ) + + def warn_env_vars_change_after_setup(self) -> _PrinterCallback: + """Warn if environment variables changed after `wandb.setup()`. + + Returns: + A callback to print any generated warnings. + """ + exclude_env_vars = {"WANDB_SERVICE", "WANDB_KUBEFLOW_URL"} + # check if environment variables have changed + singleton_env = { + k: v + for k, v in wandb_setup.singleton()._environ.items() + if k.startswith("WANDB_") and k not in exclude_env_vars + } + os_env = { + k: v + for k, v in os.environ.items() + if k.startswith("WANDB_") and k not in exclude_env_vars + } + + if ( + set(singleton_env.keys()) == set(os_env.keys()) # + and set(singleton_env.values()) == set(os_env.values()) + ): + return _noop_printer_callback() + + def print_warning(run_printer: printer.Printer) -> None: + line = ( + "Changes to your `wandb` environment variables will be ignored " + "because your `wandb` session has already started. " + "For more information on how to modify your settings with " + "`wandb.init()` arguments, please refer to " + f"{run_printer.link(url_registry.url('wandb-init'), 'the W&B docs')}." + ) + run_printer.display(line, level="warn") + + return print_warning + + def clear_run_path_if_sweep_or_launch( + self, + init_settings: Settings, + ) -> _PrinterCallback: + """Clear project/entity/run_id keys if in a Sweep or a Launch context. + + Args: + init_settings: Settings specified in the call to `wandb.init()`. + + Returns: + A callback to print any generated warnings. + """ + when_doing_thing = "" + + if self._wl.settings.sweep_id: + when_doing_thing = "when running a sweep" + elif self._wl.settings.launch: + when_doing_thing = "when running from a wandb launch context" + + if not when_doing_thing: + return _noop_printer_callback() + + warnings = [] + + def warn(key: str, value: str) -> None: + warnings.append(f"Ignoring {key} {value!r} {when_doing_thing}.") + + if init_settings.project is not None: + warn("project", init_settings.project) + init_settings.project = None + if init_settings.entity is not None: + warn("entity", init_settings.entity) + init_settings.entity = None + if init_settings.run_id is not None: + warn("run_id", init_settings.run_id) + init_settings.run_id = None + + def print_warnings(run_printer: printer.Printer) -> None: + for warning in warnings: + run_printer.display(warning, level="warn") + + return print_warnings + + def make_run_settings( + self, + init_settings: Settings, + ) -> tuple[Settings, _PrinterCallback]: + """Returns the run's settings and any warnings. + + Args: + init_settings: Settings passed to `wandb.init()` or set via + keyword arguments. + """ + warning_callbacks: list[_PrinterCallback] = [ + self.warn_env_vars_change_after_setup(), + self.clear_run_path_if_sweep_or_launch(init_settings), + ] + + # Inherit global settings. + settings = self._wl.settings.model_copy() + + # Apply settings from wandb.init() call. + settings.update_from_settings(init_settings) + + # Infer the run ID from SageMaker. + if not settings.sagemaker_disable and sagemaker.is_using_sagemaker(): + if sagemaker.set_run_id(settings): + self._logger.info("set run ID and group based on SageMaker") + self._telemetry.feature.sagemaker = True + + # get status of code saving before applying user settings + save_code_pre_user_settings = settings.save_code + if not settings._offline and not settings._noop: + user_settings = self._wl._load_user_settings() + if user_settings is not None: + settings.update_from_dict(user_settings) + + # ensure that user settings don't set saving to true + # if user explicitly set these to false in UI + if save_code_pre_user_settings is False: + settings.save_code = False + + # TODO: remove this once we refactor the client. This is a temporary + # fix to make sure that we use the same project name for wandb-core. + # The reason this is not going through the settings object is to + # avoid failure cases in other parts of the code that will be + # removed with the switch to wandb-core. + if settings.project is None: + settings.project = wandb.util.auto_project_name(settings.program) + + settings.x_start_time = time.time() + + # In shared mode, generate a unique label if not provided. + # The label is used to distinguish between system metrics and console logs + # from different writers to the same run. + if settings._shared and not settings.x_label: + # TODO: If executed in a known distributed environment (e.g. Ray or SLURM), + # use the env vars to generate a label (e.g. SLURM_JOB_ID or RANK) + prefix = settings.host or "" + label = runid.generate_id() + settings.x_label = f"{prefix}-{label}" if prefix else label + + return settings, _concat_printer_callbacks(warning_callbacks) + + def _load_autoresume_run_id(self, resume_file: pathlib.Path) -> str | None: + """Returns the run_id stored in the auto-resume file, if any. + + Returns None if the file does not exist or is not in a valid format. + + Args: + resume_file: The file path to use for resume='auto' mode. + """ + if not resume_file.exists(): + return None + + with resume_file.open() as f: + try: + return json.load(f)["run_id"] + + except json.JSONDecodeError as e: + self._logger.exception( + f"could not decode {resume_file}, ignoring", + exc_info=e, + ) + return None + + except KeyError: + self._logger.exception( + f"resume file at {resume_file} did not store a run_id" + ) + return None + + def _save_autoresume_run_id( + self, + *, + resume_file: pathlib.Path, + run_id: str, + ) -> None: + """Write the run ID to the auto-resume file.""" + resume_file.parent.mkdir(exist_ok=True) + with resume_file.open("w") as f: + json.dump({"run_id": run_id}, f) + + def set_run_id(self, settings: Settings) -> None: + """Set the run ID and possibly save it to the auto-resume file. + + After this, `settings.run_id` is guaranteed to be set. + + If a `resume_from` is provided and `run_id` is not set, initialize + `run_id` with the `resume_from` run's `run_id`. + + Args: + settings: The run's settings derived from the environment + and explicit values passed to `wandb.init()`. + """ + if settings.resume == "auto" and settings.resume_fname: + resume_path = pathlib.Path(settings.resume_fname) + else: + resume_path = None + + if resume_path: + previous_id = self._load_autoresume_run_id(resume_path) + + if not previous_id: + pass + elif settings.run_id is None: + self._logger.info(f"loaded run ID from {resume_path}") + settings.run_id = previous_id + elif settings.run_id != previous_id: + wandb.termwarn( + f"Ignoring ID {previous_id} loaded due to resume='auto'" + f" because the run ID is set to {settings.run_id}.", + ) + + # If no run ID was inferred, explicitly set, or loaded from an + # auto-resume file, then we generate a new ID. + if settings.run_id is None: + # If resume_from is provided and run_id is not already set, + # initialize run_id with the value from resume_from. + if settings.resume_from: + settings.run_id = settings.resume_from.run + else: + settings.run_id = runid.generate_id() + + if resume_path: + self._save_autoresume_run_id( + resume_file=resume_path, + run_id=settings.run_id, + ) + + def make_run_config( + self, + settings: Settings, + config: dict | str | None = None, + config_exclude_keys: list[str] | None = None, + config_include_keys: list[str] | None = None, + ) -> _ConfigParts: + """Construct the run's config. + + Args: + settings: The run's finalized settings. + config: The config passed to `init()`. + config_exclude_keys: Deprecated. Keys to filter out from `config`. + config_include_keys: Deprecated. Keys to include from `config`. + + Returns: + Initial values for the run's config. + """ + if config_exclude_keys: + self.deprecated_features_used["init__config_exclude_keys"] = ( + "config_exclude_keys is deprecated. Use" + " `config=wandb.helper.parse_config(config_object," + " exclude=('key',))` instead." + ) + if config_include_keys: + self.deprecated_features_used["init__config_include_keys"] = ( + "config_include_keys is deprecated. Use" + " `config=wandb.helper.parse_config(config_object," + " include=('key',))` instead." + ) + config = parse_config( + config or dict(), + include=config_include_keys, + exclude=config_exclude_keys, + ) + + result = _ConfigParts( + base_no_artifacts=dict(), + sweep_no_artifacts=dict(), + launch_no_artifacts=dict(), + artifacts=dict(), + ) + + if not settings.sagemaker_disable and sagemaker.is_using_sagemaker(): + sagemaker_config = sagemaker.parse_sm_config() + self._split_artifacts_from_config( + sagemaker_config, + config_target=result.base_no_artifacts, + artifacts=result.artifacts, + ) + self._telemetry.feature.sagemaker = True + + if self._wl._config: + self._split_artifacts_from_config( + self._wl._config, + config_target=result.base_no_artifacts, + artifacts=result.artifacts, + ) + + if config and isinstance(config, dict): + self._split_artifacts_from_config( + config, + config_target=result.base_no_artifacts, + artifacts=result.artifacts, + ) + + if self._wl._sweep_config: + self._split_artifacts_from_config( + self._wl._sweep_config, + config_target=result.sweep_no_artifacts, + artifacts=result.artifacts, + ) + + if launch_config := _handle_launch_config(settings): + self._split_artifacts_from_config( + launch_config, + config_target=result.launch_no_artifacts, + artifacts=result.artifacts, + ) + + wandb_internal = result.base_no_artifacts.setdefault("_wandb", dict()) + + if settings.save_code and settings.program_relpath: + wandb_internal["code_path"] = paths.LogicalPath( + os.path.join("code", settings.program_relpath) + ) + if settings.fork_from is not None: + wandb_internal["branch_point"] = { + "run_id": settings.fork_from.run, + "step": settings.fork_from.value, + } + if settings.resume_from is not None: + wandb_internal["branch_point"] = { + "run_id": settings.resume_from.run, + "step": settings.resume_from.value, + } + + return result + + def teardown(self) -> None: + # TODO: currently this is only called on failed wandb.init attempts + # normally this happens on the run object + self._logger.info("tearing down wandb.init") + for hook in self._teardown_hooks: + hook.call() + + def _split_artifacts_from_config( + self, + config_source: dict, + config_target: dict, + artifacts: dict, + ) -> None: + for k, v in config_source.items(): + if _is_artifact_representation(v): + artifacts[k] = v + else: + config_target.setdefault(k, v) + + def _safe_symlink( + self, base: str, target: str, name: str, delete: bool = False + ) -> None: + # TODO(jhr): do this with relpaths, but i can't figure it out on no sleep + if not hasattr(os, "symlink"): + return + + pid = os.getpid() + tmp_name = os.path.join(base, f"{name}.{pid}") + + if delete: + try: + os.remove(os.path.join(base, name)) + except OSError: + pass + target = os.path.relpath(target, base) + try: + os.symlink(target, tmp_name) + os.rename(tmp_name, os.path.join(base, name)) + except OSError: + pass + + def _pre_run_cell_hook(self, *args, **kwargs) -> None: + """Hook for the IPython pre_run_cell event. + + This pauses a run, preventing system metrics from being collected + the run's runtime from increasing. It also uploads the notebook's code. + """ + if not self.backend: + return + + if self.notebook and self.notebook.save_ipynb(): + assert self.run is not None + res = self.run.log_code(root=None) + self._logger.info("saved code: %s", res) + + if self.backend.interface is not None: + self._logger.info("pausing backend") + self.backend.interface.publish_pause() + + def _post_run_cell_hook(self, *args, **kwargs) -> None: + """Hook for the IPython post_run_cell event. + + Resumes collection of system metrics and the run's timer. + """ + if self.backend is None or self.backend.interface is None: + return + + self._logger.info("resuming backend") + self.backend.interface.publish_resume() + + def _jupyter_teardown(self) -> None: + """Teardown hooks and display saving, called with wandb.finish.""" + assert self.notebook + ipython = self.notebook.shell + + if self.run: + self.notebook.save_history(self.run) + + if self.notebook.save_ipynb(): + assert self.run is not None + res = self.run.log_code(root=None) + self._logger.info("saved code and history: %s", res) + self._logger.info("cleaning up jupyter logic") + + ipython.events.unregister("pre_run_cell", self._pre_run_cell_hook) + ipython.events.unregister("post_run_cell", self._post_run_cell_hook) + + ipython.display_pub.publish = ipython.display_pub._orig_publish + del ipython.display_pub._orig_publish + + def monkeypatch_ipython(self, settings: Settings) -> None: + """Add hooks, and session history saving.""" + self.notebook = wandb.jupyter.Notebook(settings) + ipython = self.notebook.shell + + # Monkey patch ipython publish to capture displayed outputs + if not hasattr(ipython.display_pub, "_orig_publish"): + self._logger.info("configuring jupyter hooks %s", self) + ipython.display_pub._orig_publish = ipython.display_pub.publish + + ipython.events.register("pre_run_cell", self._pre_run_cell_hook) + ipython.events.register("post_run_cell", self._post_run_cell_hook) + + self._teardown_hooks.append( + TeardownHook(self._jupyter_teardown, TeardownStage.EARLY) + ) + + def publish(data, metadata=None, **kwargs) -> None: + ipython.display_pub._orig_publish(data, metadata=metadata, **kwargs) + assert self.notebook is not None + self.notebook.save_display( + ipython.execution_count, {"data": data, "metadata": metadata} + ) + + ipython.display_pub.publish = publish + + @contextlib.contextmanager + def setup_run_log_directory(self, settings: Settings) -> Iterator[None]: + """Set up the run's log directory. + + This is a context manager that closes and unregisters the log handler + in case of an uncaught exception, so that future logged messages do not + modify this run's log file. + """ + filesystem.mkdir_exists_ok(os.path.dirname(settings.log_user)) + filesystem.mkdir_exists_ok(os.path.dirname(settings.log_internal)) + filesystem.mkdir_exists_ok(os.path.dirname(settings.sync_file)) + filesystem.mkdir_exists_ok(settings.files_dir) + filesystem.mkdir_exists_ok(settings._tmp_code_dir) + + if settings.symlink: + self._safe_symlink( + os.path.dirname(settings.sync_symlink_latest), + os.path.dirname(settings.sync_file), + os.path.basename(settings.sync_symlink_latest), + delete=True, + ) + self._safe_symlink( + os.path.dirname(settings.log_symlink_user), + settings.log_user, + os.path.basename(settings.log_symlink_user), + delete=True, + ) + self._safe_symlink( + os.path.dirname(settings.log_symlink_internal), + settings.log_internal, + os.path.basename(settings.log_symlink_internal), + delete=True, + ) + + assert settings.run_id + handler = wb_logging.add_file_handler( + settings.run_id, + pathlib.Path(settings.log_user), + ) + + if env.is_debug(): + handler.setLevel(logging.DEBUG) + + disposed = False + + def dispose_handler() -> None: + nonlocal disposed + + if not disposed: + disposed = True + logging.getLogger("wandb").removeHandler(handler) + handler.close() + + try: + self._teardown_hooks.append( + TeardownHook( + call=dispose_handler, + stage=TeardownStage.LATE, + ) + ) + + self._wl._early_logger_flush(logging.getLogger("wandb")) + self._logger.info(f"Logging user logs to {settings.log_user}") + self._logger.info(f"Logging internal logs to {settings.log_internal}") + + yield + except Exception: + dispose_handler() + raise + + def make_disabled_run(self, config: _ConfigParts) -> Run: + """Returns a Run-like object where all methods are no-ops. + + This method is used when the `mode` setting is set to "disabled", such as + by wandb.init(mode="disabled") or by setting the WANDB_MODE environment + variable to "disabled". + + It creates a Run object that mimics the behavior of a normal Run but doesn't + communicate with the W&B servers. + + The returned Run object has all expected attributes and methods, but they + are no-op versions that don't perform any actual logging or communication. + """ + run_id = runid.generate_id() + drun = Run( + settings=Settings( + mode="disabled", + x_files_dir=tempfile.gettempdir(), + run_id=run_id, + run_tags=tuple(), + run_notes=None, + run_group=None, + run_name=f"dummy-{run_id}", + project="dummy", + entity="dummy", + ) + ) + # config, summary, and metadata objects + drun._config = wandb.sdk.wandb_config.Config() + drun._config.update(config.sweep_no_artifacts) + drun._config.update(config.base_no_artifacts) + drun.summary = SummaryDisabled() # type: ignore + + # methods + drun.log = lambda data, *_, **__: drun.summary.update(data) # type: ignore[method-assign] + drun.finish = lambda *_, **__: module.unset_globals() # type: ignore[method-assign] + drun.join = drun.finish # type: ignore[method-assign] + drun.define_metric = lambda *_, **__: wandb.sdk.wandb_metric.Metric("dummy") # type: ignore[method-assign] + drun.save = lambda *_, **__: False # type: ignore[method-assign] + for symbol in ( + "alert", + "finish_artifact", + "get_project_url", + "get_sweep_url", + "get_url", + "link_artifact", + "link_model", + "use_artifact", + "log_code", + "log_model", + "use_model", + "mark_preempting", + "restore", + "status", + "watch", + "unwatch", + "upsert_artifact", + "_finish", + ): + setattr(drun, symbol, lambda *_, **__: None) # type: ignore + + # set properties to None + for attr in ("url", "project_url", "sweep_url"): + setattr(type(drun), attr, property(lambda _: None)) + + class _ChainableNoOp: + """An object that allows chaining arbitrary attributes and method calls.""" + + def __getattr__(self, _: str) -> Self: + return self + + def __call__(self, *_: Any, **__: Any) -> Self: + return self + + class _ChainableNoOpField: + # This is used to chain arbitrary attributes and method calls. + # For example, `run.log_artifact().state` will work in disabled mode. + def __init__(self) -> None: + self._value = None + + def __set__(self, instance: Any, value: Any) -> None: + self._value = value + + def __get__(self, instance: Any, owner: type) -> Any: + return _ChainableNoOp() if (self._value is None) else self._value + + def __call__(self, *args: Any, **kwargs: Any) -> _ChainableNoOp: + return _ChainableNoOp() + + drun.log_artifact = _ChainableNoOpField() # type: ignore + # attributes + drun._start_time = time.time() + drun._starting_step = 0 + drun._step = 0 + drun._attach_id = None + drun._backend = None + + # set the disabled run as the global run + module.set_global( + run=drun, + config=drun.config, + log=drun.log, + summary=drun.summary, + save=drun.save, + use_artifact=drun.use_artifact, + log_artifact=drun.log_artifact, + define_metric=drun.define_metric, + alert=drun.alert, + watch=drun.watch, + unwatch=drun.unwatch, + ) + return drun + + def init( # noqa: C901 + self, + settings: Settings, + config: _ConfigParts, + run_printer: printer.Printer, + ) -> Run: + self._logger.info("calling init triggers") + trigger.call("on_init") + + assert self._wl is not None + + self._logger.info( + f"wandb.init called with sweep_config: {config.sweep_no_artifacts}" + f"\nconfig: {config.base_no_artifacts}" + ) + + if previous_run := self._wl.most_recent_active_run: + if ( + settings.reinit in (True, "finish_previous") + # calling wandb.init() in notebooks finishes previous runs + # by default for user convenience. + or (settings.reinit == "default" and wb_ipython.in_notebook()) + ): + run_printer.display( + "Finishing previous runs because reinit is set" + f" to {settings.reinit!r}." + ) + self._wl.finish_all_active_runs() + + elif settings.reinit == "create_new": + self._logger.info( + "wandb.init() called while a run is active," + " and reinit is set to 'create_new', so continuing" + ) + + else: + run_printer.display( + "wandb.init() called while a run is active and reinit is" + f" set to {settings.reinit!r}, so returning the previous" + " run." + ) + + with telemetry.context(run=previous_run) as tel: + tel.feature.init_return_run = True + + return previous_run + + self._logger.info("starting backend") + + service = self._wl.ensure_service() + self._logger.info("sending inform_init request") + service.inform_init( + settings=settings.to_proto(), + run_id=settings.run_id, # type: ignore + ) + + backend = Backend(settings=settings, service=service) + backend.ensure_launched() + self._logger.info("backend started and connected") + + # resuming needs access to the server, check server_status()? + run = Run( + config=config.base_no_artifacts, + settings=settings, + sweep_config=config.sweep_no_artifacts, + launch_config=config.launch_no_artifacts, + ) + + # Populate initial telemetry + with telemetry.context(run=run, obj=self._telemetry) as tel: + tel.cli_version = wandb.__version__ + tel.python_version = platform.python_version() + tel.platform = f"{platform.system()}-{platform.machine()}".lower() + hf_version = _huggingface_version() + if hf_version: + tel.huggingface_version = hf_version + if settings._jupyter: + tel.env.jupyter = True + if settings._ipython: + tel.env.ipython = True + if settings._colab: + tel.env.colab = True + if settings._kaggle: + tel.env.kaggle = True + if settings._windows: + tel.env.windows = True + + if settings.launch: + tel.feature.launch = True + + for module_name in telemetry.list_telemetry_imports(only_imported=True): + setattr(tel.imports_init, module_name, True) + + if os.environ.get("PEX"): + tel.env.pex = True + + if settings._aws_lambda: + tel.env.aws_lambda = True + + if settings.x_flow_control_disabled: + tel.feature.flow_control_disabled = True + if settings.x_flow_control_custom: + tel.feature.flow_control_custom = True + if settings._shared: + wandb.termwarn( + "The `shared` mode feature is experimental and may change. " + "Please contact support@wandb.com for guidance and to report any issues." + ) + tel.feature.shared_mode = True + + if settings.x_label: + tel.feature.user_provided_label = True + + if wandb.env.dcgm_profiling_enabled(): + tel.feature.dcgm_profiling_enabled = True + + if not settings.label_disable: + if self.notebook: + run._label_probe_notebook(self.notebook) + else: + run._label_probe_main() + + for deprecated_feature, msg in self.deprecated_features_used.items(): + deprecate( + field_name=getattr(Deprecated, deprecated_feature), + warning_message=msg, + run=run, + ) + + self._logger.info("updated telemetry") + + run._set_library(self._wl) + run._set_backend(backend) + run._set_teardown_hooks(self._teardown_hooks) + + assert backend.interface + backend.interface.publish_header() + + # Using GitRepo() blocks & can be slow, depending on user's current git setup. + # We don't want to block run initialization/start request, so populate run's git + # info beforehand. + if not (settings.disable_git or settings.x_disable_machine_info): + run._populate_git_info() + + if settings._offline and settings.resume: + wandb.termwarn( + "`resume` will be ignored since W&B syncing is set to `offline`. " + f"Starting a new run with run id {run.id}." + ) + error: wandb.Error | None = None + + timeout = settings.init_timeout + + self._logger.info( + f"communicating run to backend with {timeout} second timeout", + ) + + run_init_handle = backend.interface.deliver_run(run) + + async def display_init_message() -> None: + assert backend.interface + + with progress.progress_printer( + run_printer, + default_text="Waiting for wandb.init()...", + ) as progress_printer: + await progress.loop_printing_operation_stats( + progress_printer, + backend.interface, + ) + + try: + result = wait_with_progress( + run_init_handle, + timeout=timeout, + progress_after=1, + display_progress=display_init_message, + ) + + except TimeoutError: + run_init_handle.cancel(backend.interface) + + # This may either be an issue with the W&B server (a CommError) + # or a bug in the SDK (an Error). We cannot distinguish between + # the two causes here. + raise CommError( + f"Run initialization has timed out after {timeout} sec." + " Please try increasing the timeout with the `init_timeout`" + " setting: `wandb.init(settings=wandb.Settings(init_timeout=120))`." + ) + + assert result.run_result + + if error := ProtobufErrorHandler.to_exception(result.run_result.error): + raise error + + if not result.run_result.HasField("run"): + raise Error("Assertion failed: run_result is missing the run field") + + if result.run_result.run.resumed: + self._logger.info("run resumed") + with telemetry.context(run=run) as tel: + tel.feature.resumed = result.run_result.run.resumed + run._set_run_obj(result.run_result.run) + + self._logger.info("starting run threads in backend") + # initiate run (stats and metadata probing) + + if service: + assert settings.run_id + service.inform_start( + settings=settings.to_proto(), + run_id=settings.run_id, + ) + + assert backend.interface + + run_start_handle = backend.interface.deliver_run_start(run) + try: + # TODO: add progress to let user know we are doing something + run_start_handle.wait_or(timeout=30) + except TimeoutError: + pass + + assert self._wl is not None + self.run = run + + run._handle_launch_artifact_overrides() + if ( + settings.launch + and settings.launch_config_path + and os.path.exists(settings.launch_config_path) + ): + run.save(settings.launch_config_path) + # put artifacts in run config here + # since doing so earlier will cause an error + # as the run is not upserted + for k, v in config.artifacts.items(): + run.config.update({k: v}, allow_val_change=True) + job_artifact = run._launch_artifact_mapping.get( + wandb.util.LAUNCH_JOB_ARTIFACT_SLOT_NAME + ) + if job_artifact: + run.use_artifact(job_artifact) + + self.backend = backend + + if settings.reinit != "create_new": + _set_global_run(run) + + run._on_start() + self._logger.info("run started, returning control to user process") + return run + + +def _attach( + attach_id: str | None = None, + run_id: str | None = None, + *, + run: Run | None = None, +) -> Run | None: + """Attach to a run currently executing in another process/thread. + + Args: + attach_id: (str, optional) The id of the run or an attach identifier + that maps to a run. + run_id: (str, optional) The id of the run to attach to. + run: (Run, optional) The run instance to attach + """ + attach_id = attach_id or run_id + if not ((attach_id is None) ^ (run is None)): + raise UsageError("Either (`attach_id` or `run_id`) or `run` must be specified") + + attach_id = attach_id or (run._attach_id if run else None) + + if attach_id is None: + raise UsageError( + "Either `attach_id` or `run_id` must be specified or `run` must have `_attach_id`" + ) + wandb._assert_is_user_process() # type: ignore + + _wl = wandb_setup.singleton() + logger = _wl._get_logger() + + service = _wl.ensure_service() + + try: + attach_settings = service.inform_attach(attach_id=attach_id) + except Exception as e: + raise UsageError(f"Unable to attach to run {attach_id}") from e + + settings: Settings = copy.copy(_wl._settings) + + settings.update_from_dict( + { + "run_id": attach_id, + "x_start_time": attach_settings.x_start_time.value, + "mode": attach_settings.mode.value, + } + ) + + # TODO: consolidate this codepath with wandb.init() + backend = Backend(settings=settings, service=service) + backend.ensure_launched() + logger.info("attach backend started and connected") + + if run is None: + run = Run(settings=settings) + else: + run._init(settings=settings) + run._set_library(_wl) + run._set_backend(backend) + assert backend.interface + + attach_handle = backend.interface.deliver_attach(attach_id) + try: + # TODO: add progress to let user know we are doing something + attach_result = attach_handle.wait_or(timeout=30) + except TimeoutError: + raise UsageError("Timeout attaching to run") + + attach_response = attach_result.response.attach_response + if attach_response.error and attach_response.error.message: + raise UsageError(f"Failed to attach to run: {attach_response.error.message}") + + run._set_run_obj(attach_response.run) + _set_global_run(run) + run._on_attach() + return run + + +def _set_global_run(run: Run) -> None: + """Set `wandb.run` and point some top-level functions to its methods. + + Args: + run: The run to make global. + """ + module.set_global( + run=run, + config=run.config, + log=run.log, + summary=run.summary, + save=run.save, + use_artifact=run.use_artifact, + log_artifact=run.log_artifact, + define_metric=run.define_metric, + alert=run.alert, + watch=run.watch, + unwatch=run.unwatch, + mark_preempting=run.mark_preempting, + log_model=run.log_model, + use_model=run.use_model, + link_model=run.link_model, + ) + + +def _monkeypatch_openai_gym() -> None: + """Patch OpenAI gym to log to the global `wandb.run`.""" + if len(wandb.patched["gym"]) > 0: + return + + from wandb.integration import gym + + gym.monitor() + + +def _monkeypatch_tensorboard() -> None: + """Patch TensorBoard to log to the global `wandb.run`.""" + if len(wandb.patched["tensorboard"]) > 0: + return + + from wandb.integration import tensorboard as tb_module + + tb_module.patch() + + +def try_create_root_dir(settings: Settings) -> None: + """Try to create the root directory specified in settings. + + If creation fails due to permissions or other errors, + falls back to using the system temp directory. + + Args: + settings: The runs settings containing root_dir configuration. + This function may update the root_dir to a temporary directory + if the parent directory is not writable. + """ + fallback_to_temp_dir = False + + try: + os.makedirs(settings.root_dir, exist_ok=True) + except OSError: + wandb.termwarn( + f"Unable to create root directory {settings.root_dir}", + repeat=False, + ) + fallback_to_temp_dir = True + else: + if not os.access(settings.root_dir, os.W_OK | os.R_OK): + wandb.termwarn( + f"Path {settings.root_dir} wasn't read/writable", + repeat=False, + ) + fallback_to_temp_dir = True + + if not fallback_to_temp_dir: + return + + tmp_dir = tempfile.gettempdir() + if not os.access(tmp_dir, os.W_OK | os.R_OK): + raise ValueError( + f"System temp directory ({tmp_dir}) is not writable/readable, " + "please set the `dir` argument in `wandb.init()` to a writable/readable directory." + ) + + settings.root_dir = tmp_dir + wandb.termwarn( + f"Falling back to temporary directory {tmp_dir}.", + repeat=False, + ) + os.makedirs(settings.root_dir, exist_ok=True) + + +def init( # noqa: C901 + entity: str | None = None, + project: str | None = None, + dir: StrPath | None = None, + id: str | None = None, + name: str | None = None, + notes: str | None = None, + tags: Sequence[str] | None = None, + config: dict[str, Any] | str | None = None, + config_exclude_keys: list[str] | None = None, + config_include_keys: list[str] | None = None, + allow_val_change: bool | None = None, + group: str | None = None, + job_type: str | None = None, + mode: Literal["online", "offline", "disabled"] | None = None, + force: bool | None = None, + anonymous: Literal["never", "allow", "must"] | None = None, + reinit: ( + bool + | Literal[ + None, + "default", + "return_previous", + "finish_previous", + "create_new", + ] + ) = None, + resume: bool | Literal["allow", "never", "must", "auto"] | None = None, + resume_from: str | None = None, + fork_from: str | None = None, + save_code: bool | None = None, + tensorboard: bool | None = None, + sync_tensorboard: bool | None = None, + monitor_gym: bool | None = None, + settings: Settings | dict[str, Any] | None = None, +) -> Run: + r"""Start a new run to track and log to W&B. + + In an ML training pipeline, you could add `wandb.init()` to the beginning of + your training script as well as your evaluation script, and each piece would + be tracked as a run in W&B. + + `wandb.init()` spawns a new background process to log data to a run, and it + also syncs data to https://wandb.ai by default, so you can see your results + in real-time. + + Call `wandb.init()` to start a run before logging data with `wandb.log()`. + When you're done logging data, call `wandb.finish()` to end the run. If you + don't call `wandb.finish()`, the run will end when your script exits. + + For more on using `wandb.init()`, including detailed examples, check out our + [guide and FAQs](https://docs.wandb.ai/guides/track/launch). + + Examples: + ### Explicitly set the entity and project and choose a name for the run: + + ```python + import wandb + + run = wandb.init( + entity="geoff", + project="capsules", + name="experiment-2021-10-31", + ) + + # ... your training code here ... + + run.finish() + ``` + + ### Add metadata about the run using the `config` argument: + + ```python + import wandb + + config = {"lr": 0.01, "batch_size": 32} + with wandb.init(config=config) as run: + run.config.update({"architecture": "resnet", "depth": 34}) + + # ... your training code here ... + ``` + + Note that you can use `wandb.init()` as a context manager to automatically + call `wandb.finish()` at the end of the block. + + Args: + entity: The username or team name under which the runs will be logged. + The entity must already exist, so ensure you’ve created your account + or team in the UI before starting to log runs. If not specified, the + run will default your default entity. To change the default entity, + go to [your settings](https://wandb.ai/settings) and update the + "Default location to create new projects" under "Default team". + project: The name of the project under which this run will be logged. + If not specified, we use a heuristic to infer the project name based + on the system, such as checking the git root or the current program + file. If we can't infer the project name, the project will default to + `"uncategorized"`. + dir: The absolute path to the directory where experiment logs and + metadata files are stored. If not specified, this defaults + to the `./wandb` directory. Note that this does not affect the + location where artifacts are stored when calling `download()`. + id: A unique identifier for this run, used for resuming. It must be unique + within the project and cannot be reused once a run is deleted. The + identifier must not contain any of the following special characters: + `/ \ # ? % :`. For a short descriptive name, use the `name` field, + or for saving hyperparameters to compare across runs, use `config`. + name: A short display name for this run, which appears in the UI to help + you identify it. By default, we generate a random two-word name + allowing easy cross-reference runs from table to charts. Keeping these + run names brief enhances readability in chart legends and tables. For + saving hyperparameters, we recommend using the `config` field. + notes: A detailed description of the run, similar to a commit message in + Git. Use this argument to capture any context or details that may + help you recall the purpose or setup of this run in the future. + tags: A list of tags to label this run in the UI. Tags are helpful for + organizing runs or adding temporary identifiers like "baseline" or + "production." You can easily add, remove tags, or filter by tags in + the UI. + If resuming a run, the tags provided here will replace any existing + tags. To add tags to a resumed run without overwriting the current + tags, use `run.tags += ("new_tag",)` after calling `run = wandb.init()`. + config: Sets `wandb.config`, a dictionary-like object for storing input + parameters to your run, such as model hyperparameters or data + preprocessing settings. + The config appears in the UI in an overview page, allowing you to + group, filter, and sort runs based on these parameters. + Keys should not contain periods (`.`), and values should be + smaller than 10 MB. + If a dictionary, `argparse.Namespace`, or `absl.flags.FLAGS` is + provided, the key-value pairs will be loaded directly into + `wandb.config`. + If a string is provided, it is interpreted as a path to a YAML file, + from which configuration values will be loaded into `wandb.config`. + config_exclude_keys: A list of specific keys to exclude from `wandb.config`. + config_include_keys: A list of specific keys to include in `wandb.config`. + allow_val_change: Controls whether config values can be modified after their + initial set. By default, an exception is raised if a config value is + overwritten. For tracking variables that change during training, such as + a learning rate, consider using `wandb.log()` instead. By default, this + is `False` in scripts and `True` in Notebook environments. + group: Specify a group name to organize individual runs as part of a larger + experiment. This is useful for cases like cross-validation or running + multiple jobs that train and evaluate a model on different test sets. + Grouping allows you to manage related runs collectively in the UI, + making it easy to toggle and review results as a unified experiment. + For more information, refer to our + [guide to grouping runs](https://docs.wandb.com/guides/runs/grouping). + job_type: Specify the type of run, especially helpful when organizing runs + within a group as part of a larger experiment. For example, in a group, + you might label runs with job types such as "train" and "eval". + Defining job types enables you to easily filter and group similar runs + in the UI, facilitating direct comparisons. + mode: Specifies how run data is managed, with the following options: + - `"online"` (default): Enables live syncing with W&B when a network + connection is available, with real-time updates to visualizations. + - `"offline"`: Suitable for air-gapped or offline environments; data + is saved locally and can be synced later. Ensure the run folder + is preserved to enable future syncing. + - `"disabled"`: Disables all W&B functionality, making the run’s methods + no-ops. Typically used in testing to bypass W&B operations. + force: Determines if a W&B login is required to run the script. If `True`, + the user must be logged in to W&B; otherwise, the script will not + proceed. If `False` (default), the script can proceed without a login, + switching to offline mode if the user is not logged in. + anonymous: Specifies the level of control over anonymous data logging. + Available options are: + - `"never"` (default): Requires you to link your W&B account before + tracking the run. This prevents unintentional creation of anonymous + runs by ensuring each run is associated with an account. + - `"allow"`: Enables a logged-in user to track runs with their account, + but also allows someone running the script without a W&B account + to view the charts and data in the UI. + - `"must"`: Forces the run to be logged to an anonymous account, even + if the user is logged in. + reinit: Shorthand for the "reinit" setting. Determines the behavior of + `wandb.init()` when a run is active. + resume: Controls the behavior when resuming a run with the specified `id`. + Available options are: + - `"allow"`: If a run with the specified `id` exists, it will resume + from the last step; otherwise, a new run will be created. + - `"never"`: If a run with the specified `id` exists, an error will + be raised. If no such run is found, a new run will be created. + - `"must"`: If a run with the specified `id` exists, it will resume + from the last step. If no run is found, an error will be raised. + - `"auto"`: Automatically resumes the previous run if it crashed on + this machine; otherwise, starts a new run. + - `True`: Deprecated. Use `"auto"` instead. + - `False`: Deprecated. Use the default behavior (leaving `resume` + unset) to always start a new run. + Note: If `resume` is set, `fork_from` and `resume_from` cannot be + used. When `resume` is unset, the system will always start a new run. + For more details, see our + [guide to resuming runs](https://docs.wandb.com/guides/runs/resuming). + resume_from: Specifies a moment in a previous run to resume a run from, + using the format `{run_id}?_step={step}`. This allows users to truncate + the history logged to a run at an intermediate step and resume logging + from that step. The target run must be in the same project. + If an `id` argument is also provided, the `resume_from` argument will + take precedence. + `resume`, `resume_from` and `fork_from` cannot be used together, only + one of them can be used at a time. + Note: This feature is in beta and may change in the future. + fork_from: Specifies a point in a previous run from which to fork a new + run, using the format `{id}?_step={step}`. This creates a new run that + resumes logging from the specified step in the target run’s history. + The target run must be part of the current project. + If an `id` argument is also provided, it must be different from the + `fork_from` argument, an error will be raised if they are the same. + `resume`, `resume_from` and `fork_from` cannot be used together, only + one of them can be used at a time. + Note: This feature is in beta and may change in the future. + save_code: Enables saving the main script or notebook to W&B, aiding in + experiment reproducibility and allowing code comparisons across runs in + the UI. By default, this is disabled, but you can change the default to + enable on your [settings page](https://wandb.ai/settings). + tensorboard: Deprecated. Use `sync_tensorboard` instead. + sync_tensorboard: Enables automatic syncing of W&B logs from TensorBoard + or TensorBoardX, saving relevant event files for viewing in the W&B UI. + saving relevant event files for viewing in the W&B UI. (Default: `False`) + monitor_gym: Enables automatic logging of videos of the environment when + using OpenAI Gym. For additional details, see our + [guide for gym integration](https://docs.wandb.com/guides/integrations/openai-gym). + settings: Specifies a dictionary or `wandb.Settings` object with advanced + settings for the run. + + Returns: + A `Run` object, which is a handle to the current run. Use this object + to perform operations like logging data, saving files, and finishing + the run. See the [Run API](https://docs.wandb.ai/ref/python/run) for + more details. + + Raises: + Error: If some unknown or internal error happened during the run + initialization. + AuthenticationError: If the user failed to provide valid credentials. + CommError: If there was a problem communicating with the W&B server. + UsageError: If the user provided invalid arguments to the function. + KeyboardInterrupt: If the user interrupts the run initialization process. + If the user interrupts the run initialization process. + """ + wandb._assert_is_user_process() # type: ignore + + init_telemetry = telemetry.TelemetryRecord() + + init_settings = Settings() + if isinstance(settings, dict): + init_settings = Settings(**settings) + elif isinstance(settings, Settings): + init_settings = settings + + # Explicit function arguments take precedence over settings + if job_type is not None: + init_settings.run_job_type = job_type + if dir is not None: + init_settings.root_dir = dir # type: ignore + if project is not None: + init_settings.project = project + if entity is not None: + init_settings.entity = entity + if reinit is not None: + init_settings.reinit = reinit + if tags is not None: + init_settings.run_tags = tuple(tags) + if group is not None: + init_settings.run_group = group + if name is not None: + init_settings.run_name = name + if notes is not None: + init_settings.run_notes = notes + if anonymous is not None: + init_settings.anonymous = anonymous # type: ignore + if mode is not None: + init_settings.mode = mode # type: ignore + if resume is not None: + init_settings.resume = resume # type: ignore + if force is not None: + init_settings.force = force + # TODO: deprecate "tensorboard" in favor of "sync_tensorboard" + if tensorboard is not None: + init_settings.sync_tensorboard = tensorboard + if sync_tensorboard is not None: + init_settings.sync_tensorboard = sync_tensorboard + if save_code is not None: + init_settings.save_code = save_code + if id is not None: + init_settings.run_id = id + if fork_from is not None: + init_settings.fork_from = fork_from # type: ignore + if resume_from is not None: + init_settings.resume_from = resume_from # type: ignore + + if config is not None: + init_telemetry.feature.set_init_config = True + + wl: wandb_setup._WandbSetup | None = None + + try: + wl = wandb_setup.singleton() + + wi = _WandbInit(wl, init_telemetry) + + wi.maybe_login(init_settings) + run_settings, show_warnings = wi.make_run_settings(init_settings) + + if isinstance(run_settings.reinit, bool): + wi.deprecated_features_used["run__reinit_bool"] = ( + "Using a boolean value for 'reinit' is deprecated." + " Use 'return_previous' or 'finish_previous' instead." + ) + + if run_settings.run_id is not None: + init_telemetry.feature.set_init_id = True + if run_settings.run_name is not None: + init_telemetry.feature.set_init_name = True + if run_settings.run_tags is not None: + init_telemetry.feature.set_init_tags = True + if run_settings._offline: + init_telemetry.feature.offline = True + if run_settings.fork_from is not None: + init_telemetry.feature.fork_mode = True + if run_settings.resume_from is not None: + init_telemetry.feature.rewind_mode = True + + wi.set_run_id(run_settings) + run_printer = printer.new_printer(run_settings) + show_warnings(run_printer) + + with contextlib.ExitStack() as exit_stack: + exit_stack.enter_context(wb_logging.log_to_run(run_settings.run_id)) + + run_config = wi.make_run_config( + settings=run_settings, + config=config, + config_exclude_keys=config_exclude_keys, + config_include_keys=config_include_keys, + ) + + if run_settings._noop: + return wi.make_disabled_run(run_config) + + try_create_root_dir(run_settings) + exit_stack.enter_context(wi.setup_run_log_directory(run_settings)) + + if run_settings._jupyter: + wi.monkeypatch_ipython(run_settings) + + if monitor_gym: + _monkeypatch_openai_gym() + + if wandb.patched["tensorboard"]: + # NOTE: The user may have called the patch function directly. + init_telemetry.feature.tensorboard_patch = True + if run_settings.sync_tensorboard: + _monkeypatch_tensorboard() + init_telemetry.feature.tensorboard_sync = True + + if run_settings.x_server_side_derived_summary: + init_telemetry.feature.server_side_derived_summary = True + + return wi.init(run_settings, run_config, run_printer) + + except KeyboardInterrupt as e: + if wl: + wl._get_logger().warning("interrupted", exc_info=e) + + raise + + except Exception as e: + if wl: + wl._get_logger().exception("error in wandb.init()", exc_info=e) + + # Need to build delay into this sentry capture because our exit hooks + # mess with sentry's ability to send out errors before the program ends. + wandb._sentry.reraise(e) diff --git a/lib/python3.12/site-packages/wandb/sdk/wandb_login.py b/lib/python3.12/site-packages/wandb/sdk/wandb_login.py new file mode 100644 index 0000000000000000000000000000000000000000..86ebfaee69a8c9814c663fee31290c3886eab159 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/wandb_login.py @@ -0,0 +1,328 @@ +"""Log in to Weights & Biases. + +This authenticates your machine to log data to your account. +""" + +import enum +import os +from typing import Literal, Optional, Tuple + +import click +from requests.exceptions import ConnectionError + +import wandb +from wandb.errors import AuthenticationError, UsageError +from wandb.old.settings import Settings as OldSettings +from wandb.sdk import wandb_setup + +from ..apis import InternalApi +from .internal.internal_api import Api +from .lib import apikey + + +def _handle_host_wandb_setting(host: Optional[str], cloud: bool = False) -> None: + """Write the host parameter to the global settings file. + + This takes the parameter from wandb.login or wandb login for use by the + application's APIs. + """ + _api = InternalApi() + if host == "https://api.wandb.ai" or (host is None and cloud): + _api.clear_setting("base_url", globally=True, persist=True) + # To avoid writing an empty local settings file, we only clear if it exists + if os.path.exists(OldSettings._local_path()): + _api.clear_setting("base_url", persist=True) + elif host: + host = host.rstrip("/") + # force relogin if host is specified + _api.set_setting("base_url", host, globally=True, persist=True) + + +def login( + anonymous: Optional[Literal["must", "allow", "never"]] = None, + key: Optional[str] = None, + relogin: Optional[bool] = None, + host: Optional[str] = None, + force: Optional[bool] = None, + timeout: Optional[int] = None, + verify: bool = False, + referrer: Optional[str] = None, +) -> bool: + """Set up W&B login credentials. + + By default, this will only store credentials locally without + verifying them with the W&B server. To verify credentials, pass + `verify=True`. + + Args: + anonymous: (string, optional) Can be "must", "allow", or "never". + If set to "must", always log a user in anonymously. If set to + "allow", only create an anonymous user if the user + isn't already logged in. If set to "never", never log a + user anonymously. Default set to "never". + key: (string, optional) The API key to use. + relogin: (bool, optional) If true, will re-prompt for API key. + host: (string, optional) The host to connect to. + force: (bool, optional) If true, will force a relogin. + timeout: (int, optional) Number of seconds to wait for user input. + verify: (bool) Verify the credentials with the W&B server. + referrer: (string, optional) The referrer to use in the URL login request. + + Returns: + bool: if key is configured + + Raises: + AuthenticationError - if api_key fails verification with the server + UsageError - if api_key cannot be configured and no tty + """ + _handle_host_wandb_setting(host) + return _login( + anonymous=anonymous, + key=key, + relogin=relogin, + host=host, + force=force, + timeout=timeout, + verify=verify, + referrer=referrer, + ) + + +class ApiKeyStatus(enum.Enum): + VALID = 1 + NOTTY = 2 + OFFLINE = 3 + DISABLED = 4 + + +class _WandbLogin: + def __init__( + self, + anonymous: Optional[Literal["must", "allow", "never"]] = None, + force: Optional[bool] = None, + host: Optional[str] = None, + key: Optional[str] = None, + relogin: Optional[bool] = None, + timeout: Optional[int] = None, + ): + self._relogin = relogin + + login_settings = { + "anonymous": anonymous, + "api_key": key, + "base_url": host, + "force": force, + "login_timeout": timeout, + } + self.is_anonymous = anonymous == "must" + + self._wandb_setup = wandb_setup.singleton() + self._wandb_setup.settings.update_from_dict(login_settings) + self._settings = self._wandb_setup.settings + + def _update_global_anonymous_setting(self) -> None: + api = InternalApi() + if self.is_anonymous: + api.set_setting("anonymous", "must", globally=True, persist=True) + else: + api.clear_setting("anonymous", globally=True, persist=True) + + def is_apikey_configured(self) -> bool: + """Returns whether an API key is set or can be inferred.""" + return apikey.api_key(settings=self._settings) is not None + + def _print_logged_in_message(self) -> None: + """Prints a message telling the user they are logged in.""" + username = self._wandb_setup._get_username() + + if username: + host_str = ( + f" to {click.style(self._settings.base_url, fg='green')}" + if self._settings.base_url + else "" + ) + + # check to see if we got an entity from the setup call or from the user + entity = self._settings.entity or self._wandb_setup._get_entity() + + entity_str = "" + # check if entity exist, valid (is part of a certain team) and different from the username + if ( + entity + and entity in self._wandb_setup._get_teams() + and entity != username + ): + entity_str = f" ({click.style(entity, fg='yellow')})" + + login_state_str = f"Currently logged in as: {click.style(username, fg='yellow')}{entity_str}{host_str}" + else: + login_state_str = "W&B API key is configured" + + login_info_str = ( + f"Use {click.style('`wandb login --relogin`', bold=True)} to force relogin" + ) + wandb.termlog( + f"{login_state_str}. {login_info_str}", + repeat=False, + ) + + def try_save_api_key(self, key: str) -> None: + """Saves the API key to disk for future use.""" + if self._settings._notebook and not self._settings.silent: + wandb.termwarn( + "If you're specifying your api key in code, ensure this " + "code is not shared publicly.\nConsider setting the " + "WANDB_API_KEY environment variable, or running " + "`wandb login` from the command line." + ) + if key: + try: + apikey.write_key(self._settings, key) + except apikey.WriteNetrcError as e: + wandb.termwarn(str(e)) + + def update_session( + self, + key: Optional[str], + status: ApiKeyStatus = ApiKeyStatus.VALID, + ) -> None: + """Updates mode and API key settings on the global setup object. + + If we're online, this also pulls in user settings from the server. + """ + login_settings = dict() + if status == ApiKeyStatus.OFFLINE: + login_settings = dict(mode="offline") + elif status == ApiKeyStatus.DISABLED: + login_settings = dict(mode="disabled") + elif key: + login_settings = dict(api_key=key) + self._wandb_setup.settings.update_from_dict(login_settings) + # Whenever the key changes, make sure to pull in user settings + # from server. + if not self._wandb_setup.settings._offline: + self._wandb_setup._update_user_settings() + + def _prompt_api_key( + self, referrer: Optional[str] = None + ) -> Tuple[Optional[str], ApiKeyStatus]: + api = Api(self._settings) + while True: + try: + key = apikey.prompt_api_key( + self._settings, + api=api, + no_offline=self._settings.force if self._settings else None, + no_create=self._settings.force if self._settings else None, + referrer=referrer, + ) + except ValueError as e: + # invalid key provided, try again + wandb.termerror(e.args[0]) + continue + except TimeoutError: + wandb.termlog("W&B disabled due to login timeout.") + return None, ApiKeyStatus.DISABLED + if key is False: + return None, ApiKeyStatus.NOTTY + if not key: + return None, ApiKeyStatus.OFFLINE + return key, ApiKeyStatus.VALID + + def prompt_api_key( + self, referrer: Optional[str] = None + ) -> Tuple[Optional[str], ApiKeyStatus]: + """Updates the global API key by prompting the user.""" + key, status = self._prompt_api_key(referrer) + if status == ApiKeyStatus.NOTTY: + directive = ( + "wandb login [your_api_key]" + if self._settings.x_cli_only_mode + else "wandb.login(key=[your_api_key])" + ) + raise UsageError("api_key not configured (no-tty). call " + directive) + + return key, status + + def _verify_login(self, key: str) -> None: + api = InternalApi(api_key=key) + + try: + is_api_key_valid = api.validate_api_key() + except ConnectionError: + raise AuthenticationError( + "Unable to connect to server to verify API token." + ) + except Exception: + raise AuthenticationError("An error occurred while verifying the API key.") + + if not is_api_key_valid: + raise AuthenticationError( + "API key verification failed. Make sure your API key is valid." + ) + + +def _login( + *, + anonymous: Optional[Literal["allow", "must", "never"]] = None, + key: Optional[str] = None, + relogin: Optional[bool] = None, + host: Optional[str] = None, + force: Optional[bool] = None, + timeout: Optional[int] = None, + verify: bool = False, + referrer: str = "models", + _silent: Optional[bool] = None, + _disable_warning: Optional[bool] = None, +) -> bool: + if wandb.run is not None: + if not _disable_warning: + wandb.termwarn("Calling wandb.login() after wandb.init() has no effect.") + return True + + wlogin = _WandbLogin( + anonymous=anonymous, + force=force, + host=host, + key=key, + relogin=relogin, + timeout=timeout, + ) + + if wlogin._settings._noop: + return True + + if wlogin._settings._offline and not wlogin._settings.x_cli_only_mode: + wandb.termwarn("Unable to verify login in offline mode.") + return False + elif wandb.util._is_kaggle() and not wandb.util._has_internet(): + wandb.termerror( + "To use W&B in kaggle you must enable internet in the settings panel on the right." + ) + return False + + if wlogin._settings.identity_token_file: + return True + + key_is_pre_configured = False + key_status = None + if key is None: + # Check if key is already set in the settings, or configured in the users .netrc file. + key = apikey.api_key(settings=wlogin._settings) + if key and not relogin: + key_is_pre_configured = True + else: + key, key_status = wlogin.prompt_api_key(referrer=referrer) + + if verify: + wlogin._verify_login(key) + + if not key_is_pre_configured: + wlogin.try_save_api_key(key) + wlogin.update_session(key, status=key_status) + wlogin._update_global_anonymous_setting() + + if key and not _silent: + wlogin._print_logged_in_message() + + return key is not None diff --git a/lib/python3.12/site-packages/wandb/sdk/wandb_metric.py b/lib/python3.12/site-packages/wandb/sdk/wandb_metric.py new file mode 100644 index 0000000000000000000000000000000000000000..f01e4b7a30bbbc070a2a00b174d7c2f9f4fa27e0 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/wandb_metric.py @@ -0,0 +1,110 @@ +"""metric.""" + +import logging +from typing import Callable, Optional, Sequence, Tuple + +from wandb.proto import wandb_internal_pb2 as pb + +logger = logging.getLogger("wandb") + + +class Metric: + """Metric object.""" + + _callback: Optional[Callable[[pb.MetricRecord], None]] + _name: str + _step_metric: Optional[str] + _step_sync: Optional[bool] + _hidden: Optional[bool] + _summary: Optional[Sequence[str]] + _goal: Optional[str] + _overwrite: Optional[bool] + + def __init__( + self, + name: str, + step_metric: Optional[str] = None, + step_sync: Optional[bool] = None, + hidden: Optional[bool] = None, + summary: Optional[Sequence[str]] = None, + goal: Optional[str] = None, + overwrite: Optional[bool] = None, + ) -> None: + self._callback = None + self._name = name + self._step_metric = step_metric + # default to step_sync=True if step metric is set + step_sync = step_sync if step_sync is not None else step_metric is not None + self._step_sync = step_sync + self._hidden = hidden + self._summary = summary + self._goal = goal + self._overwrite = overwrite + + def _set_callback(self, cb: Callable[[pb.MetricRecord], None]) -> None: + self._callback = cb + + @property + def name(self) -> str: + return self._name + + @property + def step_metric(self) -> Optional[str]: + return self._step_metric + + @property + def step_sync(self) -> Optional[bool]: + return self._step_sync + + @property + def summary(self) -> Optional[Tuple[str, ...]]: + if self._summary is None: + return None + return tuple(self._summary) + + @property + def hidden(self) -> Optional[bool]: + return self._hidden + + @property + def goal(self) -> Optional[str]: + goal_dict = dict(min="minimize", max="maximize") + return goal_dict[self._goal] if self._goal else None + + def _commit(self) -> None: + m = pb.MetricRecord() + m.options.defined = True + if self._name.endswith("*"): + m.glob_name = self._name + else: + m.name = self._name + if self._step_metric: + m.step_metric = self._step_metric + if self._step_sync: + m.options.step_sync = self._step_sync + if self._hidden: + m.options.hidden = self._hidden + if self._summary: + summary_set = set(self._summary) + if "min" in summary_set: + m.summary.min = True + if "max" in summary_set: + m.summary.max = True + if "mean" in summary_set: + m.summary.mean = True + if "last" in summary_set: + m.summary.last = True + if "copy" in summary_set: + m.summary.copy = True + if "none" in summary_set: + m.summary.none = True + if "best" in summary_set: + m.summary.best = True + if self._goal == "min": + m.goal = m.GOAL_MINIMIZE + if self._goal == "max": + m.goal = m.GOAL_MAXIMIZE + if self._overwrite: + m._control.overwrite = self._overwrite + if self._callback: + self._callback(m) diff --git a/lib/python3.12/site-packages/wandb/sdk/wandb_require.py b/lib/python3.12/site-packages/wandb/sdk/wandb_require.py new file mode 100644 index 0000000000000000000000000000000000000000..b1540e29afe1156b5ed994d662203c68855ac46a --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/wandb_require.py @@ -0,0 +1,88 @@ +"""Feature Flags Module. + +This module implements a feature flag system for the wandb library to require experimental features +and notify the user when features have been deprecated. + +Example: + import wandb + wandb.require("wandb-service@beta") + wandb.require("incremental-artifacts@beta") +""" + +from __future__ import annotations + +from typing import Iterable + +import wandb +from wandb.errors import UnsupportedError + + +class _Requires: + """Internal feature class.""" + + _features: tuple[str, ...] + + def __init__(self, features: str | Iterable[str]) -> None: + self._features = ( + tuple([features]) if isinstance(features, str) else tuple(features) + ) + + def require_require(self) -> None: + pass + + def require_service(self) -> None: + # Legacy no-op kept solely for backward compatibility: + # some integrations (e.g. PyTorch Lightning) still call + # `wandb.require('service')`, which routes here. + wandb.termwarn( + "`wandb.require('service')` is a no-op as it is now the default behavior." + ) + + def require_core(self) -> None: + # Legacy no-op kept solely for backward compatibility: + # many public codebases still call `wandb.require('core')`. + wandb.termwarn( + "`wandb.require('core')` is a no-op as it is now the default behavior." + ) + + def apply(self) -> None: + """Call require_* method for supported features.""" + last_message: str = "" + for feature_item in self._features: + full_feature = feature_item.split("@", 2)[0] + feature = full_feature.split(":", 2)[0] + func_str = "require_{}".format(feature.replace("-", "_")) + func = getattr(self, func_str, None) + if not func: + last_message = f"require() unsupported requirement: {feature}" + wandb.termwarn(last_message) + continue + func() + + if last_message: + raise UnsupportedError(last_message) + + +def require( + requirement: str | Iterable[str] | None = None, + experiment: str | Iterable[str] | None = None, +) -> None: + """Indicate which experimental features are used by the script. + + This should be called before any other `wandb` functions, ideally right + after importing `wandb`. + + Args: + requirement: The name of a feature to require or an iterable of + feature names. + experiment: An alias for `requirement`. + + Raises: + wandb.errors.UnsupportedError: If a feature name is unknown. + """ + features = requirement or experiment + if not features: + return + + f = _Requires(features=features) + f.apply() diff --git a/lib/python3.12/site-packages/wandb/sdk/wandb_require_helpers.py b/lib/python3.12/site-packages/wandb/sdk/wandb_require_helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..3f14583ac53efd1b2b49b76ea1a0a27e65e88627 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/wandb_require_helpers.py @@ -0,0 +1,44 @@ +import os +from functools import wraps +from typing import Any, Callable, Dict, TypeVar, cast + +FuncT = TypeVar("FuncT", bound=Callable[..., Any]) + +requirement_env_var_mapping: Dict[str, str] = { + "report-editing:v0": "WANDB_REQUIRE_REPORT_EDITING_V0" +} + + +def requires(requirement: str) -> FuncT: # type: ignore + """Decorate functions to gate features with wandb.require.""" + env_var = requirement_env_var_mapping[requirement] + + def deco(func: FuncT) -> FuncT: + @wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> Any: + if not os.getenv(env_var): + raise Exception( + f"You need to enable this feature with `wandb.require({requirement!r})`" + ) + return func(*args, **kwargs) + + return cast(FuncT, wrapper) + + return cast(FuncT, deco) + + +class RequiresMixin: + requirement = "" + + def __init__(self) -> None: + self._check_if_requirements_met() + + def __post_init__(self) -> None: + self._check_if_requirements_met() + + def _check_if_requirements_met(self) -> None: + env_var = requirement_env_var_mapping[self.requirement] + if not os.getenv(env_var): + raise Exception( + f'You must explicitly enable this feature with `wandb.require("{self.requirement})"' + ) diff --git a/lib/python3.12/site-packages/wandb/sdk/wandb_run.py b/lib/python3.12/site-packages/wandb/sdk/wandb_run.py new file mode 100644 index 0000000000000000000000000000000000000000..54fd44606a2b9aa9b54192c07b1ac2a7ca8611fb --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/wandb_run.py @@ -0,0 +1,4060 @@ +from __future__ import annotations + +import asyncio +import functools +import glob +import json +import logging +import numbers +import os +import pathlib +import re +import sys +import threading +import time +import traceback +from collections.abc import Mapping +from dataclasses import dataclass, field +from datetime import datetime, timedelta, timezone +from enum import IntEnum +from types import TracebackType +from typing import TYPE_CHECKING, Callable, Sequence, TextIO, TypeVar + +import requests +from typing_extensions import Any, Concatenate, Literal, NamedTuple, ParamSpec + +import wandb +import wandb.env +import wandb.util +from wandb import trigger +from wandb.apis import internal, public +from wandb.apis.public import Api as PublicApi +from wandb.errors import CommError, UsageError +from wandb.errors.links import url_registry +from wandb.integration.torch import wandb_torch +from wandb.plot import CustomChart, Visualize +from wandb.proto.wandb_deprecated import Deprecated +from wandb.proto.wandb_internal_pb2 import ( + MetricRecord, + PollExitResponse, + Result, + RunRecord, +) +from wandb.sdk.artifacts._internal_artifact import InternalArtifact +from wandb.sdk.artifacts._validators import is_artifact_registry_project +from wandb.sdk.artifacts.artifact import Artifact +from wandb.sdk.internal import job_builder +from wandb.sdk.lib import asyncio_compat, wb_logging +from wandb.sdk.lib.import_hooks import ( + register_post_import_hook, + unregister_post_import_hook, +) +from wandb.sdk.lib.paths import FilePathStr, StrPath +from wandb.util import ( + _is_artifact_object, + _is_artifact_string, + _is_artifact_version_weave_dict, + _is_py_requirements_or_dockerfile, + _resolve_aliases, + add_import_hook, + parse_artifact_string, +) + +from . import wandb_config, wandb_metric, wandb_summary +from .artifacts._validators import ( + MAX_ARTIFACT_METADATA_KEYS, + ArtifactPath, + validate_aliases, + validate_tags, +) +from .data_types._dtypes import TypeRegistry +from .interface.interface import FilesDict, GlobStr, InterfaceBase, PolicyName +from .interface.summary_record import SummaryRecord +from .lib import ( + config_util, + deprecate, + filenames, + filesystem, + interrupt, + ipython, + module, + printer, + progress, + proto_util, + redirect, + telemetry, +) +from .lib.exit_hooks import ExitHooks +from .mailbox import ( + HandleAbandonedError, + MailboxClosedError, + MailboxHandle, + wait_with_progress, +) +from .wandb_alerts import AlertLevel +from .wandb_settings import Settings +from .wandb_setup import _WandbSetup + +if TYPE_CHECKING: + from typing import TypedDict + + import torch # type: ignore [import-not-found] + + import wandb.apis.public + import wandb.sdk.backend.backend + import wandb.sdk.interface.interface_queue + from wandb.proto.wandb_internal_pb2 import ( + GetSummaryResponse, + InternalMessagesResponse, + SampledHistoryResponse, + ) + + class GitSourceDict(TypedDict): + remote: str + commit: str + entrypoint: list[str] + args: Sequence[str] + + class ArtifactSourceDict(TypedDict): + artifact: str + entrypoint: list[str] + args: Sequence[str] + + class ImageSourceDict(TypedDict): + image: str + args: Sequence[str] + + class JobSourceDict(TypedDict, total=False): + _version: str + source_type: str + source: GitSourceDict | ArtifactSourceDict | ImageSourceDict + input_types: dict[str, Any] + output_types: dict[str, Any] + runtime: str | None + + +logger = logging.getLogger("wandb") +EXIT_TIMEOUT = 60 +RE_LABEL = re.compile(r"[a-zA-Z0-9_-]+$") + + +class TeardownStage(IntEnum): + EARLY = 1 + LATE = 2 + + +class TeardownHook(NamedTuple): + call: Callable[[], None] + stage: TeardownStage + + +class RunStatusChecker: + """Periodically polls the background process for relevant updates. + + - check if the user has requested a stop. + - check the network status. + - check the run sync status. + """ + + _stop_status_lock: threading.Lock + _stop_status_handle: MailboxHandle[Result] | None + _network_status_lock: threading.Lock + _network_status_handle: MailboxHandle[Result] | None + _internal_messages_lock: threading.Lock + _internal_messages_handle: MailboxHandle[Result] | None + + def __init__( + self, + run_id: str, + interface: InterfaceBase, + settings: Settings, + stop_polling_interval: int = 15, + retry_polling_interval: int = 5, + internal_messages_polling_interval: int = 10, + ) -> None: + self._run_id = run_id + self._interface = interface + self._stop_polling_interval = stop_polling_interval + self._retry_polling_interval = retry_polling_interval + self._internal_messages_polling_interval = internal_messages_polling_interval + self._settings = settings + + self._join_event = threading.Event() + + self._stop_status_lock = threading.Lock() + self._stop_status_handle = None + self._stop_thread = threading.Thread( + target=self.check_stop_status, + name="ChkStopThr", + daemon=True, + ) + + self._network_status_lock = threading.Lock() + self._network_status_handle = None + self._network_status_thread = threading.Thread( + target=self.check_network_status, + name="NetStatThr", + daemon=True, + ) + + self._internal_messages_lock = threading.Lock() + self._internal_messages_handle = None + self._internal_messages_thread = threading.Thread( + target=self.check_internal_messages, + name="IntMsgThr", + daemon=True, + ) + + def start(self) -> None: + self._stop_thread.start() + self._network_status_thread.start() + self._internal_messages_thread.start() + + @staticmethod + def _abandon_status_check( + lock: threading.Lock, + handle: MailboxHandle[Result] | None, + ): + with lock: + if handle: + handle.abandon() + + def _loop_check_status( + self, + *, + lock: threading.Lock, + set_handle: Any, + timeout: int, + request: Any, + process: Any, + ) -> None: + local_handle: MailboxHandle[Result] | None = None + join_requested = False + while not join_requested: + time_probe = time.monotonic() + if not local_handle: + try: + local_handle = request() + except MailboxClosedError: + # This can happen if the service process dies. + break + assert local_handle + + with lock: + if self._join_event.is_set(): + break + set_handle(local_handle) + + try: + result = local_handle.wait_or(timeout=timeout) + except HandleAbandonedError: + # This can happen if the service process dies. + break + except TimeoutError: + result = None + + with lock: + set_handle(None) + + if result: + process(result) + local_handle = None + + time_elapsed = time.monotonic() - time_probe + wait_time = max(timeout - time_elapsed, 0) + join_requested = self._join_event.wait(timeout=wait_time) + + def check_network_status(self) -> None: + def _process_network_status(result: Result) -> None: + network_status = result.response.network_status_response + for hr in network_status.network_responses: + if ( + hr.http_status_code == 200 or hr.http_status_code == 0 + ): # we use 0 for non-http errors (eg wandb errors) + wandb.termlog(f"{hr.http_response_text}") + else: + wandb.termlog( + f"{hr.http_status_code} encountered ({hr.http_response_text.rstrip()}), retrying request" + ) + + with wb_logging.log_to_run(self._run_id): + try: + self._loop_check_status( + lock=self._network_status_lock, + set_handle=lambda x: setattr(self, "_network_status_handle", x), + timeout=self._retry_polling_interval, + request=self._interface.deliver_network_status, + process=_process_network_status, + ) + except BrokenPipeError: + self._abandon_status_check( + self._network_status_lock, + self._network_status_handle, + ) + + def check_stop_status(self) -> None: + def _process_stop_status(result: Result) -> None: + stop_status = result.response.stop_status_response + if stop_status.run_should_stop: + # TODO(frz): This check is required + # until WB-3606 is resolved on server side. + if not wandb.agents.pyagent.is_running(): # type: ignore + interrupt.interrupt_main() + return + + with wb_logging.log_to_run(self._run_id): + try: + self._loop_check_status( + lock=self._stop_status_lock, + set_handle=lambda x: setattr(self, "_stop_status_handle", x), + timeout=self._stop_polling_interval, + request=self._interface.deliver_stop_status, + process=_process_stop_status, + ) + except BrokenPipeError: + self._abandon_status_check( + self._stop_status_lock, + self._stop_status_handle, + ) + + def check_internal_messages(self) -> None: + def _process_internal_messages(result: Result) -> None: + if ( + not self._settings.show_warnings + or self._settings.quiet + or self._settings.silent + ): + return + internal_messages = result.response.internal_messages_response + for msg in internal_messages.messages.warning: + wandb.termwarn(msg, repeat=False) + + with wb_logging.log_to_run(self._run_id): + try: + self._loop_check_status( + lock=self._internal_messages_lock, + set_handle=lambda x: setattr(self, "_internal_messages_handle", x), + timeout=self._internal_messages_polling_interval, + request=self._interface.deliver_internal_messages, + process=_process_internal_messages, + ) + except BrokenPipeError: + self._abandon_status_check( + self._internal_messages_lock, + self._internal_messages_handle, + ) + + def stop(self) -> None: + self._join_event.set() + self._abandon_status_check( + self._stop_status_lock, + self._stop_status_handle, + ) + self._abandon_status_check( + self._network_status_lock, + self._network_status_handle, + ) + self._abandon_status_check( + self._internal_messages_lock, + self._internal_messages_handle, + ) + + def join(self) -> None: + self.stop() + self._stop_thread.join() + self._network_status_thread.join() + self._internal_messages_thread.join() + + +_P = ParamSpec("_P") +_T = TypeVar("_T") + + +def _log_to_run( + func: Callable[Concatenate[Run, _P], _T], +) -> Callable[Concatenate[Run, _P], _T]: + """Decorate a Run method to set the run ID in the logging context. + + Any logs during the execution of the method go to the run's log file + and not to other runs' log files. + + This is meant for use on all public methods and some callbacks. Private + methods can be assumed to be called from some public method somewhere. + The general rule is to use it on methods that can be called from a + context that isn't specific to this run (such as all user code or + internal methods that aren't run-specific). + """ + + @functools.wraps(func) + def wrapper(self: Run, *args, **kwargs) -> _T: + # In "attach" usage, many properties of the Run are not initially + # populated. + if hasattr(self, "_settings"): + run_id = self._settings.run_id + else: + run_id = self._attach_id + + with wb_logging.log_to_run(run_id): + return func(self, *args, **kwargs) + + return wrapper + + +_is_attaching: str = "" + + +def _attach( + func: Callable[Concatenate[Run, _P], _T], +) -> Callable[Concatenate[Run, _P], _T]: + """Decorate a Run method to auto-attach when in a new process. + + When in a forked process or using a pickled Run instance, this automatically + connects to the service process to "attach" to the existing run. + """ + + @functools.wraps(func) + def wrapper(self: Run, *args, **kwargs) -> _T: + global _is_attaching + + # The _attach_id attribute is only None when running in the "disable + # service" mode. + # + # Since it is set early in `__init__` and included in the run's pickled + # state, the attribute always exists. + is_using_service = self._attach_id is not None + + # The _attach_pid attribute is not pickled, so it might not exist. + # It is set when the run is initialized. + attach_pid = getattr(self, "_attach_pid", None) + + if is_using_service and attach_pid != os.getpid(): + if _is_attaching: + raise RuntimeError( + f"Trying to attach `{func.__name__}`" + f" while in the middle of attaching `{_is_attaching}`" + ) + + _is_attaching = func.__name__ + try: + wandb._attach(run=self) # type: ignore + finally: + _is_attaching = "" + + return func(self, *args, **kwargs) + + return wrapper + + +def _raise_if_finished( + func: Callable[Concatenate[Run, _P], _T], +) -> Callable[Concatenate[Run, _P], _T]: + """Decorate a Run method to raise an error after the run is finished.""" + + @functools.wraps(func) + def wrapper_fn(self: Run, *args, **kwargs) -> _T: + if not getattr(self, "_is_finished", False): + return func(self, *args, **kwargs) + + message = ( + f"Run ({self.id}) is finished. The call to" + f" `{func.__name__}` will be ignored." + f" Please make sure that you are using an active run." + ) + + raise UsageError(message) + + return wrapper_fn + + +@dataclass +class RunStatus: + sync_items_total: int = field(default=0) + sync_items_pending: int = field(default=0) + sync_time: datetime | None = field(default=None) + + +class Run: + """A unit of computation logged by wandb. Typically, this is an ML experiment. + + Create a run with `wandb.init()`: + ```python + import wandb + + run = wandb.init() + ``` + + There is only ever at most one active `wandb.Run` in any process, + and it is accessible as `wandb.run`: + ```python + import wandb + + assert wandb.run is None + + wandb.init() + + assert wandb.run is not None + ``` + anything you log with `wandb.log` will be sent to that run. + + If you want to start more runs in the same script or notebook, you'll need to + finish the run that is in-flight. Runs can be finished with `wandb.finish` or + by using them in a `with` block: + ```python + import wandb + + wandb.init() + wandb.finish() + + assert wandb.run is None + + with wandb.init() as run: + pass # log data here + + assert wandb.run is None + ``` + + See the documentation for `wandb.init` for more on creating runs, or check out + [our guide to `wandb.init`](https://docs.wandb.ai/guides/track/launch). + + In distributed training, you can either create a single run in the rank 0 process + and then log information only from that process, or you can create a run in each process, + logging from each separately, and group the results together with the `group` argument + to `wandb.init`. For more details on distributed training with W&B, check out + [our guide](https://docs.wandb.ai/guides/track/log/distributed-training). + + Currently, there is a parallel `Run` object in the `wandb.Api`. Eventually these + two objects will be merged. + + Attributes: + summary: (Summary) Single values set for each `wandb.log()` key. By + default, summary is set to the last value logged. You can manually + set summary to the best value, like max accuracy, instead of the + final value. + """ + + _telemetry_obj: telemetry.TelemetryRecord + _telemetry_obj_active: bool + _telemetry_obj_dirty: bool + _telemetry_obj_flushed: bytes + + _teardown_hooks: list[TeardownHook] + + _backend: wandb.sdk.backend.backend.Backend | None + _internal_run_interface: wandb.sdk.interface.interface_queue.InterfaceQueue | None + _wl: _WandbSetup | None + + _out_redir: redirect.RedirectBase | None + _err_redir: redirect.RedirectBase | None + _redirect_cb: Callable[[str, str], None] | None + _redirect_raw_cb: Callable[[str, str], None] | None + _output_writer: filesystem.CRDedupedFile | None + + _atexit_cleanup_called: bool + _hooks: ExitHooks | None + _exit_code: int | None + + _run_status_checker: RunStatusChecker | None + + _sampled_history: SampledHistoryResponse | None + _final_summary: GetSummaryResponse | None + _poll_exit_handle: MailboxHandle[Result] | None + _poll_exit_response: PollExitResponse | None + _internal_messages_response: InternalMessagesResponse | None + + _stdout_slave_fd: int | None + _stderr_slave_fd: int | None + _artifact_slots: list[str] + + _init_pid: int + _attach_pid: int + + _attach_id: str | None + _is_attached: bool + _is_finished: bool + _settings: Settings + + _forked: bool + + _launch_artifacts: dict[str, Any] | None + _printer: printer.Printer + + summary: wandb_summary.Summary + + def __init__( + self, + settings: Settings, + config: dict[str, Any] | None = None, + sweep_config: dict[str, Any] | None = None, + launch_config: dict[str, Any] | None = None, + ) -> None: + # pid is set, so we know if this run object was initialized by this process + self._init_pid = os.getpid() + self._attach_id = None + + if settings._noop: + # TODO: properly handle setting for disabled mode + self._settings = settings + return + + self._init( + settings=settings, + config=config, + sweep_config=sweep_config, + launch_config=launch_config, + ) + + def _init( + self, + settings: Settings, + config: dict[str, Any] | None = None, + sweep_config: dict[str, Any] | None = None, + launch_config: dict[str, Any] | None = None, + ) -> None: + self._settings = settings + + self._config = wandb_config.Config() + self._config._set_callback(self._config_callback) + self._config._set_artifact_callback(self._config_artifact_callback) + self._config._set_settings(self._settings) + + # The _wandb key is always expected on the run config. + wandb_key = "_wandb" + self._config._update({wandb_key: dict()}) + + # TODO: perhaps this should be a property that is a noop on a finished run + self.summary = wandb_summary.Summary( + self._summary_get_current_summary_callback, + ) + self.summary._set_update_callback(self._summary_update_callback) + + self._step = 0 + self._starting_step = 0 + self._start_runtime = 0 + # TODO: eventually would be nice to make this configurable using self._settings._start_time + # need to test (jhr): if you set start time to 2 days ago and run a test for 15 minutes, + # does the total time get calculated right (not as 2 days and 15 minutes)? + self._start_time = time.time() + + self._printer = printer.new_printer(settings) + + self._torch_history: wandb_torch.TorchHistory | None = None # type: ignore + + self._backend = None + self._internal_run_interface = None + self._wl = None + + self._hooks = None + self._teardown_hooks = [] + + self._output_writer = None + self._out_redir = None + self._err_redir = None + self._stdout_slave_fd = None + self._stderr_slave_fd = None + + self._exit_code = None + self._exit_result = None + + self._used_artifact_slots: dict[str, str] = {} + + # Created when the run "starts". + self._run_status_checker = None + + self._sampled_history = None + self._final_summary = None + self._poll_exit_response = None + self._internal_messages_response = None + self._poll_exit_handle = None + + # Initialize telemetry object + self._telemetry_obj = telemetry.TelemetryRecord() + self._telemetry_obj_active = False + self._telemetry_obj_flushed = b"" + self._telemetry_obj_dirty = False + + self._atexit_cleanup_called = False + + # Initial scope setup for sentry. + # This might get updated when the actual run comes back. + wandb._sentry.configure_scope( + tags=dict(self._settings), + process_context="user", + ) + + self._launch_artifact_mapping: dict[str, Any] = {} + self._unique_launch_artifact_sequence_names: dict[str, Any] = {} + + # Populate config + config = config or dict() + self._config._update(config, allow_val_change=True, ignore_locked=True) + + if sweep_config: + self._config.merge_locked( + sweep_config, user="sweep", _allow_val_change=True + ) + + if launch_config: + self._config.merge_locked( + launch_config, user="launch", _allow_val_change=True + ) + + # if run is from a launch queue, add queue id to _wandb config + launch_queue_name = wandb.env.get_launch_queue_name() + if launch_queue_name: + self._config[wandb_key]["launch_queue_name"] = launch_queue_name + + launch_queue_entity = wandb.env.get_launch_queue_entity() + if launch_queue_entity: + self._config[wandb_key]["launch_queue_entity"] = launch_queue_entity + + launch_trace_id = wandb.env.get_launch_trace_id() + if launch_trace_id: + self._config[wandb_key]["launch_trace_id"] = launch_trace_id + + self._attach_id = None + self._is_attached = False + self._is_finished = False + + self._attach_pid = os.getpid() + self._forked = False + # for now, use runid as attach id, this could/should be versioned in the future + self._attach_id = self._settings.run_id + + def _handle_launch_artifact_overrides(self) -> None: + if self._settings.launch and (os.environ.get("WANDB_ARTIFACTS") is not None): + try: + artifacts: dict[str, Any] = json.loads( + os.environ.get("WANDB_ARTIFACTS", "{}") + ) + except (ValueError, SyntaxError): + wandb.termwarn("Malformed WANDB_ARTIFACTS, using original artifacts") + else: + self._initialize_launch_artifact_maps(artifacts) + + elif ( + self._settings.launch + and self._settings.launch_config_path + and os.path.exists(self._settings.launch_config_path) + ): + self.save(self._settings.launch_config_path) + with open(self._settings.launch_config_path) as fp: + launch_config = json.loads(fp.read()) + if launch_config.get("overrides", {}).get("artifacts") is not None: + artifacts = launch_config.get("overrides").get("artifacts") + self._initialize_launch_artifact_maps(artifacts) + + def _initialize_launch_artifact_maps(self, artifacts: dict[str, Any]) -> None: + for key, item in artifacts.items(): + self._launch_artifact_mapping[key] = item + artifact_sequence_tuple_or_slot = key.split(":") + + if len(artifact_sequence_tuple_or_slot) == 2: + sequence_name = artifact_sequence_tuple_or_slot[0].split("/")[-1] + if self._unique_launch_artifact_sequence_names.get(sequence_name): + self._unique_launch_artifact_sequence_names.pop(sequence_name) + else: + self._unique_launch_artifact_sequence_names[sequence_name] = item + + def _telemetry_callback(self, telem_obj: telemetry.TelemetryRecord) -> None: + if not hasattr(self, "_telemetry_obj") or self._is_finished: + return + + self._telemetry_obj.MergeFrom(telem_obj) + self._telemetry_obj_dirty = True + self._telemetry_flush() + + def _telemetry_flush(self) -> None: + if not hasattr(self, "_telemetry_obj"): + return + if not self._telemetry_obj_active: + return + if not self._telemetry_obj_dirty: + return + if self._backend and self._backend.interface: + serialized = self._telemetry_obj.SerializeToString() + if serialized == self._telemetry_obj_flushed: + return + self._backend.interface._publish_telemetry(self._telemetry_obj) + self._telemetry_obj_flushed = serialized + self._telemetry_obj_dirty = False + + def _freeze(self) -> None: + self._frozen = True + + def __setattr__(self, attr: str, value: object) -> None: + if getattr(self, "_frozen", None) and not hasattr(self, attr): + raise Exception(f"Attribute {attr} is not supported on Run object.") + super().__setattr__(attr, value) + + def __deepcopy__(self, memo: dict[int, Any]) -> Run: + return self + + def __getstate__(self) -> Any: + """Return run state as a custom pickle.""" + # We only pickle in service mode + if not self._settings: + return + + _attach_id = self._attach_id + if not _attach_id: + return + + return dict( + _attach_id=_attach_id, + _init_pid=self._init_pid, + _is_finished=self._is_finished, + ) + + def __setstate__(self, state: Any) -> None: + """Set run state from a custom pickle.""" + if not state: + return + + _attach_id = state.get("_attach_id") + if not _attach_id: + return + + if state["_init_pid"] == os.getpid(): + raise RuntimeError("attach in the same process is not supported currently") + + self.__dict__.update(state) + + @property + def _torch(self) -> wandb_torch.TorchHistory: # type: ignore + if self._torch_history is None: + self._torch_history = wandb_torch.TorchHistory() # type: ignore + return self._torch_history + + @property + @_log_to_run + @_attach + def settings(self) -> Settings: + """A frozen copy of run's Settings object.""" + return self._settings.model_copy(deep=True) + + @property + @_log_to_run + @_attach + def dir(self) -> str: + """The directory where files associated with the run are saved.""" + return self._settings.files_dir + + @property + @_log_to_run + @_attach + def config(self) -> wandb_config.Config: + """Config object associated with this run.""" + return self._config + + @property + @_log_to_run + @_attach + def config_static(self) -> wandb_config.ConfigStatic: + return wandb_config.ConfigStatic(self._config) + + @property + @_log_to_run + @_attach + def name(self) -> str | None: + """Display name of the run. + + Display names are not guaranteed to be unique and may be descriptive. + By default, they are randomly generated. + """ + return self._settings.run_name + + @name.setter + @_log_to_run + @_raise_if_finished + def name(self, name: str) -> None: + with telemetry.context(run=self) as tel: + tel.feature.set_run_name = True + self._settings.run_name = name + if self._backend and self._backend.interface: + self._backend.interface.publish_run(self) + + @property + @_log_to_run + @_attach + def notes(self) -> str | None: + """Notes associated with the run, if there are any. + + Notes can be a multiline string and can also use markdown and latex + equations inside `$$`, like `$x + 3$`. + """ + return self._settings.run_notes + + @notes.setter + @_log_to_run + @_raise_if_finished + def notes(self, notes: str) -> None: + self._settings.run_notes = notes + if self._backend and self._backend.interface: + self._backend.interface.publish_run(self) + + @property + @_log_to_run + @_attach + def tags(self) -> tuple | None: + """Tags associated with the run, if there are any.""" + return self._settings.run_tags or () + + @tags.setter + @_log_to_run + @_raise_if_finished + def tags(self, tags: Sequence) -> None: + with telemetry.context(run=self) as tel: + tel.feature.set_run_tags = True + self._settings.run_tags = tuple(tags) + if self._backend and self._backend.interface: + self._backend.interface.publish_run(self) + + @property + @_log_to_run + @_attach + def id(self) -> str: + """Identifier for this run.""" + assert self._settings.run_id is not None + return self._settings.run_id + + @property + @_log_to_run + @_attach + def sweep_id(self) -> str | None: + """Identifier for the sweep associated with the run, if there is one.""" + return self._settings.sweep_id + + def _get_path(self) -> str: + return "/".join( + e + for e in [ + self._settings.entity, + self._settings.project, + self._settings.run_id, + ] + if e is not None + ) + + @property + @_log_to_run + @_attach + def path(self) -> str: + """Path to the run. + + Run paths include entity, project, and run ID, in the format + `entity/project/run_id`. + """ + return self._get_path() + + @property + @_log_to_run + @_attach + def start_time(self) -> float: + """Unix timestamp (in seconds) of when the run started.""" + return self._start_time + + @property + @_log_to_run + @_attach + def starting_step(self) -> int: + """The first step of the run.""" + return self._starting_step + + @property + @_log_to_run + @_attach + def resumed(self) -> bool: + """True if the run was resumed, False otherwise.""" + return self._settings.resumed + + @property + @_log_to_run + @_attach + def step(self) -> int: + """Current value of the step. + + This counter is incremented by `wandb.log`. + """ + return self._step + + @property + @_log_to_run + @_attach + def offline(self) -> bool: + return self._settings._offline + + @property + @_log_to_run + @_attach + def disabled(self) -> bool: + return self._settings._noop + + @property + @_log_to_run + @_attach + def group(self) -> str: + """Name of the group associated with the run. + + Setting a group helps the W&B UI organize runs in a sensible way. + + If you are doing a distributed training you should give all of the + runs in the training the same group. + If you are doing cross-validation you should give all the cross-validation + folds the same group. + """ + return self._settings.run_group or "" + + @property + @_log_to_run + @_attach + def job_type(self) -> str: + return self._settings.run_job_type or "" + + def project_name(self) -> str: + """Name of the W&B project associated with the run. + + Note: this method is deprecated and will be removed in a future release. + Please use `run.project` instead. + """ + deprecate.deprecate( + field_name=Deprecated.run__project_name, + warning_message=( + "The project_name method is deprecated and will be removed in a" + " future release. Please use `run.project` instead." + ), + ) + return self.project + + @property + @_log_to_run + @_attach + def project(self) -> str: + """Name of the W&B project associated with the run.""" + assert self._settings.project is not None + return self._settings.project + + @_log_to_run + def get_project_url(self) -> str | None: + """URL of the W&B project associated with the run, if there is one. + + Offline runs do not have a project URL. + + Note: this method is deprecated and will be removed in a future release. + Please use `run.project_url` instead. + """ + deprecate.deprecate( + field_name=Deprecated.run__get_project_url, + warning_message=( + "The get_project_url method is deprecated and will be removed in a" + " future release. Please use `run.project_url` instead." + ), + ) + return self.project_url + + @property + @_log_to_run + @_attach + def project_url(self) -> str | None: + """URL of the W&B project associated with the run, if there is one. + + Offline runs do not have a project URL. + """ + if self._settings._offline: + wandb.termwarn("URL not available in offline run") + return None + return self._settings.project_url + + @_raise_if_finished + @_log_to_run + @_attach + def log_code( + self, + root: str | None = ".", + name: str | None = None, + include_fn: Callable[[str, str], bool] + | Callable[[str], bool] = _is_py_requirements_or_dockerfile, + exclude_fn: Callable[[str, str], bool] + | Callable[[str], bool] = filenames.exclude_wandb_fn, + ) -> Artifact | None: + """Save the current state of your code to a W&B Artifact. + + By default, it walks the current directory and logs all files that end with `.py`. + + Args: + root: The relative (to `os.getcwd()`) or absolute path to recursively find code from. + name: (str, optional) The name of our code artifact. By default, we'll name + the artifact `source-$PROJECT_ID-$ENTRYPOINT_RELPATH`. There may be scenarios where you want + many runs to share the same artifact. Specifying name allows you to achieve that. + include_fn: A callable that accepts a file path and (optionally) root path and + returns True when it should be included and False otherwise. This + defaults to: `lambda path, root: path.endswith(".py")` + exclude_fn: A callable that accepts a file path and (optionally) root path and + returns `True` when it should be excluded and `False` otherwise. This + defaults to a function that excludes all files within `/.wandb/` + and `/wandb/` directories. + + Examples: + Basic usage + ```python + run.log_code() + ``` + + Advanced usage + ```python + run.log_code( + "../", + include_fn=lambda path: path.endswith(".py") or path.endswith(".ipynb"), + exclude_fn=lambda path, root: os.path.relpath(path, root).startswith( + "cache/" + ), + ) + ``` + + Returns: + An `Artifact` object if code was logged + """ + if name is None: + if self.settings._jupyter: + notebook_name = None + if self.settings.notebook_name: + notebook_name = self.settings.notebook_name + elif self.settings.x_jupyter_path: + if self.settings.x_jupyter_path.startswith("fileId="): + notebook_name = self.settings.x_jupyter_name + else: + notebook_name = self.settings.x_jupyter_path + name_string = f"{self._settings.project}-{notebook_name}" + else: + name_string = ( + f"{self._settings.project}-{self._settings.program_relpath}" + ) + name = wandb.util.make_artifact_name_safe(f"source-{name_string}") + art = InternalArtifact(name, "code") + files_added = False + if root is not None: + root = os.path.abspath(root) + for file_path in filenames.filtered_dir(root, include_fn, exclude_fn): + files_added = True + save_name = os.path.relpath(file_path, root) + art.add_file(file_path, name=save_name) + # Add any manually staged files such as ipynb notebooks + for dirpath, _, files in os.walk(self._settings._tmp_code_dir): + for fname in files: + file_path = os.path.join(dirpath, fname) + save_name = os.path.relpath(file_path, self._settings._tmp_code_dir) + files_added = True + art.add_file(file_path, name=save_name) + if not files_added: + wandb.termwarn( + "No relevant files were detected in the specified directory. No code will be logged to your run." + ) + return None + + artifact = self._log_artifact(art) + + self._config.update( + {"_wandb": {"code_path": artifact.name}}, + allow_val_change=True, + ) + + return artifact + + @_log_to_run + def get_sweep_url(self) -> str | None: + """The URL of the sweep associated with the run, if there is one. + + Offline runs do not have a sweep URL. + + Note: this method is deprecated and will be removed in a future release. + Please use `run.sweep_url` instead. + """ + deprecate.deprecate( + field_name=Deprecated.run__get_sweep_url, + warning_message=( + "The get_sweep_url method is deprecated and will be removed in a" + " future release. Please use `run.sweep_url` instead." + ), + ) + return self.sweep_url + + @property + @_attach + def sweep_url(self) -> str | None: + """URL of the sweep associated with the run, if there is one. + + Offline runs do not have a sweep URL. + """ + if self._settings._offline: + wandb.termwarn("URL not available in offline run") + return None + return self._settings.sweep_url + + @_log_to_run + def get_url(self) -> str | None: + """URL of the W&B run, if there is one. + + Offline runs do not have a URL. + + Note: this method is deprecated and will be removed in a future release. + Please use `run.url` instead. + """ + deprecate.deprecate( + field_name=Deprecated.run__get_url, + warning_message=( + "The get_url method is deprecated and will be removed in a" + " future release. Please use `run.url` instead." + ), + ) + return self.url + + @property + @_log_to_run + @_attach + def url(self) -> str | None: + """The url for the W&B run, if there is one. + + Offline runs will not have a url. + """ + if self._settings._offline: + wandb.termwarn("URL not available in offline run") + return None + return self._settings.run_url + + @property + @_log_to_run + @_attach + def entity(self) -> str: + """The name of the W&B entity associated with the run. + + Entity can be a username or the name of a team or organization. + """ + return self._settings.entity or "" + + def _label_internal( + self, + code: str | None = None, + repo: str | None = None, + code_version: str | None = None, + ) -> None: + with telemetry.context(run=self) as tel: + if code and RE_LABEL.match(code): + tel.label.code_string = code + if repo and RE_LABEL.match(repo): + tel.label.repo_string = repo + if code_version and RE_LABEL.match(code_version): + tel.label.code_version = code_version + + def _label( + self, + code: str | None = None, + repo: str | None = None, + code_version: str | None = None, + **kwargs: str, + ) -> None: + if self._settings.label_disable: + return + for k, v in (("code", code), ("repo", repo), ("code_version", code_version)): + if v and not RE_LABEL.match(v): + wandb.termwarn( + f"Label added for '{k}' with invalid identifier '{v}' (ignored).", + repeat=False, + ) + for v in kwargs: + wandb.termwarn( + f"Label added for unsupported key {v!r} (ignored).", + repeat=False, + ) + + self._label_internal(code=code, repo=repo, code_version=code_version) + + # update telemetry in the backend immediately for _label() callers + self._telemetry_flush() + + def _label_probe_lines(self, lines: list[str]) -> None: + if not lines: + return + parsed = telemetry._parse_label_lines(lines) + if not parsed: + return + label_dict = {} + code = parsed.get("code") or parsed.get("c") + if code: + label_dict["code"] = code + repo = parsed.get("repo") or parsed.get("r") + if repo: + label_dict["repo"] = repo + code_ver = parsed.get("version") or parsed.get("v") + if code_ver: + label_dict["code_version"] = code_ver + self._label_internal(**label_dict) + + def _label_probe_main(self) -> None: + m = sys.modules.get("__main__") + if not m: + return + doc = getattr(m, "__doc__", None) + if not doc: + return + + doclines = doc.splitlines() + self._label_probe_lines(doclines) + + # TODO: annotate jupyter Notebook class + def _label_probe_notebook(self, notebook: Any) -> None: + logger.info("probe notebook") + lines = None + try: + data = notebook.probe_ipynb() + cell0 = data.get("cells", [])[0] + lines = cell0.get("source") + # kaggle returns a string instead of a list + if isinstance(lines, str): + lines = lines.split() + except Exception as e: + logger.info(f"Unable to probe notebook: {e}") + return + if lines: + self._label_probe_lines(lines) + + @_log_to_run + @_attach + def display(self, height: int = 420, hidden: bool = False) -> bool: + """Display this run in jupyter.""" + if self._settings.silent: + return False + + if not ipython.in_jupyter(): + return False + + try: + from IPython import display + except ImportError: + wandb.termwarn(".display() only works in jupyter environments") + return False + + display.display(display.HTML(self.to_html(height, hidden))) + return True + + @_log_to_run + @_attach + def to_html(self, height: int = 420, hidden: bool = False) -> str: + """Generate HTML containing an iframe displaying the current run.""" + url = self._settings.run_url + "?jupyter=true" + style = f"border:none;width:100%;height:{height}px;" + prefix = "" + if hidden: + style += "display:none;" + prefix = ipython.toggle_button() + return prefix + f"" + + def _repr_mimebundle_( + self, include: Any | None = None, exclude: Any | None = None + ) -> dict[str, str]: + return {"text/html": self.to_html(hidden=True)} + + @_log_to_run + @_raise_if_finished + def _config_callback( + self, + key: tuple[str, ...] | str | None = None, + val: Any | None = None, + data: dict[str, object] | None = None, + ) -> None: + logger.info(f"config_cb {key} {val} {data}") + if self._backend and self._backend.interface: + self._backend.interface.publish_config(key=key, val=val, data=data) + + @_log_to_run + def _config_artifact_callback( + self, key: str, val: str | Artifact | dict + ) -> Artifact: + # artifacts can look like dicts as they are passed into the run config + # since the run config stores them on the backend as a dict with fields shown + # in wandb.util.artifact_to_json + if _is_artifact_version_weave_dict(val): + assert isinstance(val, dict) + public_api = self._public_api() + artifact = Artifact._from_id(val["id"], public_api.client) + + assert artifact + return self.use_artifact(artifact) + elif _is_artifact_string(val): + # this will never fail, but is required to make mypy happy + assert isinstance(val, str) + artifact_string, base_url, is_id = parse_artifact_string(val) + overrides = {} + if base_url is not None: + overrides = {"base_url": base_url} + public_api = public.Api(overrides) + else: + public_api = self._public_api() + if is_id: + artifact = Artifact._from_id(artifact_string, public_api._client) + else: + artifact = public_api._artifact(name=artifact_string) + # in the future we'll need to support using artifacts from + # different instances of wandb. + + assert artifact + return self.use_artifact(artifact) + elif _is_artifact_object(val): + return self.use_artifact(val) + else: + raise ValueError( + f"Cannot call _config_artifact_callback on type {type(val)}" + ) + + def _set_config_wandb(self, key: str, val: Any) -> None: + self._config_callback(key=("_wandb", key), val=val) + + @_log_to_run + @_raise_if_finished + def _summary_update_callback(self, summary_record: SummaryRecord) -> None: + with telemetry.context(run=self) as tel: + tel.feature.set_summary = True + if self._backend and self._backend.interface: + self._backend.interface.publish_summary(self, summary_record) + + @_log_to_run + def _summary_get_current_summary_callback(self) -> dict[str, Any]: + if self._is_finished: + # TODO: WB-18420: fetch summary from backend and stage it before run is finished + wandb.termwarn("Summary data not available in finished run") + return {} + if not self._backend or not self._backend.interface: + return {} + handle = self._backend.interface.deliver_get_summary() + + try: + result = handle.wait_or(timeout=self._settings.summary_timeout) + except TimeoutError: + return {} + + get_summary_response = result.response.get_summary_response + return proto_util.dict_from_proto_list(get_summary_response.item) + + @_log_to_run + def _metric_callback(self, metric_record: MetricRecord) -> None: + if self._backend and self._backend.interface: + self._backend.interface._publish_metric(metric_record) + + @_log_to_run + def _publish_file(self, fname: str) -> None: + """Mark a run file to be uploaded with the run. + + This is a W&B-internal function: it can be used by other internal + wandb code. + + Args: + fname: The path to the file in the run's files directory, relative + to the run's files directory. + """ + if not self._backend or not self._backend.interface: + return + files: FilesDict = dict(files=[(GlobStr(fname), "now")]) + self._backend.interface.publish_files(files) + + def _pop_all_charts( + self, + data: dict[str, Any], + key_prefix: str | None = None, + ) -> dict[str, Any]: + """Pops all charts from a dictionary including nested charts. + + This function will return a mapping of the charts and a dot-separated + key for each chart. Indicating the path to the chart in the data dictionary. + """ + keys_to_remove = set() + charts: dict[str, Any] = {} + for k, v in data.items(): + key = f"{key_prefix}.{k}" if key_prefix else k + if isinstance(v, Visualize): + keys_to_remove.add(k) + charts[key] = v + elif isinstance(v, CustomChart): + keys_to_remove.add(k) + charts[key] = v + elif isinstance(v, dict): + nested_charts = self._pop_all_charts(v, key) + charts.update(nested_charts) + + for k in keys_to_remove: + data.pop(k) + + return charts + + def _serialize_custom_charts( + self, + data: dict[str, Any], + ) -> dict[str, Any]: + """Process and replace chart objects with their underlying table values. + + This processes the chart objects passed to `run.log()`, replacing their entries + in the given dictionary (which is saved to the run's history) and adding them + to the run's config. + + Args: + data: Dictionary containing data that may include plot objects + Plot objects can be nested in dictionaries, which will be processed recursively. + + Returns: + The processed dictionary with custom charts transformed into tables. + """ + if not data: + return data + + charts = self._pop_all_charts(data) + for k, v in charts.items(): + v.set_key(k) + self._config_callback( + val=v.spec.config_value, + key=v.spec.config_key, + ) + + if isinstance(v, CustomChart): + data[v.spec.table_key] = v.table + elif isinstance(v, Visualize): + data[k] = v.table + + return data + + @_log_to_run + def _partial_history_callback( + self, + data: dict[str, Any], + step: int | None = None, + commit: bool | None = None, + ) -> None: + if not (self._backend and self._backend.interface): + return + + data = data.copy() # avoid modifying the original data + + # Serialize custom charts before publishing + data = self._serialize_custom_charts(data) + + not_using_tensorboard = len(wandb.patched["tensorboard"]) == 0 + self._backend.interface.publish_partial_history( + self, + data, + user_step=self._step, + step=step, + flush=commit, + publish_step=not_using_tensorboard, + ) + + @_log_to_run + def _console_callback(self, name: str, data: str) -> None: + # logger.info("console callback: %s, %s", name, data) + if self._backend and self._backend.interface: + self._backend.interface.publish_output(name, data) + + @_log_to_run + @_raise_if_finished + def _console_raw_callback(self, name: str, data: str) -> None: + # logger.info("console callback: %s, %s", name, data) + + # NOTE: console output is only allowed on the process which installed the callback + # this will prevent potential corruption in the socket to the service. Other methods + # are protected by the _attach run decorator, but this callback was installed on the + # write function of stdout and stderr streams. + console_pid = getattr(self, "_attach_pid", 0) + if console_pid != os.getpid(): + return + + if self._backend and self._backend.interface: + self._backend.interface.publish_output_raw(name, data) + + @_log_to_run + def _tensorboard_callback( + self, logdir: str, save: bool = True, root_logdir: str = "" + ) -> None: + logger.info("tensorboard callback: %s, %s", logdir, save) + if self._backend and self._backend.interface: + self._backend.interface.publish_tbdata(logdir, save, root_logdir) + + def _set_library(self, library: _WandbSetup) -> None: + self._wl = library + + def _set_backend(self, backend: wandb.sdk.backend.backend.Backend) -> None: + self._backend = backend + + def _set_internal_run_interface( + self, + interface: wandb.sdk.interface.interface_queue.InterfaceQueue, + ) -> None: + self._internal_run_interface = interface + + def _set_teardown_hooks(self, hooks: list[TeardownHook]) -> None: + self._teardown_hooks = hooks + + def _set_run_obj(self, run_obj: RunRecord) -> None: # noqa: C901 + if run_obj.starting_step: + self._starting_step = run_obj.starting_step + self._step = run_obj.starting_step + + if run_obj.start_time: + self._start_time = run_obj.start_time.ToMicroseconds() / 1e6 + + if run_obj.runtime: + self._start_runtime = run_obj.runtime + + # Grab the config from resuming + if run_obj.config: + c_dict = config_util.dict_no_value_from_proto_list(run_obj.config.update) + # We update the config object here without triggering the callback + self._config._update(c_dict, allow_val_change=True, ignore_locked=True) + # Update the summary, this will trigger an un-needed graphql request :( + if run_obj.summary: + summary_dict = {} + for orig in run_obj.summary.update: + summary_dict[orig.key] = json.loads(orig.value_json) + if summary_dict: + self.summary.update(summary_dict) + + # update settings from run_obj + if run_obj.run_id: + self._settings.run_id = run_obj.run_id + if run_obj.entity: + self._settings.entity = run_obj.entity + if run_obj.project: + self._settings.project = run_obj.project + if run_obj.run_group: + self._settings.run_group = run_obj.run_group + if run_obj.job_type: + self._settings.run_job_type = run_obj.job_type + if run_obj.display_name: + self._settings.run_name = run_obj.display_name + if run_obj.notes: + self._settings.run_notes = run_obj.notes + if run_obj.tags: + self._settings.run_tags = tuple(run_obj.tags) + if run_obj.sweep_id: + self._settings.sweep_id = run_obj.sweep_id + if run_obj.host: + self._settings.host = run_obj.host + if run_obj.resumed: + self._settings.resumed = run_obj.resumed + if run_obj.git: + if run_obj.git.remote_url: + self._settings.git_remote_url = run_obj.git.remote_url + if run_obj.git.commit: + self._settings.git_commit = run_obj.git.commit + + if run_obj.forked: + self._forked = run_obj.forked + + wandb._sentry.configure_scope( + process_context="user", + tags=dict(self._settings), + ) + + def _populate_git_info(self) -> None: + from .lib.gitlib import GitRepo + + # Use user-provided git info if available, otherwise resolve it from the environment + try: + repo = GitRepo( + root=self._settings.git_root, + remote=self._settings.git_remote, + remote_url=self._settings.git_remote_url, + commit=self._settings.git_commit, + lazy=False, + ) + self._settings.git_remote_url = repo.remote_url + self._settings.git_commit = repo.last_commit + except Exception: + wandb.termwarn("Cannot find valid git repo associated with this directory.") + + def _add_singleton( + self, data_type: str, key: str, value: dict[int | str, str] + ) -> None: + """Store a singleton item to wandb config. + + A singleton in this context is a piece of data that is continually + logged with the same value in each history step, but represented + as a single item in the config. + + We do this to avoid filling up history with a lot of repeated unnecessary data + + Add singleton can be called many times in one run, and it will only be + updated when the value changes. The last value logged will be the one + persisted to the server. + """ + value_extra = {"type": data_type, "key": key, "value": value} + + if data_type not in self._config["_wandb"]: + self._config["_wandb"][data_type] = {} + + if data_type in self._config["_wandb"][data_type]: + old_value = self._config["_wandb"][data_type][key] + else: + old_value = None + + if value_extra != old_value: + self._config["_wandb"][data_type][key] = value_extra + self._config.persist() + + def _log( + self, + data: dict[str, Any], + step: int | None = None, + commit: bool | None = None, + ) -> None: + if not isinstance(data, Mapping): + raise TypeError("wandb.log must be passed a dictionary") + + if any(not isinstance(key, str) for key in data.keys()): + raise TypeError("Key values passed to `wandb.log` must be strings.") + + self._partial_history_callback(data, step, commit) + + if step is not None: + if os.getpid() != self._init_pid or self._is_attached: + wandb.termwarn( + "Note that setting step in multiprocessing can result in data loss. " + "Please use `run.define_metric(...)` to define a custom metric " + "to log your step values.", + repeat=False, + ) + # if step is passed in when tensorboard_sync is used we honor the step passed + # to make decisions about how to close out the history record, but will strip + # this history later on in publish_history() + if len(wandb.patched["tensorboard"]) > 0: + wandb.termwarn( + "Step cannot be set when using tensorboard syncing. " + "Please use `run.define_metric(...)` to define a custom metric " + "to log your step values.", + repeat=False, + ) + if step > self._step: + self._step = step + + if (step is None and commit is None) or commit: + self._step += 1 + + @_log_to_run + @_raise_if_finished + @_attach + def log( + self, + data: dict[str, Any], + step: int | None = None, + commit: bool | None = None, + ) -> None: + """Upload run data. + + Use `log` to log data from runs, such as scalars, images, video, + histograms, plots, and tables. + + See our [guides to logging](https://docs.wandb.ai/guides/track/log) for + live examples, code snippets, best practices, and more. + + The most basic usage is `run.log({"train-loss": 0.5, "accuracy": 0.9})`. + This will save the loss and accuracy to the run's history and update + the summary values for these metrics. + + Visualize logged data in the workspace at [wandb.ai](https://wandb.ai), + or locally on a [self-hosted instance](https://docs.wandb.ai/guides/hosting) + of the W&B app, or export data to visualize and explore locally, e.g. in + Jupyter notebooks, with [our API](https://docs.wandb.ai/guides/track/public-api-guide). + + Logged values don't have to be scalars. Logging any wandb object is supported. + For example `run.log({"example": wandb.Image("myimage.jpg")})` will log an + example image which will be displayed nicely in the W&B UI. + See the [reference documentation](https://docs.wandb.com/ref/python/data-types) + for all of the different supported types or check out our + [guides to logging](https://docs.wandb.ai/guides/track/log) for examples, + from 3D molecular structures and segmentation masks to PR curves and histograms. + You can use `wandb.Table` to log structured data. See our + [guide to logging tables](https://docs.wandb.ai/guides/models/tables/tables-walkthrough) + for details. + + The W&B UI organizes metrics with a forward slash (`/`) in their name + into sections named using the text before the final slash. For example, + the following results in two sections named "train" and "validate": + + ``` + run.log( + { + "train/accuracy": 0.9, + "train/loss": 30, + "validate/accuracy": 0.8, + "validate/loss": 20, + } + ) + ``` + + Only one level of nesting is supported; `run.log({"a/b/c": 1})` + produces a section named "a/b". + + `run.log` is not intended to be called more than a few times per second. + For optimal performance, limit your logging to once every N iterations, + or collect data over multiple iterations and log it in a single step. + + ### The W&B step + + With basic usage, each call to `log` creates a new "step". + The step must always increase, and it is not possible to log + to a previous step. + + Note that you can use any metric as the X axis in charts. + In many cases, it is better to treat the W&B step like + you'd treat a timestamp rather than a training step. + + ``` + # Example: log an "epoch" metric for use as an X axis. + run.log({"epoch": 40, "train-loss": 0.5}) + ``` + + See also [define_metric](https://docs.wandb.ai/ref/python/run#define_metric). + + It is possible to use multiple `log` invocations to log to + the same step with the `step` and `commit` parameters. + The following are all equivalent: + + ``` + # Normal usage: + run.log({"train-loss": 0.5, "accuracy": 0.8}) + run.log({"train-loss": 0.4, "accuracy": 0.9}) + + # Implicit step without auto-incrementing: + run.log({"train-loss": 0.5}, commit=False) + run.log({"accuracy": 0.8}) + run.log({"train-loss": 0.4}, commit=False) + run.log({"accuracy": 0.9}) + + # Explicit step: + run.log({"train-loss": 0.5}, step=current_step) + run.log({"accuracy": 0.8}, step=current_step) + current_step += 1 + run.log({"train-loss": 0.4}, step=current_step) + run.log({"accuracy": 0.9}, step=current_step) + ``` + + Args: + data: A `dict` with `str` keys and values that are serializable + Python objects including: `int`, `float` and `string`; + any of the `wandb.data_types`; lists, tuples and NumPy arrays + of serializable Python objects; other `dict`s of this + structure. + step: The step number to log. If `None`, then an implicit + auto-incrementing step is used. See the notes in + the description. + commit: If true, finalize and upload the step. If false, then + accumulate data for the step. See the notes in the description. + If `step` is `None`, then the default is `commit=True`; + otherwise, the default is `commit=False`. + + Examples: + For more and more detailed examples, see + [our guides to logging](https://docs.wandb.com/guides/track/log). + + ### Basic usage + ```python + import wandb + + run = wandb.init() + run.log({"accuracy": 0.9, "epoch": 5}) + ``` + + ### Incremental logging + ```python + import wandb + + run = wandb.init() + run.log({"loss": 0.2}, commit=False) + # Somewhere else when I'm ready to report this step: + run.log({"accuracy": 0.8}) + ``` + + ### Histogram + ```python + import numpy as np + import wandb + + # sample gradients at random from normal distribution + gradients = np.random.randn(100, 100) + run = wandb.init() + run.log({"gradients": wandb.Histogram(gradients)}) + ``` + + ### Image from numpy + ```python + import numpy as np + import wandb + + run = wandb.init() + examples = [] + for i in range(3): + pixels = np.random.randint(low=0, high=256, size=(100, 100, 3)) + image = wandb.Image(pixels, caption=f"random field {i}") + examples.append(image) + run.log({"examples": examples}) + ``` + + ### Image from PIL + ```python + import numpy as np + from PIL import Image as PILImage + import wandb + + run = wandb.init() + examples = [] + for i in range(3): + pixels = np.random.randint( + low=0, + high=256, + size=(100, 100, 3), + dtype=np.uint8, + ) + pil_image = PILImage.fromarray(pixels, mode="RGB") + image = wandb.Image(pil_image, caption=f"random field {i}") + examples.append(image) + run.log({"examples": examples}) + ``` + + ### Video from numpy + ```python + import numpy as np + import wandb + + run = wandb.init() + # axes are (time, channel, height, width) + frames = np.random.randint( + low=0, + high=256, + size=(10, 3, 100, 100), + dtype=np.uint8, + ) + run.log({"video": wandb.Video(frames, fps=4)}) + ``` + + ### Matplotlib Plot + ```python + from matplotlib import pyplot as plt + import numpy as np + import wandb + + run = wandb.init() + fig, ax = plt.subplots() + x = np.linspace(0, 10) + y = x * x + ax.plot(x, y) # plot y = x^2 + run.log({"chart": fig}) + ``` + + ### PR Curve + ```python + import wandb + + run = wandb.init() + run.log({"pr": wandb.plot.pr_curve(y_test, y_probas, labels)}) + ``` + + ### 3D Object + ```python + import wandb + + run = wandb.init() + run.log( + { + "generated_samples": [ + wandb.Object3D(open("sample.obj")), + wandb.Object3D(open("sample.gltf")), + wandb.Object3D(open("sample.glb")), + ] + } + ) + ``` + + Raises: + wandb.Error: if called before `wandb.init` + ValueError: if invalid data is passed + + """ + if step is not None: + with telemetry.context(run=self) as tel: + tel.feature.set_step_log = True + + if self._settings._shared and step is not None: + wandb.termwarn( + "In shared mode, the use of `wandb.log` with the step argument is not supported " + f"and will be ignored. Please refer to {url_registry.url('define-metric')} " + "on how to customize your x-axis.", + repeat=False, + ) + self._log(data=data, step=step, commit=commit) + + @_log_to_run + @_raise_if_finished + @_attach + def save( + self, + glob_str: str | os.PathLike, + base_path: str | os.PathLike | None = None, + policy: PolicyName = "live", + ) -> bool | list[str]: + """Sync one or more files to W&B. + + Relative paths are relative to the current working directory. + + A Unix glob, such as "myfiles/*", is expanded at the time `save` is + called regardless of the `policy`. In particular, new files are not + picked up automatically. + + A `base_path` may be provided to control the directory structure of + uploaded files. It should be a prefix of `glob_str`, and the directory + structure beneath it is preserved. It's best understood through + examples: + + ``` + wandb.save("these/are/myfiles/*") + # => Saves files in a "these/are/myfiles/" folder in the run. + + wandb.save("these/are/myfiles/*", base_path="these") + # => Saves files in an "are/myfiles/" folder in the run. + + wandb.save("/User/username/Documents/run123/*.txt") + # => Saves files in a "run123/" folder in the run. See note below. + + wandb.save("/User/username/Documents/run123/*.txt", base_path="/User") + # => Saves files in a "username/Documents/run123/" folder in the run. + + wandb.save("files/*/saveme.txt") + # => Saves each "saveme.txt" file in an appropriate subdirectory + # of "files/". + ``` + + Note: when given an absolute path or glob and no `base_path`, one + directory level is preserved as in the example above. + + Args: + glob_str: A relative or absolute path or Unix glob. + base_path: A path to use to infer a directory structure; see examples. + policy: One of `live`, `now`, or `end`. + * live: upload the file as it changes, overwriting the previous version + * now: upload the file once now + * end: upload file when the run ends + + Returns: + Paths to the symlinks created for the matched files. + + For historical reasons, this may return a boolean in legacy code. + """ + if isinstance(glob_str, bytes): + # Preserved for backward compatibility: allow bytes inputs. + glob_str = glob_str.decode("utf-8") + if isinstance(glob_str, str) and (glob_str.startswith(("gs://", "s3://"))): + # Provide a better error message for a common misuse. + wandb.termlog(f"{glob_str} is a cloud storage url, can't save file to W&B.") + return [] + # NOTE: We use PurePath instead of Path because WindowsPath doesn't + # like asterisks and errors out in resolve(). It also makes logical + # sense: globs aren't real paths, they're just path-like strings. + glob_path = pathlib.PurePath(glob_str) + resolved_glob_path = pathlib.PurePath(os.path.abspath(glob_path)) + + if base_path is not None: + base_path = pathlib.Path(base_path) + elif not glob_path.is_absolute(): + base_path = pathlib.Path(".") + else: + # Absolute glob paths with no base path get special handling. + wandb.termwarn( + "Saving files without folders. If you want to preserve " + "subdirectories pass base_path to wandb.save, i.e. " + 'wandb.save("/mnt/folder/file.h5", base_path="/mnt")', + repeat=False, + ) + base_path = resolved_glob_path.parent.parent + + if policy not in ("live", "end", "now"): + raise ValueError( + 'Only "live", "end" and "now" policies are currently supported.' + ) + + resolved_base_path = pathlib.PurePath(os.path.abspath(base_path)) + + return self._save( + resolved_glob_path, + resolved_base_path, + policy, + ) + + def _save( + self, + glob_path: pathlib.PurePath, + base_path: pathlib.PurePath, + policy: PolicyName, + ) -> list[str]: + # Can't use is_relative_to() because that's added in Python 3.9, + # but we support down to Python 3.8. + if not str(glob_path).startswith(str(base_path)): + raise ValueError("Glob may not walk above the base path") + + if glob_path == base_path: + raise ValueError("Glob cannot be the same as the base path") + + relative_glob = glob_path.relative_to(base_path) + if relative_glob.parts[0] == "*": + raise ValueError("Glob may not start with '*' relative to the base path") + relative_glob_str = GlobStr(str(relative_glob)) + + with telemetry.context(run=self) as tel: + tel.feature.save = True + + # Files in the files directory matched by the glob, including old and + # new ones. + globbed_files = set( + pathlib.Path( + self._settings.files_dir, + ).glob(relative_glob_str) + ) + + had_symlinked_files = len(globbed_files) > 0 + is_star_glob = "*" in relative_glob_str + + # The base_path may itself be a glob, so we can't do + # base_path.glob(relative_glob_str) + for path_str in glob.glob(str(base_path / relative_glob_str)): + source_path = pathlib.Path(path_str).absolute() + + # We can't use relative_to() because base_path may be a glob. + relative_path = pathlib.Path(*source_path.parts[len(base_path.parts) :]) + + target_path = pathlib.Path(self._settings.files_dir, relative_path) + globbed_files.add(target_path) + + # If the file is already where it needs to be, don't create a symlink. + if source_path.resolve() == target_path.resolve(): + continue + + target_path.parent.mkdir(parents=True, exist_ok=True) + + # Delete the symlink if it exists. + target_path.unlink(missing_ok=True) + + target_path.symlink_to(source_path) + + # Inform users that new files aren't detected automatically. + if not had_symlinked_files and is_star_glob: + file_str = f"{len(globbed_files)} file" + if len(globbed_files) > 1: + file_str += "s" + wandb.termwarn( + f"Symlinked {file_str} into the W&B run directory, " + "call wandb.save again to sync new files." + ) + + files_dict: FilesDict = { + "files": [ + ( + GlobStr(str(f.relative_to(self._settings.files_dir))), + policy, + ) + for f in globbed_files + ] + } + if self._backend and self._backend.interface: + self._backend.interface.publish_files(files_dict) + + return [str(f) for f in globbed_files] + + @_log_to_run + @_attach + def restore( + self, + name: str, + run_path: str | None = None, + replace: bool = False, + root: str | None = None, + ) -> None | TextIO: + return restore( + name, + run_path or self._get_path(), + replace, + root or self._settings.files_dir, + ) + + @_log_to_run + @_attach + def finish( + self, + exit_code: int | None = None, + quiet: bool | None = None, + ) -> None: + """Finish a run and upload any remaining data. + + Marks the completion of a W&B run and ensures all data is synced to the server. + The run's final state is determined by its exit conditions and sync status. + + Run States: + - Running: Active run that is logging data and/or sending heartbeats. + - Crashed: Run that stopped sending heartbeats unexpectedly. + - Finished: Run completed successfully (`exit_code=0`) with all data synced. + - Failed: Run completed with errors (`exit_code!=0`). + + Args: + exit_code: Integer indicating the run's exit status. Use 0 for success, + any other value marks the run as failed. + quiet: Deprecated. Configure logging verbosity using `wandb.Settings(quiet=...)`. + """ + if quiet is not None: + deprecate.deprecate( + field_name=Deprecated.run__finish_quiet, + warning_message=( + "The `quiet` argument to `wandb.run.finish()` is deprecated, " + "use `wandb.Settings(quiet=...)` to set this instead." + ), + run=self, + ) + return self._finish(exit_code) + + @_log_to_run + def _finish( + self, + exit_code: int | None = None, + ) -> None: + if self._is_finished: + return + + assert self._wl + + logger.info(f"finishing run {self._get_path()}") + with telemetry.context(run=self) as tel: + tel.feature.finish = True + + # Run hooks that need to happen before the last messages to the + # internal service, like Jupyter hooks. + for hook in self._teardown_hooks: + if hook.stage == TeardownStage.EARLY: + hook.call() + + # Early-stage hooks may use methods that require _is_finished + # to be False, so we set this after running those hooks. + self._is_finished = True + self._wl.remove_active_run(self) + + try: + self._atexit_cleanup(exit_code=exit_code) + + # Run hooks that should happen after the last messages to the + # internal service, like detaching the logger. + for hook in self._teardown_hooks: + if hook.stage == TeardownStage.LATE: + hook.call() + self._teardown_hooks = [] + + # Inform the service that we're done sending messages for this run. + # + # TODO: Why not do this in _atexit_cleanup()? + if self._settings.run_id: + service = self._wl.assert_service() + service.inform_finish(run_id=self._settings.run_id) + + finally: + if wandb.run is self: + module.unset_globals() + wandb._sentry.end_session() + + @_log_to_run + @_raise_if_finished + @_attach + def status( + self, + ) -> RunStatus: + """Get sync info from the internal backend, about the current run's sync status.""" + if not self._backend or not self._backend.interface: + return RunStatus() + + handle_run_status = self._backend.interface.deliver_request_run_status() + result = handle_run_status.wait_or(timeout=None) + sync_data = result.response.run_status_response + + sync_time = None + if sync_data.sync_time.seconds: + sync_time = datetime.fromtimestamp( + sync_data.sync_time.seconds + sync_data.sync_time.nanos / 1e9 + ) + return RunStatus( + sync_items_total=sync_data.sync_items_total, + sync_items_pending=sync_data.sync_items_pending, + sync_time=sync_time, + ) + + def _add_panel( + self, visualize_key: str, panel_type: str, panel_config: dict + ) -> None: + config = { + "panel_type": panel_type, + "panel_config": panel_config, + } + self._config_callback(val=config, key=("_wandb", "visualize", visualize_key)) + + def _redirect( + self, + stdout_slave_fd: int | None, + stderr_slave_fd: int | None, + console: str | None = None, + ) -> None: + if console is None: + console = self._settings.console + # only use raw for service to minimize potential changes + if console == "wrap": + console = "wrap_raw" + logger.info("redirect: %s", console) + + out_redir: redirect.RedirectBase + err_redir: redirect.RedirectBase + + # raw output handles the output_log writing in the internal process + if console in {"redirect", "wrap_emu"}: + output_log_path = os.path.join( + self._settings.files_dir, filenames.OUTPUT_FNAME + ) + # output writer might have been set up, see wrap_fallback case + if not self._output_writer: + self._output_writer = filesystem.CRDedupedFile( + open(output_log_path, "wb") + ) + + if console == "redirect": + logger.info("Redirecting console.") + out_redir = redirect.Redirect( + src="stdout", + cbs=[ + lambda data: self._console_callback("stdout", data), + self._output_writer.write, # type: ignore + ], + flush_periodically=(self._settings.mode == "online"), + ) + err_redir = redirect.Redirect( + src="stderr", + cbs=[ + lambda data: self._console_callback("stderr", data), + self._output_writer.write, # type: ignore + ], + flush_periodically=(self._settings.mode == "online"), + ) + if os.name == "nt": + + def wrap_fallback() -> None: + if self._out_redir: + self._out_redir.uninstall() + if self._err_redir: + self._err_redir.uninstall() + msg = ( + "Tensorflow detected. Stream redirection is not supported " + "on Windows when tensorflow is imported. Falling back to " + "wrapping stdout/err." + ) + wandb.termlog(msg) + self._redirect(None, None, console="wrap") + + add_import_hook("tensorflow", wrap_fallback) + elif console == "wrap_emu": + logger.info("Wrapping output streams.") + out_redir = redirect.StreamWrapper( + src="stdout", + cbs=[ + lambda data: self._console_callback("stdout", data), + self._output_writer.write, # type: ignore + ], + flush_periodically=(self._settings.mode == "online"), + ) + err_redir = redirect.StreamWrapper( + src="stderr", + cbs=[ + lambda data: self._console_callback("stderr", data), + self._output_writer.write, # type: ignore + ], + flush_periodically=(self._settings.mode == "online"), + ) + elif console == "wrap_raw": + logger.info("Wrapping output streams.") + out_redir = redirect.StreamRawWrapper( + src="stdout", + cbs=[ + lambda data: self._console_raw_callback("stdout", data), + ], + ) + err_redir = redirect.StreamRawWrapper( + src="stderr", + cbs=[ + lambda data: self._console_raw_callback("stderr", data), + ], + ) + elif console == "off": + return + else: + raise ValueError("unhandled console") + try: + # save stdout and stderr before installing new write functions + out_redir.install() + err_redir.install() + self._out_redir = out_redir + self._err_redir = err_redir + logger.info("Redirects installed.") + except Exception as e: + wandb.termwarn(f"Failed to redirect: {e}") + logger.exception("Failed to redirect.") + return + + def _restore(self) -> None: + logger.info("restore") + # TODO(jhr): drain and shutdown all threads + if self._out_redir: + self._out_redir.uninstall() + if self._err_redir: + self._err_redir.uninstall() + logger.info("restore done") + + def _atexit_cleanup(self, exit_code: int | None = None) -> None: + if self._backend is None: + logger.warning("process exited without backend configured") + return + if self._atexit_cleanup_called: + return + self._atexit_cleanup_called = True + + exit_code = exit_code or (self._hooks and self._hooks.exit_code) or 0 + self._exit_code = exit_code + logger.info(f"got exitcode: {exit_code}") + + # Delete this run's "resume" file if the run finished successfully. + # + # This is used by the "auto" resume mode, which resumes from the last + # failed (or unfinished/crashed) run. If we reach this line, then this + # run shouldn't be a candidate for "auto" resume. + if exit_code == 0: + if os.path.exists(self._settings.resume_fname): + os.remove(self._settings.resume_fname) + + try: + self._on_finish() + + except KeyboardInterrupt: + if not wandb.wandb_agent._is_running(): # type: ignore + wandb.termerror("Control-C detected -- Run data was not synced") + raise + + except Exception: + self._console_stop() + logger.exception("Problem finishing run") + wandb.termerror("Problem finishing run") + raise + + Run._footer( + sampled_history=self._sampled_history, + final_summary=self._final_summary, + poll_exit_response=self._poll_exit_response, + internal_messages_response=self._internal_messages_response, + settings=self._settings, + printer=self._printer, + ) + + def _console_start(self) -> None: + logger.info("atexit reg") + self._hooks = ExitHooks() + + self._redirect(self._stdout_slave_fd, self._stderr_slave_fd) + + def _console_stop(self) -> None: + self._restore() + if self._output_writer: + self._output_writer.close() + self._output_writer = None + + def _on_start(self) -> None: + self._header() + + if self._settings.save_code and self._settings.code_dir is not None: + self.log_code(self._settings.code_dir) + + if self._settings.x_save_requirements: + if self._backend and self._backend.interface: + from wandb.util import working_set + + logger.debug( + "Saving list of pip packages installed into the current environment" + ) + self._backend.interface.publish_python_packages(working_set()) + + if self._backend and self._backend.interface and not self._settings._offline: + assert self._settings.run_id + self._run_status_checker = RunStatusChecker( + self._settings.run_id, + interface=self._backend.interface, + settings=self._settings, + ) + self._run_status_checker.start() + + self._console_start() + self._on_ready() + + def _on_attach(self) -> None: + """Event triggered when run is attached to another run.""" + with telemetry.context(run=self) as tel: + tel.feature.attach = True + + self._is_attached = True + self._on_ready() + + def _register_telemetry_import_hooks( + self, + ) -> None: + def _telemetry_import_hook( + run: Run, + module: Any, + ) -> None: + with telemetry.context(run=run) as tel: + try: + name = getattr(module, "__name__", None) + if name is not None: + setattr(tel.imports_finish, name, True) + except AttributeError: + return + + import_telemetry_set = telemetry.list_telemetry_imports() + import_hook_fn = functools.partial(_telemetry_import_hook, self) + if not self._settings.run_id: + return + for module_name in import_telemetry_set: + register_post_import_hook( + import_hook_fn, + self._settings.run_id, + module_name, + ) + + def _on_ready(self) -> None: + """Event triggered when run is ready for the user.""" + assert self._wl + self._wl.add_active_run(self) + + self._register_telemetry_import_hooks() + + # start reporting any telemetry changes + self._telemetry_obj_active = True + self._telemetry_flush() + + try: + self._detect_and_apply_job_inputs() + except Exception: + logger.exception("Problem applying launch job inputs") + + # object is about to be returned to the user, don't let them modify it + self._freeze() + + if not self._settings.resume: + if os.path.exists(self._settings.resume_fname): + os.remove(self._settings.resume_fname) + + def _detect_and_apply_job_inputs(self) -> None: + """If the user has staged launch inputs, apply them to the run.""" + from wandb.sdk.launch.inputs.internal import StagedLaunchInputs + + StagedLaunchInputs().apply(self) + + def _make_job_source_reqs(self) -> tuple[list[str], dict[str, Any], dict[str, Any]]: + from wandb.util import working_set + + installed_packages_list = sorted(f"{d.key}=={d.version}" for d in working_set()) + input_types = TypeRegistry.type_of(self.config.as_dict()).to_json() + output_types = TypeRegistry.type_of(self.summary._as_dict()).to_json() + + return installed_packages_list, input_types, output_types + + def _construct_job_artifact( + self, + name: str, + source_dict: JobSourceDict, + installed_packages_list: list[str], + patch_path: os.PathLike | None = None, + ) -> Artifact: + job_artifact = InternalArtifact(name, job_builder.JOB_ARTIFACT_TYPE) + if patch_path and os.path.exists(patch_path): + job_artifact.add_file(FilePathStr(str(patch_path)), "diff.patch") + with job_artifact.new_file("requirements.frozen.txt") as f: + f.write("\n".join(installed_packages_list)) + with job_artifact.new_file("wandb-job.json") as f: + f.write(json.dumps(source_dict)) + + return job_artifact + + def _create_image_job( + self, + input_types: dict[str, Any], + output_types: dict[str, Any], + installed_packages_list: list[str], + docker_image_name: str | None = None, + args: list[str] | None = None, + ) -> Artifact | None: + docker_image_name = docker_image_name or os.getenv("WANDB_DOCKER") + + if not docker_image_name: + return None + + name = wandb.util.make_artifact_name_safe(f"job-{docker_image_name}") + s_args: Sequence[str] = args if args is not None else self._settings._args + source_info: JobSourceDict = { + "_version": "v0", + "source_type": "image", + "source": {"image": docker_image_name, "args": s_args}, + "input_types": input_types, + "output_types": output_types, + "runtime": self._settings._python, + } + job_artifact = self._construct_job_artifact( + name, source_info, installed_packages_list + ) + + return job_artifact + + def _log_job_artifact_with_image( + self, docker_image_name: str, args: list[str] | None = None + ) -> Artifact: + packages, in_types, out_types = self._make_job_source_reqs() + job_artifact = self._create_image_job( + in_types, + out_types, + packages, + args=args, + docker_image_name=docker_image_name, + ) + + assert job_artifact + artifact = self.log_artifact(job_artifact) + + if not artifact: + raise wandb.Error(f"Job Artifact log unsuccessful: {artifact}") + else: + return artifact + + async def _display_finish_stats( + self, + progress_printer: progress.ProgressPrinter, + ) -> None: + last_result: Result | None = None + + async def loop_update_printer() -> None: + while True: + if last_result: + progress_printer.update( + [last_result.response.poll_exit_response], + ) + await asyncio.sleep(0.1) + + async def loop_poll_exit() -> None: + nonlocal last_result + assert self._backend and self._backend.interface + + while True: + handle = self._backend.interface.deliver_poll_exit() + + time_start = time.monotonic() + last_result = await handle.wait_async(timeout=None) + + # Update at most once a second. + time_elapsed = time.monotonic() - time_start + if time_elapsed < 1: + await asyncio.sleep(1 - time_elapsed) + + async with asyncio_compat.open_task_group() as task_group: + task_group.start_soon(loop_update_printer()) + task_group.start_soon(loop_poll_exit()) + + def _on_finish(self) -> None: + trigger.call("on_finished") + + if self._run_status_checker is not None: + self._run_status_checker.stop() + + self._console_stop() # TODO: there's a race here with jupyter console logging + + assert self._backend and self._backend.interface + + if self._settings.x_update_finish_state: + exit_handle = self._backend.interface.deliver_exit(self._exit_code) + else: + exit_handle = self._backend.interface.deliver_finish_without_exit() + + with progress.progress_printer( + self._printer, + default_text="Finishing up...", + ) as progress_printer: + # Wait for the run to complete. + wait_with_progress( + exit_handle, + timeout=None, + progress_after=1, + display_progress=functools.partial( + self._display_finish_stats, + progress_printer, + ), + ) + + # Print some final statistics. + poll_exit_handle = self._backend.interface.deliver_poll_exit() + result = poll_exit_handle.wait_or(timeout=None) + progress.print_sync_dedupe_stats( + self._printer, + result.response.poll_exit_response, + ) + + self._poll_exit_response = result.response.poll_exit_response + internal_messages_handle = self._backend.interface.deliver_internal_messages() + result = internal_messages_handle.wait_or(timeout=None) + self._internal_messages_response = result.response.internal_messages_response + + # dispatch all our final requests + + final_summary_handle = self._backend.interface.deliver_get_summary() + sampled_history_handle = ( + self._backend.interface.deliver_request_sampled_history() + ) + + result = sampled_history_handle.wait_or(timeout=None) + self._sampled_history = result.response.sampled_history_response + + result = final_summary_handle.wait_or(timeout=None) + self._final_summary = result.response.get_summary_response + + if self._backend: + self._backend.cleanup() + + if self._run_status_checker: + self._run_status_checker.join() + + if self._settings.run_id: + self._unregister_telemetry_import_hooks(self._settings.run_id) + + @staticmethod + def _unregister_telemetry_import_hooks(run_id: str) -> None: + import_telemetry_set = telemetry.list_telemetry_imports() + for module_name in import_telemetry_set: + unregister_post_import_hook(module_name, run_id) + + @_log_to_run + @_raise_if_finished + @_attach + def define_metric( + self, + name: str, + step_metric: str | wandb_metric.Metric | None = None, + step_sync: bool | None = None, + hidden: bool | None = None, + summary: str | None = None, + goal: str | None = None, + overwrite: bool | None = None, + ) -> wandb_metric.Metric: + """Customize metrics logged with `wandb.log()`. + + Args: + name: The name of the metric to customize. + step_metric: The name of another metric to serve as the X-axis + for this metric in automatically generated charts. + step_sync: Automatically insert the last value of step_metric into + `run.log()` if it is not provided explicitly. Defaults to True + if step_metric is specified. + hidden: Hide this metric from automatic plots. + summary: Specify aggregate metrics added to summary. + Supported aggregations include "min", "max", "mean", "last", + "best", "copy" and "none". "best" is used together with the + goal parameter. "none" prevents a summary from being generated. + "copy" is deprecated and should not be used. + goal: Specify how to interpret the "best" summary type. + Supported options are "minimize" and "maximize". + overwrite: If false, then this call is merged with previous + `define_metric` calls for the same metric by using their + values for any unspecified parameters. If true, then + unspecified parameters overwrite values specified by + previous calls. + + Returns: + An object that represents this call but can otherwise be discarded. + """ + if summary and "copy" in summary: + deprecate.deprecate( + Deprecated.run__define_metric_copy, + "define_metric(summary='copy') is deprecated and will be removed.", + self, + ) + + if (summary and "best" in summary) or goal is not None: + deprecate.deprecate( + Deprecated.run__define_metric_best_goal, + "define_metric(summary='best', goal=...) is deprecated and will be removed. " + "Use define_metric(summary='min') or define_metric(summary='max') instead.", + self, + ) + + return self._define_metric( + name, + step_metric, + step_sync, + hidden, + summary, + goal, + overwrite, + ) + + def _define_metric( + self, + name: str, + step_metric: str | wandb_metric.Metric | None = None, + step_sync: bool | None = None, + hidden: bool | None = None, + summary: str | None = None, + goal: str | None = None, + overwrite: bool | None = None, + ) -> wandb_metric.Metric: + if not name: + raise wandb.Error("define_metric() requires non-empty name argument") + if isinstance(step_metric, wandb_metric.Metric): + step_metric = step_metric.name + for arg_name, arg_val, exp_type in ( + ("name", name, str), + ("step_metric", step_metric, str), + ("step_sync", step_sync, bool), + ("hidden", hidden, bool), + ("summary", summary, str), + ("goal", goal, str), + ("overwrite", overwrite, bool), + ): + # NOTE: type checking is broken for isinstance and str + if arg_val is not None and not isinstance(arg_val, exp_type): + arg_type = type(arg_val).__name__ + raise wandb.Error( + f"Unhandled define_metric() arg: {arg_name} type: {arg_type}" + ) + stripped = name[:-1] if name.endswith("*") else name + if "*" in stripped: + raise wandb.Error( + f"Unhandled define_metric() arg: name (glob suffixes only): {name}" + ) + summary_ops: Sequence[str] | None = None + if summary: + summary_items = [s.lower() for s in summary.split(",")] + summary_ops = [] + valid = {"min", "max", "mean", "best", "last", "copy", "none"} + # TODO: deprecate copy and best + for i in summary_items: + if i not in valid: + raise wandb.Error(f"Unhandled define_metric() arg: summary op: {i}") + summary_ops.append(i) + with telemetry.context(run=self) as tel: + tel.feature.metric_summary = True + # TODO: deprecate goal + goal_cleaned: str | None = None + if goal is not None: + goal_cleaned = goal[:3].lower() + valid_goal = {"min", "max"} + if goal_cleaned not in valid_goal: + raise wandb.Error(f"Unhandled define_metric() arg: goal: {goal}") + with telemetry.context(run=self) as tel: + tel.feature.metric_goal = True + if hidden: + with telemetry.context(run=self) as tel: + tel.feature.metric_hidden = True + if step_sync: + with telemetry.context(run=self) as tel: + tel.feature.metric_step_sync = True + + with telemetry.context(run=self) as tel: + tel.feature.metric = True + + m = wandb_metric.Metric( + name=name, + step_metric=step_metric, + step_sync=step_sync, + summary=summary_ops, + hidden=hidden, + goal=goal_cleaned, + overwrite=overwrite, + ) + m._set_callback(self._metric_callback) + m._commit() + return m + + @_log_to_run + @_attach + def watch( + self, + models: torch.nn.Module | Sequence[torch.nn.Module], + criterion: torch.F | None = None, # type: ignore + log: Literal["gradients", "parameters", "all"] | None = "gradients", + log_freq: int = 1000, + idx: int | None = None, + log_graph: bool = False, + ) -> None: + """Hooks into the given PyTorch model(s) to monitor gradients and the model's computational graph. + + This function can track parameters, gradients, or both during training. It should be + extended to support arbitrary machine learning models in the future. + + Args: + models (Union[torch.nn.Module, Sequence[torch.nn.Module]]): + A single model or a sequence of models to be monitored. + criterion (Optional[torch.F]): + The loss function being optimized (optional). + log (Optional[Literal["gradients", "parameters", "all"]]): + Specifies whether to log "gradients", "parameters", or "all". + Set to None to disable logging. (default="gradients") + log_freq (int): + Frequency (in batches) to log gradients and parameters. (default=1000) + idx (Optional[int]): + Index used when tracking multiple models with `wandb.watch`. (default=None) + log_graph (bool): + Whether to log the model's computational graph. (default=False) + + Raises: + ValueError: + If `wandb.init` has not been called or if any of the models are not instances + of `torch.nn.Module`. + """ + wandb.sdk._watch(self, models, criterion, log, log_freq, idx, log_graph) + + @_log_to_run + @_attach + def unwatch( + self, models: torch.nn.Module | Sequence[torch.nn.Module] | None = None + ) -> None: + """Remove pytorch model topology, gradient and parameter hooks. + + Args: + models (torch.nn.Module | Sequence[torch.nn.Module]): + Optional list of pytorch models that have had watch called on them + """ + wandb.sdk._unwatch(self, models=models) + + @_log_to_run + @_raise_if_finished + @_attach + def link_artifact( + self, + artifact: Artifact, + target_path: str, + aliases: list[str] | None = None, + ) -> Artifact | None: + """Link the given artifact to a portfolio (a promoted collection of artifacts). + + The linked artifact will be visible in the UI for the specified portfolio. + + Args: + artifact: the (public or local) artifact which will be linked + target_path: `str` - takes the following forms: `{portfolio}`, `{project}/{portfolio}`, + or `{entity}/{project}/{portfolio}` + aliases: `List[str]` - optional alias(es) that will only be applied on this linked artifact + inside the portfolio. + The alias "latest" will always be applied to the latest version of an artifact that is linked. + + Returns: + The linked artifact if linking was successful, otherwise None. + + """ + if artifact.is_draft() and not artifact._is_draft_save_started(): + artifact = self._log_artifact(artifact) + + if self._settings._offline: + # TODO: implement offline mode + sync + raise NotImplementedError + + # Normalize the target "entity/project/collection" with defaults + # inferred from this run's entity and project, if needed. + # + # HOWEVER, if the target path is a registry collection, avoid setting + # the target entity to the run's entity. Instead, delegate to + # Artifact.link() to resolve the required org entity. + target = ArtifactPath.from_str(target_path) + if not (target.project and is_artifact_registry_project(target.project)): + target = target.with_defaults(prefix=self.entity, project=self.project) + + return artifact.link(target.to_str(), aliases) + + @_log_to_run + @_raise_if_finished + @_attach + def use_artifact( + self, + artifact_or_name: str | Artifact, + type: str | None = None, + aliases: list[str] | None = None, + use_as: str | None = None, + ) -> Artifact: + """Declare an artifact as an input to a run. + + Call `download` or `file` on the returned object to get the contents locally. + + Args: + artifact_or_name: (str or Artifact) An artifact name. + May be prefixed with project/ or entity/project/. + If no entity is specified in the name, the Run or API setting's entity is used. + Valid names can be in the following forms: + - name:version + - name:alias + You can also pass an Artifact object created by calling `wandb.Artifact` + type: (str, optional) The type of artifact to use. + aliases: (list, optional) Aliases to apply to this artifact + use_as: This argument is deprecated and does nothing. + + Returns: + An `Artifact` object. + """ + if self._settings._offline: + raise TypeError("Cannot use artifact when in offline mode.") + + api = internal.Api( + default_settings={ + "entity": self._settings.entity, + "project": self._settings.project, + } + ) + api.set_current_run_id(self._settings.run_id) + + if use_as is not None: + deprecate.deprecate( + field_name=Deprecated.run__use_artifact_use_as, + warning_message=( + "`use_as` argument is deprecated and does not affect the behaviour of `run.use_artifact`" + ), + ) + + if isinstance(artifact_or_name, str): + name = artifact_or_name + public_api = self._public_api() + artifact = public_api._artifact(type=type, name=name) + if type is not None and type != artifact.type: + raise ValueError( + f"Supplied type {type} does not match type {artifact.type} of artifact {artifact.name}" + ) + api.use_artifact( + artifact.id, + entity_name=self._settings.entity, + project_name=self._settings.project, + artifact_entity_name=artifact.entity, + artifact_project_name=artifact.project, + ) + else: + artifact = artifact_or_name + if aliases is None: + aliases = [] + elif isinstance(aliases, str): + aliases = [aliases] + if isinstance(artifact_or_name, Artifact) and artifact.is_draft(): + if use_as is not None: + wandb.termwarn( + "Indicating use_as is not supported when using a draft artifact" + ) + self._log_artifact( + artifact, + aliases=aliases, + is_user_created=True, + use_after_commit=True, + ) + artifact.wait() + elif isinstance(artifact, Artifact) and not artifact.is_draft(): + api.use_artifact( + artifact.id, + artifact_entity_name=artifact.entity, + artifact_project_name=artifact.project, + ) + else: + raise ValueError( + 'You must pass an artifact name (e.g. "pedestrian-dataset:v1"), ' + "an instance of `wandb.Artifact`, or `wandb.Api().artifact()` to `use_artifact`" + ) + if self._backend and self._backend.interface: + self._backend.interface.publish_use_artifact(artifact) + return artifact + + @_log_to_run + @_raise_if_finished + @_attach + def log_artifact( + self, + artifact_or_path: Artifact | StrPath, + name: str | None = None, + type: str | None = None, + aliases: list[str] | None = None, + tags: list[str] | None = None, + ) -> Artifact: + """Declare an artifact as an output of a run. + + Args: + artifact_or_path: (str or Artifact) A path to the contents of this artifact, + can be in the following forms: + - `/local/directory` + - `/local/directory/file.txt` + - `s3://bucket/path` + You can also pass an Artifact object created by calling + `wandb.Artifact`. + name: (str, optional) An artifact name. Valid names can be in the following forms: + - name:version + - name:alias + - digest + This will default to the basename of the path prepended with the current + run id if not specified. + type: (str) The type of artifact to log, examples include `dataset`, `model` + aliases: (list, optional) Aliases to apply to this artifact, + defaults to `["latest"]` + tags: (list, optional) Tags to apply to this artifact, if any. + + Returns: + An `Artifact` object. + """ + return self._log_artifact( + artifact_or_path, + name=name, + type=type, + aliases=aliases, + tags=tags, + ) + + @_log_to_run + @_raise_if_finished + @_attach + def upsert_artifact( + self, + artifact_or_path: Artifact | str, + name: str | None = None, + type: str | None = None, + aliases: list[str] | None = None, + distributed_id: str | None = None, + ) -> Artifact: + """Declare (or append to) a non-finalized artifact as output of a run. + + Note that you must call run.finish_artifact() to finalize the artifact. + This is useful when distributed jobs need to all contribute to the same artifact. + + Args: + artifact_or_path: (str or Artifact) A path to the contents of this artifact, + can be in the following forms: + - `/local/directory` + - `/local/directory/file.txt` + - `s3://bucket/path` + You can also pass an Artifact object created by calling + `wandb.Artifact`. + name: (str, optional) An artifact name. May be prefixed with entity/project. + Valid names can be in the following forms: + - name:version + - name:alias + - digest + This will default to the basename of the path prepended with the current + run id if not specified. + type: (str) The type of artifact to log, examples include `dataset`, `model` + aliases: (list, optional) Aliases to apply to this artifact, + defaults to `["latest"]` + distributed_id: (string, optional) Unique string that all distributed jobs share. If None, + defaults to the run's group name. + + Returns: + An `Artifact` object. + """ + if self._settings.run_group is None and distributed_id is None: + raise TypeError( + "Cannot upsert artifact unless run is in a group or distributed_id is provided" + ) + if distributed_id is None: + distributed_id = self._settings.run_group or "" + return self._log_artifact( + artifact_or_path, + name=name, + type=type, + aliases=aliases, + distributed_id=distributed_id, + finalize=False, + ) + + @_log_to_run + @_raise_if_finished + @_attach + def finish_artifact( + self, + artifact_or_path: Artifact | str, + name: str | None = None, + type: str | None = None, + aliases: list[str] | None = None, + distributed_id: str | None = None, + ) -> Artifact: + """Finishes a non-finalized artifact as output of a run. + + Subsequent "upserts" with the same distributed ID will result in a new version. + + Args: + artifact_or_path: (str or Artifact) A path to the contents of this artifact, + can be in the following forms: + - `/local/directory` + - `/local/directory/file.txt` + - `s3://bucket/path` + You can also pass an Artifact object created by calling + `wandb.Artifact`. + name: (str, optional) An artifact name. May be prefixed with entity/project. + Valid names can be in the following forms: + - name:version + - name:alias + - digest + This will default to the basename of the path prepended with the current + run id if not specified. + type: (str) The type of artifact to log, examples include `dataset`, `model` + aliases: (list, optional) Aliases to apply to this artifact, + defaults to `["latest"]` + distributed_id: (string, optional) Unique string that all distributed jobs share. If None, + defaults to the run's group name. + + Returns: + An `Artifact` object. + """ + if self._settings.run_group is None and distributed_id is None: + raise TypeError( + "Cannot finish artifact unless run is in a group or distributed_id is provided" + ) + if distributed_id is None: + distributed_id = self._settings.run_group or "" + + return self._log_artifact( + artifact_or_path, + name, + type, + aliases, + distributed_id=distributed_id, + finalize=True, + ) + + def _log_artifact( + self, + artifact_or_path: Artifact | StrPath, + name: str | None = None, + type: str | None = None, + aliases: list[str] | None = None, + tags: list[str] | None = None, + distributed_id: str | None = None, + finalize: bool = True, + is_user_created: bool = False, + use_after_commit: bool = False, + ) -> Artifact: + if self._settings.anonymous in ["allow", "must"]: + wandb.termwarn( + "Artifacts logged anonymously cannot be claimed and expire after 7 days." + ) + + if not finalize and distributed_id is None: + raise TypeError("Must provide distributed_id if artifact is not finalize") + + if aliases is not None: + aliases = validate_aliases(aliases) + + # Check if artifact tags are supported + if tags is not None: + tags = validate_tags(tags) + + artifact, aliases = self._prepare_artifact( + artifact_or_path, name, type, aliases + ) + + if len(artifact.metadata) > MAX_ARTIFACT_METADATA_KEYS: + raise ValueError( + f"Artifact must not have more than {MAX_ARTIFACT_METADATA_KEYS} metadata keys." + ) + + artifact.distributed_id = distributed_id + self._assert_can_log_artifact(artifact) + if self._backend and self._backend.interface: + if not self._settings._offline: + handle = self._backend.interface.deliver_artifact( + self, + artifact, + aliases, + tags, + self.step, + finalize=finalize, + is_user_created=is_user_created, + use_after_commit=use_after_commit, + ) + artifact._set_save_handle(handle, self._public_api().client) + else: + self._backend.interface.publish_artifact( + self, + artifact, + aliases, + tags, + finalize=finalize, + is_user_created=is_user_created, + use_after_commit=use_after_commit, + ) + elif self._internal_run_interface: + self._internal_run_interface.publish_artifact( + self, + artifact, + aliases, + tags, + finalize=finalize, + is_user_created=is_user_created, + use_after_commit=use_after_commit, + ) + return artifact + + def _public_api(self, overrides: dict[str, str] | None = None) -> PublicApi: + overrides = {"run": self._settings.run_id} # type: ignore + if not self._settings._offline: + overrides["entity"] = self._settings.entity or "" + overrides["project"] = self._settings.project or "" + return public.Api(overrides) + + # TODO(jhr): annotate this + def _assert_can_log_artifact(self, artifact) -> None: # type: ignore + if self._settings._offline: + return + try: + public_api = self._public_api() + entity = public_api.settings["entity"] + project = public_api.settings["project"] + expected_type = Artifact._expected_type( + entity, project, artifact.name, public_api.client + ) + except requests.exceptions.RequestException: + # Just return early if there is a network error. This is + # ok, as this function is intended to help catch an invalid + # type early, but not a hard requirement for valid operation. + return + if expected_type is not None and artifact.type != expected_type: + raise ValueError( + f"Artifact {artifact.name} already exists with type '{expected_type}'; " + f"cannot create another with type '{artifact.type}'" + ) + if entity and artifact._source_entity and entity != artifact._source_entity: + raise ValueError( + f"Artifact {artifact.name} is owned by entity " + f"'{artifact._source_entity}'; it can't be moved to '{entity}'" + ) + if project and artifact._source_project and project != artifact._source_project: + raise ValueError( + f"Artifact {artifact.name} exists in project " + f"'{artifact._source_project}'; it can't be moved to '{project}'" + ) + + def _prepare_artifact( + self, + artifact_or_path: Artifact | StrPath, + name: str | None = None, + type: str | None = None, + aliases: list[str] | None = None, + ) -> tuple[Artifact, list[str]]: + if isinstance(artifact_or_path, (str, os.PathLike)): + name = ( + name + or f"run-{self._settings.run_id}-{os.path.basename(artifact_or_path)}" + ) + artifact = Artifact(name, type or "unspecified") + if os.path.isfile(artifact_or_path): + artifact.add_file(str(artifact_or_path)) + elif os.path.isdir(artifact_or_path): + artifact.add_dir(str(artifact_or_path)) + elif "://" in str(artifact_or_path): + artifact.add_reference(str(artifact_or_path)) + else: + raise ValueError( + "path must be a file, directory or external" + "reference like s3://bucket/path" + ) + else: + artifact = artifact_or_path + if not isinstance(artifact, Artifact): + raise TypeError( + "You must pass an instance of wandb.Artifact or a " + "valid file path to log_artifact" + ) + + artifact.finalize() + return artifact, _resolve_aliases(aliases) + + @_log_to_run + @_raise_if_finished + @_attach + def log_model( + self, + path: StrPath, + name: str | None = None, + aliases: list[str] | None = None, + ) -> None: + """Logs a model artifact containing the contents inside the 'path' to a run and marks it as an output to this run. + + Args: + path: (str) A path to the contents of this model, + can be in the following forms: + - `/local/directory` + - `/local/directory/file.txt` + - `s3://bucket/path` + name: (str, optional) A name to assign to the model artifact that the file contents will be added to. + The string must contain only the following alphanumeric characters: dashes, underscores, and dots. + This will default to the basename of the path prepended with the current + run id if not specified. + aliases: (list, optional) Aliases to apply to the created model artifact, + defaults to `["latest"]` + + Examples: + ```python + run.log_model( + path="/local/directory", + name="my_model_artifact", + aliases=["production"], + ) + ``` + + Invalid usage + ```python + run.log_model( + path="/local/directory", + name="my_entity/my_project/my_model_artifact", + aliases=["production"], + ) + ``` + + Raises: + ValueError: if name has invalid special characters + + Returns: + None + """ + self._log_artifact( + artifact_or_path=path, name=name, type="model", aliases=aliases + ) + + @_log_to_run + @_raise_if_finished + @_attach + def use_model(self, name: str) -> FilePathStr: + """Download the files logged in a model artifact 'name'. + + Args: + name: (str) A model artifact name. 'name' must match the name of an existing logged + model artifact. + May be prefixed with entity/project/. Valid names + can be in the following forms: + - model_artifact_name:version + - model_artifact_name:alias + + Examples: + ```python + run.use_model( + name="my_model_artifact:latest", + ) + + run.use_model( + name="my_project/my_model_artifact:v0", + ) + + run.use_model( + name="my_entity/my_project/my_model_artifact:", + ) + ``` + + Invalid usage + ```python + run.use_model( + name="my_entity/my_project/my_model_artifact", + ) + ``` + + Raises: + AssertionError: if model artifact 'name' is of a type that does not contain the substring 'model'. + + Returns: + path: (str) path to downloaded model artifact file(s). + """ + if self._settings._offline: + # Downloading artifacts is not supported when offline. + raise RuntimeError("`use_model` not supported in offline mode.") + + artifact = self.use_artifact(artifact_or_name=name) + if "model" not in str(artifact.type.lower()): + raise AssertionError( + "You can only use this method for 'model' artifacts." + " For an artifact to be a 'model' artifact, its type property" + " must contain the substring 'model'." + ) + + path = artifact.download() + + # If returned directory contains only one file, return path to that file + dir_list = os.listdir(path) + if len(dir_list) == 1: + return FilePathStr(os.path.join(path, dir_list[0])) + return path + + @_log_to_run + @_raise_if_finished + @_attach + def link_model( + self, + path: StrPath, + registered_model_name: str, + name: str | None = None, + aliases: list[str] | None = None, + ) -> Artifact | None: + """Log a model artifact version and link it to a registered model in the model registry. + + The linked model version will be visible in the UI for the specified registered model. + + Steps: + - Check if 'name' model artifact has been logged. If so, use the artifact version that matches the files + located at 'path' or log a new version. Otherwise log files under 'path' as a new model artifact, 'name' + of type 'model'. + - Check if registered model with name 'registered_model_name' exists in the 'model-registry' project. + If not, create a new registered model with name 'registered_model_name'. + - Link version of model artifact 'name' to registered model, 'registered_model_name'. + - Attach aliases from 'aliases' list to the newly linked model artifact version. + + Args: + path: (str) A path to the contents of this model, + can be in the following forms: + - `/local/directory` + - `/local/directory/file.txt` + - `s3://bucket/path` + registered_model_name: (str) - the name of the registered model that the model is to be linked to. + A registered model is a collection of model versions linked to the model registry, typically representing a + team's specific ML Task. The entity that this registered model belongs to will be derived from the run + name: (str, optional) - the name of the model artifact that files in 'path' will be logged to. This will + default to the basename of the path prepended with the current run id if not specified. + aliases: (List[str], optional) - alias(es) that will only be applied on this linked artifact + inside the registered model. + The alias "latest" will always be applied to the latest version of an artifact that is linked. + + Examples: + ```python + run.link_model( + path="/local/directory", + registered_model_name="my_reg_model", + name="my_model_artifact", + aliases=["production"], + ) + ``` + + Invalid usage + ```python + run.link_model( + path="/local/directory", + registered_model_name="my_entity/my_project/my_reg_model", + name="my_model_artifact", + aliases=["production"], + ) + + run.link_model( + path="/local/directory", + registered_model_name="my_reg_model", + name="my_entity/my_project/my_model_artifact", + aliases=["production"], + ) + ``` + + Raises: + AssertionError: if registered_model_name is a path or + if model artifact 'name' is of a type that does not contain the substring 'model' + ValueError: if name has invalid special characters + + Returns: + The linked artifact if linking was successful, otherwise None. + """ + name_parts = registered_model_name.split("/") + if len(name_parts) != 1: + raise AssertionError( + "Please provide only the name of the registered model." + " Do not append the entity or project name." + ) + + project = "model-registry" + target_path = self.entity + "/" + project + "/" + registered_model_name + + public_api = self._public_api() + try: + artifact = public_api._artifact(name=f"{name}:latest") + if "model" not in str(artifact.type.lower()): + raise AssertionError( + "You can only use this method for 'model' artifacts." + " For an artifact to be a 'model' artifact, its type" + " property must contain the substring 'model'." + ) + + artifact = self._log_artifact( + artifact_or_path=path, name=name, type=artifact.type + ) + except (ValueError, CommError): + artifact = self._log_artifact( + artifact_or_path=path, name=name, type="model" + ) + return self.link_artifact( + artifact=artifact, target_path=target_path, aliases=aliases + ) + + @_log_to_run + @_raise_if_finished + @_attach + def alert( + self, + title: str, + text: str, + level: str | AlertLevel | None = None, + wait_duration: int | float | timedelta | None = None, + ) -> None: + """Launch an alert with the given title and text. + + Args: + title: (str) The title of the alert, must be less than 64 characters long. + text: (str) The text body of the alert. + level: (str or AlertLevel, optional) The alert level to use, either: `INFO`, `WARN`, or `ERROR`. + wait_duration: (int, float, or timedelta, optional) The time to wait (in seconds) before sending another + alert with this title. + """ + level = level or AlertLevel.INFO + level_str: str = level.value if isinstance(level, AlertLevel) else level + if level_str not in {lev.value for lev in AlertLevel}: + raise ValueError("level must be one of 'INFO', 'WARN', or 'ERROR'") + + wait_duration = wait_duration or timedelta(minutes=1) + if isinstance(wait_duration, int) or isinstance(wait_duration, float): + wait_duration = timedelta(seconds=wait_duration) + elif not callable(getattr(wait_duration, "total_seconds", None)): + raise TypeError( + "wait_duration must be an int, float, or datetime.timedelta" + ) + wait_duration = int(wait_duration.total_seconds() * 1000) + + if self._backend and self._backend.interface: + self._backend.interface.publish_alert(title, text, level_str, wait_duration) + + def __enter__(self) -> Run: + return self + + def __exit__( + self, + exc_type: type[BaseException], + exc_val: BaseException, + exc_tb: TracebackType, + ) -> bool: + exception_raised = exc_type is not None + if exception_raised: + traceback.print_exception(exc_type, exc_val, exc_tb) + exit_code = 1 if exception_raised else 0 + self._finish(exit_code=exit_code) + return not exception_raised + + @_log_to_run + @_raise_if_finished + @_attach + def mark_preempting(self) -> None: + """Mark this run as preempting. + + Also tells the internal process to immediately report this to server. + """ + if self._backend and self._backend.interface: + self._backend.interface.publish_preempting() + + @property + @_log_to_run + @_raise_if_finished + @_attach + def _system_metrics(self) -> dict[str, list[tuple[datetime, float]]]: + """Returns a dictionary of system metrics. + + Returns: + A dictionary of system metrics. + """ + from wandb.proto import wandb_internal_pb2 + + def pb_to_dict( + system_metrics_pb: wandb_internal_pb2.GetSystemMetricsResponse, + ) -> dict[str, list[tuple[datetime, float]]]: + res = {} + + for metric, records in system_metrics_pb.system_metrics.items(): + measurements = [] + for record in records.record: + # Convert timestamp to datetime + dt = datetime.fromtimestamp( + record.timestamp.seconds, tz=timezone.utc + ) + dt = dt.replace(microsecond=record.timestamp.nanos // 1000) + + measurements.append((dt, record.value)) + + res[metric] = measurements + + return res + + if not self._backend or not self._backend.interface: + return {} + + handle = self._backend.interface.deliver_get_system_metrics() + + try: + result = handle.wait_or(timeout=1) + except TimeoutError: + return {} + else: + try: + response = result.response.get_system_metrics_response + return pb_to_dict(response) if response else {} + except Exception: + logger.exception("Error getting system metrics.") + return {} + + # ------------------------------------------------------------------------------ + # HEADER + # ------------------------------------------------------------------------------ + def _header(self) -> None: + self._header_wandb_version_info() + self._header_sync_info() + self._header_run_info() + + def _header_wandb_version_info(self) -> None: + if self._settings.quiet or self._settings.silent: + return + + # TODO: add this to a higher verbosity level + self._printer.display(f"Tracking run with wandb version {wandb.__version__}") + + def _header_sync_info(self) -> None: + if self._settings._offline: + self._printer.display( + [ + f"W&B syncing is set to {self._printer.code('`offline`')}" + f" in this directory. Run {self._printer.code('`wandb online`')}" + f" or set {self._printer.code('WANDB_MODE=online')}" + " to enable cloud syncing.", + ] + ) + else: + sync_dir = self._settings.sync_dir + info = [f"Run data is saved locally in {self._printer.files(sync_dir)}"] + if not self._printer.supports_html: + info.append( + f"Run {self._printer.code('`wandb offline`')} to turn off syncing." + ) + if not self._settings.quiet and not self._settings.silent: + self._printer.display(info) + + def _header_run_info(self) -> None: + settings, printer = self._settings, self._printer + + if settings._offline or settings.silent: + return + + run_url = settings.run_url + project_url = settings.project_url + sweep_url = settings.sweep_url + + run_state_str = ( + "Resuming run" + if settings.resumed or settings.resume_from + else "Syncing run" + ) + run_name = settings.run_name + if not run_name: + return + + if printer.supports_html: + import wandb.jupyter + + if not wandb.jupyter.display_if_magic_is_used(self): + run_line = f"{printer.link(run_url, run_name)}" + project_line, sweep_line = "", "" + + if not settings.quiet: + doc_html = printer.link(url_registry.url("developer-guide"), "docs") + + project_html = printer.link(project_url, "Weights & Biases") + project_line = f"to {project_html} ({doc_html})" + + if sweep_url: + sweep_line = f"Sweep page: {printer.link(sweep_url, sweep_url)}" + + printer.display( + [f"{run_state_str} {run_line} {project_line}", sweep_line], + ) + + elif run_name: + printer.display(f"{run_state_str} {printer.name(run_name)}") + + if not settings.quiet: + # TODO: add verbosity levels and add this to higher levels + printer.display( + f"{printer.emoji('star')} View project at {printer.link(project_url)}" + ) + if sweep_url: + printer.display( + f"{printer.emoji('broom')} View sweep at {printer.link(sweep_url)}" + ) + printer.display( + f"{printer.emoji('rocket')} View run at {printer.link(run_url)}", + ) + + if run_name and settings.anonymous in ["allow", "must"]: + printer.display( + ( + "Do NOT share these links with anyone." + " They can be used to claim your runs." + ), + level="warn", + ) + + # ------------------------------------------------------------------------------ + # FOOTER + # ------------------------------------------------------------------------------ + # Note: All the footer methods are static methods since we want to share the printing logic + # with the service execution path that doesn't have access to the run instance + @staticmethod + def _footer( + sampled_history: SampledHistoryResponse | None = None, + final_summary: GetSummaryResponse | None = None, + poll_exit_response: PollExitResponse | None = None, + internal_messages_response: InternalMessagesResponse | None = None, + *, + settings: Settings, + printer: printer.Printer, + ) -> None: + Run._footer_history_summary_info( + history=sampled_history, + summary=final_summary, + settings=settings, + printer=printer, + ) + + Run._footer_sync_info( + poll_exit_response=poll_exit_response, + settings=settings, + printer=printer, + ) + Run._footer_log_dir_info(settings=settings, printer=printer) + Run._footer_internal_messages( + internal_messages_response=internal_messages_response, + settings=settings, + printer=printer, + ) + + @staticmethod + def _footer_sync_info( + poll_exit_response: PollExitResponse | None = None, + *, + settings: Settings, + printer: printer.Printer, + ) -> None: + if settings.silent: + return + + if settings._offline: + if not settings.quiet: + printer.display( + [ + "You can sync this run to the cloud by running:", + printer.code(f"wandb sync {settings.sync_dir}"), + ], + ) + return + + info = [] + if settings.run_name and settings.run_url: + info.append( + f"{printer.emoji('rocket')} View run {printer.name(settings.run_name)} at: {printer.link(settings.run_url)}" + ) + if settings.project_url: + info.append( + f"{printer.emoji('star')} View project at: {printer.link(settings.project_url)}" + ) + if poll_exit_response and poll_exit_response.file_counts: + logger.info("logging synced files") + file_counts = poll_exit_response.file_counts + info.append( + f"Synced {file_counts.wandb_count} W&B file(s), {file_counts.media_count} media file(s), " + f"{file_counts.artifact_count} artifact file(s) and {file_counts.other_count} other file(s)", + ) + printer.display(info) + + @staticmethod + def _footer_log_dir_info( + *, + settings: Settings, + printer: printer.Printer, + ) -> None: + if settings.quiet or settings.silent: + return + + log_dir = settings.log_user or settings.log_internal + if log_dir: + log_dir = os.path.dirname(log_dir.replace(os.getcwd(), ".")) + printer.display( + f"Find logs at: {printer.files(log_dir)}", + ) + + @staticmethod + def _footer_history_summary_info( + history: SampledHistoryResponse | None = None, + summary: GetSummaryResponse | None = None, + *, + settings: Settings, + printer: printer.Printer, + ) -> None: + if settings.quiet or settings.silent: + return + + panel = [] + + # Render history if available + if history: + logger.info("rendering history") + + sampled_history = { + item.key: wandb.util.downsample( + item.values_float or item.values_int, 40 + ) + for item in history.item + if not item.key.startswith("_") + } + + history_rows = [] + for key, values in sorted(sampled_history.items()): + if any(not isinstance(value, numbers.Number) for value in values): + continue + sparkline = printer.sparklines(values) + if sparkline: + history_rows.append([key, sparkline]) + if history_rows: + history_grid = printer.grid( + history_rows, + "Run history:", + ) + panel.append(history_grid) + + # Render summary if available + if summary: + final_summary = {} + for item in summary.item: + if item.key.startswith("_") or len(item.nested_key) > 0: + continue + final_summary[item.key] = json.loads(item.value_json) + + logger.info("rendering summary") + summary_rows = [] + for key, value in sorted(final_summary.items()): + # arrays etc. might be too large. for now, we just don't print them + if isinstance(value, str): + value = value[:20] + "..." * (len(value) >= 20) + summary_rows.append([key, value]) + elif isinstance(value, numbers.Number): + value = round(value, 5) if isinstance(value, float) else value + summary_rows.append([key, str(value)]) + else: + continue + + if summary_rows: + summary_grid = printer.grid( + summary_rows, + "Run summary:", + ) + panel.append(summary_grid) + + if panel: + printer.display(printer.panel(panel)) + + @staticmethod + def _footer_internal_messages( + internal_messages_response: InternalMessagesResponse | None = None, + *, + settings: Settings, + printer: printer.Printer, + ) -> None: + if settings.quiet or settings.silent: + return + + if not internal_messages_response: + return + + for message in internal_messages_response.messages.warning: + printer.display(message, level="warn") + + +# We define this outside of the run context to support restoring before init +def restore( + name: str, + run_path: str | None = None, + replace: bool = False, + root: str | None = None, +) -> None | TextIO: + """Download the specified file from cloud storage. + + File is placed into the current directory or run directory. + By default, will only download the file if it doesn't already exist. + + Args: + name: the name of the file + run_path: optional path to a run to pull files from, i.e. `username/project_name/run_id` + if wandb.init has not been called, this is required. + replace: whether to download the file even if it already exists locally + root: the directory to download the file to. Defaults to the current + directory or the run directory if wandb.init was called. + + Returns: + None if it can't find the file, otherwise a file object open for reading + + Raises: + wandb.CommError: if we can't connect to the wandb backend + ValueError: if the file is not found or can't find run_path + """ + is_disabled = wandb.run is not None and wandb.run.disabled + run = None if is_disabled else wandb.run + if run_path is None: + if run is not None: + run_path = run.path + else: + raise ValueError( + "run_path required when calling wandb.restore before wandb.init" + ) + if root is None: + if run is not None: + root = run.dir + api = public.Api() + api_run = api.run(run_path) + if root is None: + root = os.getcwd() + path = os.path.join(root, name) + if os.path.exists(path) and replace is False: + return open(path) + if is_disabled: + return None + files = api_run.files([name]) + if len(files) == 0: + return None + # if the file does not exist, the file has an md5 of 0 + if files[0].md5 == "0": + raise ValueError(f"File {name} not found in {run_path or root}.") + return files[0].download(root=root, replace=True) + + +# propagate our doc string to the runs restore method +try: + Run.restore.__doc__ = restore.__doc__ +except AttributeError: + pass + + +def finish( + exit_code: int | None = None, + quiet: bool | None = None, +) -> None: + """Finish a run and upload any remaining data. + + Marks the completion of a W&B run and ensures all data is synced to the server. + The run's final state is determined by its exit conditions and sync status. + + Run States: + - Running: Active run that is logging data and/or sending heartbeats. + - Crashed: Run that stopped sending heartbeats unexpectedly. + - Finished: Run completed successfully (`exit_code=0`) with all data synced. + - Failed: Run completed with errors (`exit_code!=0`). + + Args: + exit_code: Integer indicating the run's exit status. Use 0 for success, + any other value marks the run as failed. + quiet: Deprecated. Configure logging verbosity using `wandb.Settings(quiet=...)`. + """ + if wandb.run: + wandb.run.finish(exit_code=exit_code, quiet=quiet) diff --git a/lib/python3.12/site-packages/wandb/sdk/wandb_settings.py b/lib/python3.12/site-packages/wandb/sdk/wandb_settings.py new file mode 100644 index 0000000000000000000000000000000000000000..a26c169ac56052e580436c54035a2168539bc40a --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/wandb_settings.py @@ -0,0 +1,1964 @@ +from __future__ import annotations + +import configparser +import json +import logging +import os +import pathlib +import platform +import re +import shutil +import socket +import sys +from datetime import datetime + +# Optional and Union are used for type hinting instead of | because +# the latter is not supported in pydantic<2.6 and Python<3.10. +# Dict, List, and Tuple are used for backwards compatibility +# with pydantic v1 and Python<3.9. +from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Tuple, Union +from urllib.parse import quote, unquote, urlencode + +from google.protobuf.wrappers_pb2 import BoolValue, DoubleValue, Int32Value, StringValue +from pydantic import BaseModel, ConfigDict, Field +from typing_extensions import Self + +import wandb +from wandb import env, util +from wandb._pydantic import ( + IS_PYDANTIC_V2, + AliasChoices, + computed_field, + field_validator, + model_validator, +) +from wandb.errors import UsageError +from wandb.proto import wandb_settings_pb2 + +from .lib import apikey, credentials, ipython +from .lib.gitlib import GitRepo +from .lib.run_moment import RunMoment + +validate_url: Callable[[str], None] + +if IS_PYDANTIC_V2: + from pydantic_core import SchemaValidator, core_schema + + def validate_url(url: str) -> None: + """Validate a URL string.""" + url_validator = SchemaValidator( + core_schema.url_schema( + allowed_schemes=["http", "https"], + strict=True, + ) + ) + url_validator.validate_python(url) +else: + from pydantic import root_validator + + def validate_url(url: str) -> None: + """Validate the base url of the wandb server. + + param value: URL to validate + + Based on the Django URLValidator, but with a few additional checks. + + Copyright (c) Django Software Foundation and individual contributors. + All rights reserved. + + Redistribution and use in source and binary forms, with or without modification, + are permitted provided that the following conditions are met: + + 1. Redistributions of source code must retain the above copyright notice, + this list of conditions and the following disclaimer. + + 2. Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + + 3. Neither the name of Django nor the names of its contributors may be used + to endorse or promote products derived from this software without + specific prior written permission. + + THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND + ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED + WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR + ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES + (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; + LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON + ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS + SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + """ + from urllib.parse import urlparse, urlsplit + + if url is None: + return + + ul = "\u00a1-\uffff" # Unicode letters range (must not be a raw string). + + # IP patterns + ipv4_re = ( + r"(?:0|25[0-5]|2[0-4][0-9]|1[0-9]?[0-9]?|[1-9][0-9]?)" + r"(?:\.(?:0|25[0-5]|2[0-4][0-9]|1[0-9]?[0-9]?|[1-9][0-9]?)){3}" + ) + ipv6_re = r"\[[0-9a-f:.]+\]" # (simple regex, validated later) + + # Host patterns + hostname_re = ( + r"[a-z" + ul + r"0-9](?:[a-z" + ul + r"0-9-]{0,61}[a-z" + ul + r"0-9])?" + ) + # Max length for domain name labels is 63 characters per RFC 1034 sec. 3.1 + domain_re = r"(?:\.(?!-)[a-z" + ul + r"0-9-]{1,63}(? 253: + raise ValueError("hostname is invalid") + + +def _path_convert(*args: str) -> str: + """Join path and apply os.path.expanduser to it.""" + return os.path.expanduser(os.path.join(*args)) + + +class Settings(BaseModel, validate_assignment=True): + """Settings for the W&B SDK. + + This class manages configuration settings for the W&B SDK, + ensuring type safety and validation of all settings. Settings are accessible + as attributes and can be initialized programmatically, through environment + variables (WANDB_ prefix), and via configuration files. + + The settings are organized into three categories: + 1. Public settings: Core configuration options that users can safely modify to customize + W&B's behavior for their specific needs. + 2. Internal settings: Settings prefixed with 'x_' that handle low-level SDK behavior. + These settings are primarily for internal use and debugging. While they can be modified, + they are not considered part of the public API and may change without notice in future + versions. + 3. Computed settings: Read-only settings that are automatically derived from other settings or + the environment. + """ + + # Pydantic Model configuration. + model_config = ConfigDict( + extra="forbid", # throw an error if extra fields are provided + validate_default=True, # validate default values + use_attribute_docstrings=True, # for field descriptions + revalidate_instances="always", + ) + + # Public settings. + + allow_offline_artifacts: bool = True + """Flag to allow table artifacts to be synced in offline mode. + + To revert to the old behavior, set this to False. + """ + + allow_val_change: bool = False + """Flag to allow modification of `Config` values after they've been set.""" + + anonymous: Optional[Literal["allow", "must", "never"]] = None + """Controls anonymous data logging. + + Possible values are: + - "never": requires you to link your W&B account before + tracking the run, so you don't accidentally create an anonymous + run. + - "allow": lets a logged-in user track runs with their account, but + lets someone who is running the script without a W&B account see + the charts in the UI. + - "must": sends the run to an anonymous account instead of to a + signed-up user account. + """ + + api_key: Optional[str] = None + """The W&B API key.""" + + azure_account_url_to_access_key: Optional[Dict[str, str]] = None + """Mapping of Azure account URLs to their corresponding access keys for Azure integration.""" + + base_url: str = "https://api.wandb.ai" + """The URL of the W&B backend for data synchronization.""" + + code_dir: Optional[str] = None + """Directory containing the code to be tracked by W&B.""" + + config_paths: Optional[Sequence[str]] = None + """Paths to files to load configuration from into the `Config` object.""" + + console: Literal["auto", "off", "wrap", "redirect", "wrap_raw", "wrap_emu"] = Field( + default="auto", + validate_default=True, + ) + """The type of console capture to be applied. + + Possible values are: + "auto" - Automatically selects the console capture method based on the + system environment and settings. + + "off" - Disables console capture. + + "redirect" - Redirects low-level file descriptors for capturing output. + + "wrap" - Overrides the write methods of sys.stdout/sys.stderr. Will be + mapped to either "wrap_raw" or "wrap_emu" based on the state of the system. + + "wrap_raw" - Same as "wrap" but captures raw output directly instead of + through an emulator. Derived from the `wrap` setting and should not be set manually. + + "wrap_emu" - Same as "wrap" but captures output through an emulator. + Derived from the `wrap` setting and should not be set manually. + """ + + console_multipart: bool = False + """Whether to produce multipart console log files.""" + + credentials_file: str = Field( + default_factory=lambda: str(credentials.DEFAULT_WANDB_CREDENTIALS_FILE) + ) + """Path to file for writing temporary access tokens.""" + + disable_code: bool = False + """Whether to disable capturing the code.""" + + disable_git: bool = False + """Whether to disable capturing the git state.""" + + disable_job_creation: bool = True + """Whether to disable the creation of a job artifact for W&B Launch.""" + + docker: Optional[str] = None + """The Docker image used to execute the script.""" + + email: Optional[str] = None + """The email address of the user.""" + + entity: Optional[str] = None + """The W&B entity, such as a user or a team.""" + + organization: Optional[str] = None + """The W&B organization.""" + + force: bool = False + """Whether to pass the `force` flag to `wandb.login()`.""" + + fork_from: Optional[RunMoment] = None + """Specifies a point in a previous execution of a run to fork from. + + The point is defined by the run ID, a metric, and its value. + Currently, only the metric '_step' is supported. + """ + + git_commit: Optional[str] = None + """The git commit hash to associate with the run.""" + + git_remote: str = "origin" + """The git remote to associate with the run.""" + + git_remote_url: Optional[str] = None + """The URL of the git remote repository.""" + + git_root: Optional[str] = None + """Root directory of the git repository.""" + + heartbeat_seconds: int = 30 + """Interval in seconds between heartbeat signals sent to the W&B servers. + + + """ + + host: Optional[str] = None + """Hostname of the machine running the script.""" + + http_proxy: Optional[str] = None + """Custom proxy servers for http requests to W&B.""" + + https_proxy: Optional[str] = None + """Custom proxy servers for https requests to W&B.""" + + identity_token_file: Optional[str] = None + """Path to file containing an identity token (JWT) for authentication.""" + + ignore_globs: Sequence[str] = () + """Unix glob patterns relative to `files_dir` specifying files to exclude from upload.""" + + init_timeout: float = 90.0 + """Time in seconds to wait for the `wandb.init` call to complete before timing out.""" + + insecure_disable_ssl: bool = False + """Whether to insecurely disable SSL verification.""" + + job_name: Optional[str] = None + """Name of the Launch job running the script.""" + + job_source: Optional[Literal["repo", "artifact", "image"]] = None + """Source type for Launch.""" + + label_disable: bool = False + """Whether to disable automatic labeling features.""" + + launch: bool = False + """Flag to indicate if the run is being launched through W&B Launch. + + + """ + + launch_config_path: Optional[str] = None + """Path to the launch configuration file.""" + + login_timeout: Optional[float] = None + """Time in seconds to wait for login operations before timing out.""" + + mode: Literal["online", "offline", "shared", "disabled", "dryrun", "run"] = Field( + default="online", + validate_default=True, + ) + """The operating mode for W&B logging and synchronization.""" + + notebook_name: Optional[str] = None + """Name of the notebook if running in a Jupyter-like environment.""" + + program: Optional[str] = None + """Path to the script that created the run, if available.""" + + program_abspath: Optional[str] = None + """The absolute path from the root repository directory to the script that + created the run. + + Root repository directory is defined as the directory containing the + .git directory, if it exists. Otherwise, it's the current working directory. + """ + + program_relpath: Optional[str] = None + """The relative path to the script that created the run.""" + + project: Optional[str] = None + """The W&B project ID.""" + + quiet: bool = False + """Flag to suppress non-essential output.""" + + reinit: Union[ + Literal[ + "default", + "return_previous", + "finish_previous", + "create_new", + ], + bool, + ] = "default" + """What to do when `wandb.init()` is called while a run is active. + + Options: + - "default": Use "finish_previous" in notebooks and "return_previous" + otherwise. + - "return_previous": Return the most recently created run + that is not yet finished. This does not update `wandb.run`; see + the "create_new" option. + - "finish_previous": Finish all active runs, then return a new run. + - "create_new": Create a new run without modifying other active runs. + Does not update `wandb.run` and top-level functions like `wandb.log`. + Because of this, some older integrations that rely on the global run + will not work. + + Can also be a boolean, but this is deprecated. False is the same as + "return_previous", and True is the same as "finish_previous". + """ + + relogin: bool = False + """Flag to force a new login attempt.""" + + resume: Optional[Literal["allow", "must", "never", "auto"]] = None + """Specifies the resume behavior for the run. + + Options: + - "must": Resumes from an existing run with the same ID. If no such run exists, + it will result in failure. + - "allow": Attempts to resume from an existing run with the same ID. If none is + found, a new run will be created. + - "never": Always starts a new run. If a run with the same ID already exists, + it will result in failure. + - "auto": Automatically resumes from the most recent failed run on the same + machine. + """ + + resume_from: Optional[RunMoment] = None + """Specifies a point in a previous execution of a run to resume from. + + The point is defined by the run ID, a metric, and its value. + Currently, only the metric '_step' is supported. + """ + + resumed: bool = False + """Indication from the server about the state of the run. + + This is different from resume, a user provided flag. + + """ + + root_dir: str = Field(default_factory=lambda: os.path.abspath(os.getcwd())) + """The root directory to use as the base for all run-related paths. + + In particular, this is used to derive the wandb directory and the run directory. + """ + + run_group: Optional[str] = None + """Group identifier for related runs. + + Used for grouping runs in the UI. + """ + + run_id: Optional[str] = None + """The ID of the run.""" + + run_job_type: Optional[str] = None + """Type of job being run (e.g., training, evaluation).""" + + run_name: Optional[str] = None + """Human-readable name for the run.""" + + run_notes: Optional[str] = None + """Additional notes or description for the run.""" + + run_tags: Optional[Tuple[str, ...]] = None + """Tags to associate with the run for organization and filtering.""" + + sagemaker_disable: bool = False + """Flag to disable SageMaker-specific functionality.""" + + save_code: Optional[bool] = None + """Whether to save the code associated with the run.""" + + settings_system: Optional[str] = None + """Path to the system-wide settings file.""" + + show_colors: Optional[bool] = None + """Whether to use colored output in the console. + + + """ + + show_emoji: Optional[bool] = None + """Whether to show emoji in the console output. + + + """ + + show_errors: bool = True + """Whether to display error messages.""" + + show_info: bool = True + """Whether to display informational messages.""" + + show_warnings: bool = True + """Whether to display warning messages.""" + + silent: bool = False + """Flag to suppress all output.""" + + start_method: Optional[str] = None + """Method to use for starting subprocesses. + + This is deprecated and will be removed in a future release. + + """ + + strict: Optional[bool] = None + """Whether to enable strict mode for validation and error checking.""" + + summary_timeout: int = 60 + """Time in seconds to wait for summary operations before timing out.""" + + summary_warnings: int = 5 + """Maximum number of summary warnings to display. + + + """ + + sweep_id: Optional[str] = None + """Identifier of the sweep this run belongs to.""" + + sweep_param_path: Optional[str] = None + """Path to the sweep parameters configuration.""" + + symlink: bool = Field( + default_factory=lambda: False if platform.system() == "Windows" else True + ) + """Whether to use symlinks (True by default except on Windows).""" + + sync_tensorboard: Optional[bool] = None + """Whether to synchronize TensorBoard logs with W&B.""" + + table_raise_on_max_row_limit_exceeded: bool = False + """Whether to raise an exception when table row limits are exceeded.""" + + username: Optional[str] = None + """Username.""" + + # Internal settings. + # + # These are typically not meant to be set by the user and should not be considered + # a part of the public API as they may change or be removed in future versions. + + x_cli_only_mode: bool = False + """Flag to indicate that the SDK is running in CLI-only mode. + + + """ + + x_disable_meta: bool = False + """Flag to disable the collection of system metadata.""" + + x_disable_stats: bool = False + """Flag to disable the collection of system metrics.""" + + x_disable_viewer: bool = False + """Flag to disable the early viewer query. + + + """ + + x_disable_machine_info: bool = False + """Flag to disable automatic machine info collection. + + + """ + + x_executable: Optional[str] = None + """Path to the Python executable. + + + """ + + x_extra_http_headers: Optional[Dict[str, str]] = None + """Additional headers to add to all outgoing HTTP requests.""" + + x_file_stream_max_bytes: Optional[int] = None + """An approximate maximum request size for the filestream API. + + Its purpose is to prevent HTTP requests from failing due to + containing too much data. This number is approximate: + requests will be slightly larger. + + """ + + x_file_stream_max_line_bytes: Optional[int] = None + """Maximum line length for filestream JSONL files. + + + """ + + x_file_stream_transmit_interval: Optional[float] = None + """Interval in seconds between filestream transmissions. + + + """ + + # Filestream retry client configuration. + + x_file_stream_retry_max: Optional[int] = None + """Max number of retries for filestream operations. + + + """ + + x_file_stream_retry_wait_min_seconds: Optional[float] = None + """Minimum wait time between retries for filestream operations. + + + """ + + x_file_stream_retry_wait_max_seconds: Optional[float] = None + """Maximum wait time between retries for filestream operations. + + + """ + + x_file_stream_timeout_seconds: Optional[float] = None + """Timeout in seconds for individual filestream HTTP requests. + + + """ + + # file transfer retry client configuration + + x_file_transfer_retry_max: Optional[int] = None + """Max number of retries for file transfer operations. + + + """ + + x_file_transfer_retry_wait_min_seconds: Optional[float] = None + """Minimum wait time between retries for file transfer operations. + + + """ + + x_file_transfer_retry_wait_max_seconds: Optional[float] = None + """Maximum wait time between retries for file transfer operations. + + + """ + + x_file_transfer_timeout_seconds: Optional[float] = None + """Timeout in seconds for individual file transfer HTTP requests. + + + """ + + x_files_dir: Optional[str] = None + """Override setting for the computed files_dir. + + + """ + + x_flow_control_custom: Optional[bool] = None + """Flag indicating custom flow control for filestream. + + TODO: Not implemented in wandb-core. + + """ + + x_flow_control_disabled: Optional[bool] = None + """Flag indicating flow control is disabled for filestream. + + TODO: Not implemented in wandb-core. + + """ + + # graphql retry client configuration + + x_graphql_retry_max: Optional[int] = None + """Max number of retries for GraphQL operations. + + + """ + + x_graphql_retry_wait_min_seconds: Optional[float] = None + """Minimum wait time between retries for GraphQL operations. + + + """ + + x_graphql_retry_wait_max_seconds: Optional[float] = None + """Maximum wait time between retries for GraphQL operations. + + + """ + + x_graphql_timeout_seconds: Optional[float] = None + """Timeout in seconds for individual GraphQL requests. + + + """ + + x_internal_check_process: float = 8.0 + """Interval for internal process health checks in seconds. + + + """ + + x_jupyter_name: Optional[str] = None + """Name of the Jupyter notebook. + + + """ + + x_jupyter_path: Optional[str] = None + """Path to the Jupyter notebook. + + + """ + + x_jupyter_root: Optional[str] = None + """Root directory of the Jupyter notebook. + + + """ + + x_label: Optional[str] = None + """Label to assign to system metrics and console logs collected for the run. + + This is used to group data by on the frontend and can be used to distinguish data + from different processes in a distributed training job. + """ + + x_live_policy_rate_limit: Optional[int] = None + """Rate limit for live policy updates in seconds. + + + """ + + x_live_policy_wait_time: Optional[int] = None + """Wait time between live policy updates in seconds. + + + """ + + x_log_level: int = logging.INFO + """Logging level for internal operations. + + + """ + + x_network_buffer: Optional[int] = None + """Size of the network buffer used in flow control. + + TODO: Not implemented in wandb-core. + + """ + + x_primary: bool = Field( + default=True, validation_alias=AliasChoices("x_primary", "x_primary_node") + ) + """Determines whether to save internal wandb files and metadata. + + In a distributed setting, this is useful for avoiding file overwrites + from secondary processes when only system metrics and logs are needed, + as the primary process handles the main logging. + """ + + x_proxies: Optional[Dict[str, str]] = None + """Custom proxy servers for requests to W&B. + + This is deprecated and will be removed in a future release. + Please use `http_proxy` and `https_proxy` instead. + + """ + + x_runqueue_item_id: Optional[str] = None + """ID of the Launch run queue item being processed. + + + """ + + x_save_requirements: bool = True + """Flag to save the requirements file.""" + + x_server_side_derived_summary: bool = False + """Flag to delegate automatic computation of summary from history to the server. + + This does not disable user-provided summary updates. + """ + + x_server_side_expand_glob_metrics: bool = True + """Flag to delegate glob matching of metrics in define_metric to the server. + + If the server does not support this, the client will perform the glob matching. + + """ + + x_service_transport: Optional[str] = None + """Transport method for communication with the wandb service. + + + """ + + x_service_wait: float = 30.0 + """Time in seconds to wait for the wandb-core internal service to start.""" + + x_skip_transaction_log: bool = False + """Whether to skip saving the run events to the transaction log. + + This is only relevant for online runs. Can be used to reduce the amount of + data written to disk. + + Should be used with caution, as it removes the gurantees about + recoverability. + """ + + x_start_time: Optional[float] = None + """The start time of the run in seconds since the Unix epoch. + + + """ + + x_stats_pid: int = os.getpid() + """PID of the process that started the wandb-core process to collect system stats for. + + + """ + + x_stats_sampling_interval: float = Field(default=15.0) + """Sampling interval for the system monitor in seconds.""" + + x_stats_neuron_monitor_config_path: Optional[str] = None + """Path to the default config file for the neuron-monitor tool. + + This is used to monitor AWS Trainium devices. + + """ + + x_stats_dcgm_exporter: Optional[str] = None + """Endpoint to extract Nvidia DCGM metrics from. + + Options: + - Extract DCGM-related metrics from a query to the Prometheus `/api/v1/query` endpoint. + It is a common practice to aggregate metrics reported by the instances of the DCGM Exporter + running on different nodes in a cluster using Prometheus. + - TODO: Parse metrics directly from the `/metrics` endpoint of the DCGM Exporter. + + Examples: + - `http://localhost:9400/api/v1/query?query=DCGM_FI_DEV_GPU_TEMP{node="l1337", cluster="globular"}`. + - TODO: `http://192.168.0.1:9400/metrics`. + + """ + + x_stats_open_metrics_endpoints: Optional[Dict[str, str]] = None + """OpenMetrics `/metrics` endpoints to monitor for system metrics.""" + + x_stats_open_metrics_filters: Union[ + Dict[str, Dict[str, str]], Sequence[str], None + ] = None + """Filter to apply to metrics collected from OpenMetrics `/metrics` endpoints. + + Supports two formats: + - {"metric regex pattern, including endpoint name as prefix": {"label": "label value regex pattern"}} + - ("metric regex pattern 1", "metric regex pattern 2", ...) + """ + + x_stats_open_metrics_http_headers: Optional[Dict[str, str]] = None + """HTTP headers to add to OpenMetrics requests.""" + + x_stats_disk_paths: Optional[Sequence[str]] = ("/",) + """System paths to monitor for disk usage.""" + + x_stats_cpu_count: Optional[int] = None + """System CPU count. + + If set, overrides the auto-detected value in the run metadata. + """ + + x_stats_cpu_logical_count: Optional[int] = None + """Logical CPU count. + + If set, overrides the auto-detected value in the run metadata. + """ + + x_stats_gpu_count: Optional[int] = None + """GPU device count. + + If set, overrides the auto-detected value in the run metadata. + """ + + x_stats_gpu_type: Optional[str] = None + """GPU device type. + + If set, overrides the auto-detected value in the run metadata. + """ + + x_stats_gpu_device_ids: Optional[Sequence[int]] = None + """GPU device indices to monitor. + + If not set, the system monitor captures metrics for all GPUs. + Assumes 0-based indexing matching CUDA/ROCm device enumeration. + """ + + x_stats_buffer_size: int = 0 + """Number of system metric samples to buffer in memory in the wandb-core process. + + Can be accessed via run._system_metrics. + + """ + + x_stats_coreweave_metadata_base_url: str = "http://169.254.169.254" + """The scheme and hostname for contacting the CoreWeave metadata server. + + Only accessible from within a CoreWeave cluster. + + """ + + x_stats_coreweave_metadata_endpoint: str = "/api/v2/cloud-init/meta-data" + """The relative path on the CoreWeave metadata server to which to make requests. + + This must not include the schema and hostname prefix. + Only accessible from within a CoreWeave cluster. + + """ + + x_stats_track_process_tree: bool = False + """Monitor the entire process tree for resource usage, starting from `x_stats_pid`. + + When `True`, the system monitor aggregates the RSS, CPU%, and thread count + from the process with PID `x_stats_pid` and all of its descendants. + This can have a performance overhead and is disabled by default. + """ + + x_sync: bool = False + """Flag to indicate whether we are syncing a run from the transaction log. + + + """ + + x_update_finish_state: bool = True + """Flag to indicate whether this process can update the run's final state on the server. + + Set to False in distributed training when only the main process should determine the final state. + """ + + # Model validator to catch legacy settings. + @model_validator(mode="before") + @classmethod + def catch_private_settings(cls, values): + """Check if a private field is provided and assign to the corresponding public one. + + This is a compatibility layer to handle previous versions of the settings. + """ + new_values = {} + for key in values: + # Internal settings are prefixed with "x_" instead of "_" + # as Pydantic does not allow "_" in field names. + if key.startswith("_"): + new_values["x" + key] = values[key] + else: + new_values[key] = values[key] + return new_values + + if IS_PYDANTIC_V2: + + @model_validator(mode="after") + def validate_mutual_exclusion_of_branching_args(self) -> Self: + if ( + sum( + o is not None + for o in [self.fork_from, self.resume, self.resume_from] + ) + > 1 + ): + raise ValueError( + "`fork_from`, `resume`, or `resume_from` are mutually exclusive. " + "Please specify only one of them." + ) + return self + + @model_validator(mode="after") + def validate_skip_transaction_log(self): + if self._offline and self.x_skip_transaction_log: + raise ValueError("Cannot skip transaction log in offline mode") + return self + else: + + @root_validator(pre=False) # type: ignore [call-overload] + @classmethod + def validate_mutual_exclusion_of_branching_args(cls, values): + if ( + sum( + values.get(o) is not None + for o in ["fork_from", "resume", "resume_from"] + ) + > 1 + ): + raise ValueError( + "`fork_from`, `resume`, or `resume_from` are mutually exclusive. " + "Please specify only one of them." + ) + return values + + @root_validator(pre=False) # type: ignore [call-overload] + @classmethod + def validate_skip_transaction_log(cls, values): + if values.get("_offline") and values.get("x_skip_transaction_log"): + raise ValueError("Cannot skip transaction log in offline mode") + return values + + # Field validators. + @field_validator("api_key", mode="after") + @classmethod + def validate_api_key(cls, value): + if value is not None and (len(value) > len(value.strip())): + raise UsageError("API key cannot start or end with whitespace") + return value + + @field_validator("base_url", mode="after") + @classmethod + def validate_base_url(cls, value): + validate_url(value) + # wandb.ai-specific checks + if re.match(r".*wandb\.ai[^\.]*$", value) and "api." not in value: + # user might guess app.wandb.ai or wandb.ai is the default cloud server + raise ValueError( + f"{value} is not a valid server address, did you mean https://api.wandb.ai?" + ) + elif re.match(r".*wandb\.ai[^\.]*$", value) and not value.startswith("https"): + raise ValueError("http is not secure, please use https://api.wandb.ai") + return value.rstrip("/") + + @field_validator("code_dir", mode="before") + @classmethod + def validate_code_dir(cls, value): + # TODO: add native support for pathlib.Path + if isinstance(value, pathlib.Path): + return str(value) + return value + + @field_validator("console", mode="after") + @classmethod + def validate_console(cls, value, values): + if value != "auto": + return value + + return "wrap" + + @field_validator("x_executable", mode="before") + @classmethod + def validate_x_executable(cls, value): + # TODO: add native support for pathlib.Path + if isinstance(value, pathlib.Path): + return str(value) + return value + + @field_validator("x_extra_http_headers", mode="before") + @classmethod + def validate_x_extra_http_headers(cls, value): + if isinstance(value, str): + return json.loads(value) + return value + + @field_validator("x_file_stream_max_line_bytes", mode="after") + @classmethod + def validate_file_stream_max_line_bytes(cls, value): + if value is not None and value < 1: + raise ValueError("File stream max line bytes must be greater than 0") + return value + + @field_validator("x_files_dir", mode="before") + @classmethod + def validate_x_files_dir(cls, value): + # TODO: add native support for pathlib.Path + if isinstance(value, pathlib.Path): + return str(value) + return value + + @field_validator("fork_from", mode="before") + @classmethod + def validate_fork_from(cls, value, values) -> Optional[RunMoment]: + run_moment = cls._runmoment_preprocessor(value) + + if hasattr(values, "data"): + # pydantic v2 + values = values.data + else: + # pydantic v1 + values = values + + if ( + run_moment + and values.get("run_id") is not None + and values.get("run_id") == run_moment.run + ): + raise ValueError( + "Provided `run_id` is the same as the run to `fork_from`. " + "Please provide a different `run_id` or remove the `run_id` argument. " + "If you want to rewind the current run, please use `resume_from` instead." + ) + return run_moment + + @field_validator("http_proxy", mode="after") + @classmethod + def validate_http_proxy(cls, value): + if value is None: + return None + validate_url(value) + return value.rstrip("/") + + @field_validator("https_proxy", mode="after") + @classmethod + def validate_https_proxy(cls, value): + if value is None: + return None + validate_url(value) + return value.rstrip("/") + + @field_validator("ignore_globs", mode="after") + @classmethod + def validate_ignore_globs(cls, value): + return tuple(value) if not isinstance(value, tuple) else value + + @field_validator("program", mode="before") + @classmethod + def validate_program(cls, value): + # TODO: add native support for pathlib.Path + if isinstance(value, pathlib.Path): + return str(value) + return value + + @field_validator("program_abspath", mode="before") + @classmethod + def validate_program_abspath(cls, value): + # TODO: add native support for pathlib.Path + if isinstance(value, pathlib.Path): + return str(value) + return value + + @field_validator("program_relpath", mode="before") + @classmethod + def validate_program_relpath(cls, value): + # TODO: add native support for pathlib.Path + if isinstance(value, pathlib.Path): + return str(value) + return value + + @field_validator("project", mode="after") + @classmethod + def validate_project(cls, value, values): + if value is None: + return None + invalid_chars_list = list("/\\#?%:") + if len(value) > 128: + raise UsageError(f"Invalid project name {value!r}: exceeded 128 characters") + invalid_chars = {char for char in invalid_chars_list if char in value} + if invalid_chars: + raise UsageError( + f"Invalid project name {value!r}: " + f"cannot contain characters {','.join(invalid_chars_list)!r}, " + f"found {','.join(invalid_chars)!r}" + ) + return value + + @field_validator("resume", mode="before") + @classmethod + def validate_resume(cls, value): + if value is False: + return None + if value is True: + return "auto" + return value + + @field_validator("resume_from", mode="before") + @classmethod + def validate_resume_from(cls, value, values) -> Optional[RunMoment]: + run_moment = cls._runmoment_preprocessor(value) + + if hasattr(values, "data"): + # pydantic v2 + values = values.data + else: + # pydantic v1 + values = values + + if ( + run_moment + and values.get("run_id") is not None + and values.get("run_id") != run_moment.run + ): + raise ValueError( + "Both `run_id` and `resume_from` have been specified with different ids." + ) + return run_moment + + @field_validator("root_dir", mode="before") + @classmethod + def validate_root_dir(cls, value): + # TODO: add native support for pathlib.Path + if isinstance(value, pathlib.Path): + return str(value) + return value + + @field_validator("run_id", mode="after") + @classmethod + def validate_run_id(cls, value, values): + if value is None: + return None + + if len(value) == 0: + raise UsageError("Run ID cannot be empty") + if len(value) > len(value.strip()): + raise UsageError("Run ID cannot start or end with whitespace") + if not bool(value.strip()): + raise UsageError("Run ID cannot contain only whitespace") + + # check if the run id contains any reserved characters + reserved_chars = ":;,#?/'" + if any(char in reserved_chars for char in value): + raise UsageError(f"Run ID cannot contain the characters: {reserved_chars}") + return value + + @field_validator("settings_system", mode="after") + @classmethod + def validate_settings_system(cls, value): + if value is None: + return None + elif isinstance(value, pathlib.Path): + return str(_path_convert(value)) + else: + return _path_convert(value) + + @field_validator("x_service_wait", mode="after") + @classmethod + def validate_service_wait(cls, value): + if value < 0: + raise UsageError("Service wait time cannot be negative") + return value + + @field_validator("start_method", mode="after") + @classmethod + def validate_start_method(cls, value): + if value is None: + return value + wandb.termwarn( + "`start_method` is deprecated and will be removed in a future version " + "of wandb. This setting is currently non-functional and safely ignored.", + repeat=False, + ) + return value + + @field_validator("x_stats_coreweave_metadata_base_url", mode="after") + @classmethod + def validate_x_stats_coreweave_metadata_base_url(cls, value): + validate_url(value) + return value.rstrip("/") + + @field_validator("x_stats_gpu_device_ids", mode="before") + @classmethod + def validate_x_stats_gpu_device_ids(cls, value): + if isinstance(value, str): + return json.loads(value) + return value + + @field_validator("x_stats_neuron_monitor_config_path", mode="before") + @classmethod + def validate_x_stats_neuron_monitor_config_path(cls, value): + # TODO: add native support for pathlib.Path + if isinstance(value, pathlib.Path): + return str(value) + return value + + @field_validator("x_stats_open_metrics_endpoints", mode="before") + @classmethod + def validate_stats_open_metrics_endpoints(cls, value): + if isinstance(value, str): + return json.loads(value) + return value + + @field_validator("x_stats_open_metrics_filters", mode="before") + @classmethod + def validate_stats_open_metrics_filters(cls, value): + if isinstance(value, str): + return json.loads(value) + return value + + @field_validator("x_stats_open_metrics_http_headers", mode="before") + @classmethod + def validate_stats_open_metrics_http_headers(cls, value): + if isinstance(value, str): + return json.loads(value) + return value + + @field_validator("x_stats_sampling_interval", mode="after") + @classmethod + def validate_stats_sampling_interval(cls, value): + if value < 0.1: + raise UsageError("Stats sampling interval cannot be less than 0.1 seconds") + return value + + @field_validator("sweep_id", mode="after") + @classmethod + def validate_sweep_id(cls, value): + if value is None: + return None + if len(value) == 0: + raise UsageError("Sweep ID cannot be empty") + if len(value) > len(value.strip()): + raise UsageError("Sweep ID cannot start or end with whitespace") + if not bool(value.strip()): + raise UsageError("Sweep ID cannot contain only whitespace") + return value + + @field_validator("sweep_param_path", mode="before") + @classmethod + def validate_sweep_param_path(cls, value): + # TODO: add native support for pathlib.Path + if isinstance(value, pathlib.Path): + return str(value) + return value + + # Computed fields. + + @computed_field # type: ignore[prop-decorator] + @property + def _args(self) -> List[str]: + if not self._jupyter: + return sys.argv[1:] + return [] + + @computed_field # type: ignore[prop-decorator] + @property + def _aws_lambda(self) -> bool: + """Check if we are running in a lambda environment.""" + from sentry_sdk.integrations.aws_lambda import ( # type: ignore[import-not-found] + get_lambda_bootstrap, + ) + + lambda_bootstrap = get_lambda_bootstrap() + if not lambda_bootstrap or not hasattr( + lambda_bootstrap, "handle_event_request" + ): + return False + return True + + @computed_field # type: ignore[prop-decorator] + @property + def _code_path_local(self) -> Optional[str]: + """The relative path from the current working directory to the code path. + + For example, if the code path is /home/user/project/example.py, and the + current working directory is /home/user/project, then the code path local + is example.py. + + If couldn't find the relative path, this will be an empty string. + """ + return self._get_program_relpath(self.program) if self.program else None + + @computed_field # type: ignore[prop-decorator] + @property + def _colab(self) -> bool: + return "google.colab" in sys.modules + + @computed_field # type: ignore[prop-decorator] + @property + def _ipython(self) -> bool: + return ipython.in_ipython() + + @computed_field # type: ignore[prop-decorator] + @property + def _jupyter(self) -> bool: + return ipython.in_jupyter() + + @computed_field # type: ignore[prop-decorator] + @property + def _kaggle(self) -> bool: + return util._is_likely_kaggle() + + @computed_field # type: ignore[prop-decorator] + @property + def _noop(self) -> bool: + return self.mode == "disabled" + + @computed_field # type: ignore[prop-decorator] + @property + def _notebook(self) -> bool: + return self._ipython or self._jupyter or self._colab or self._kaggle + + @computed_field # type: ignore[prop-decorator] + @property + def _offline(self) -> bool: + return self.mode in ("offline", "dryrun") + + @computed_field # type: ignore[prop-decorator] + @property + def _os(self) -> str: + """The operating system of the machine running the script.""" + return platform.platform(aliased=True) + + @computed_field # type: ignore[prop-decorator] + @property + def _platform(self) -> str: + return f"{platform.system()}-{platform.machine()}".lower() + + @computed_field # type: ignore[prop-decorator] + @property + def _python(self) -> str: + return f"{platform.python_implementation()} {platform.python_version()}" + + @computed_field # type: ignore[prop-decorator] + @property + def _shared(self) -> bool: + """Whether we are in shared mode. + + In "shared" mode, multiple processes can write to the same run, + for example from different machines. + """ + return self.mode == "shared" + + @computed_field # type: ignore[prop-decorator] + @property + def _start_datetime(self) -> str: + if self.x_start_time is None: + return "" + datetime_now = datetime.fromtimestamp(self.x_start_time) + return datetime_now.strftime("%Y%m%d_%H%M%S") + + @computed_field # type: ignore[prop-decorator] + @property + def _tmp_code_dir(self) -> str: + return _path_convert(self.sync_dir, "tmp", "code") + + @computed_field # type: ignore[prop-decorator] + @property + def _windows(self) -> bool: + return platform.system() == "Windows" + + @computed_field # type: ignore[prop-decorator] + @property + def colab_url(self) -> Optional[str]: + """The URL to the Colab notebook, if running in Colab.""" + if not self._colab: + return None + if self.x_jupyter_path and self.x_jupyter_path.startswith("fileId="): + unescaped = unquote(self.x_jupyter_path) + return "https://colab.research.google.com/notebook#" + unescaped + return None + + @computed_field # type: ignore[prop-decorator] + @property + def deployment(self) -> Literal["local", "cloud"]: + return "local" if self.is_local else "cloud" + + @computed_field # type: ignore[prop-decorator] + @property + def files_dir(self) -> str: + """Absolute path to the local directory where the run's files are stored.""" + return self.x_files_dir or _path_convert(self.sync_dir, "files") + + @computed_field # type: ignore[prop-decorator] + @property + def is_local(self) -> bool: + return str(self.base_url) != "https://api.wandb.ai" + + @computed_field # type: ignore[prop-decorator] + @property + def log_dir(self) -> str: + """The directory for storing log files.""" + return _path_convert(self.sync_dir, "logs") + + @computed_field # type: ignore[prop-decorator] + @property + def log_internal(self) -> str: + """The path to the file to use for internal logs.""" + return _path_convert(self.log_dir, "debug-internal.log") + + @computed_field # type: ignore[prop-decorator] + @property + def log_symlink_internal(self) -> str: + """The path to the symlink to the internal log file of the most recent run.""" + return _path_convert(self.wandb_dir, "debug-internal.log") + + @computed_field # type: ignore[prop-decorator] + @property + def log_symlink_user(self) -> str: + """The path to the symlink to the user-process log file of the most recent run.""" + return _path_convert(self.wandb_dir, "debug.log") + + @computed_field # type: ignore[prop-decorator] + @property + def log_user(self) -> str: + """The path to the file to use for user-process logs.""" + return _path_convert(self.log_dir, "debug.log") + + @computed_field # type: ignore[prop-decorator] + @property + def project_url(self) -> str: + """The W&B URL where the project can be viewed.""" + project_url = self._project_url_base() + if not project_url: + return "" + + query = self._get_url_query_string() + + return f"{project_url}{query}" + + @computed_field # type: ignore[prop-decorator] + @property + def resume_fname(self) -> str: + """The path to the resume file.""" + return _path_convert(self.wandb_dir, "wandb-resume.json") + + @computed_field # type: ignore[prop-decorator] + @property + def run_mode(self) -> Literal["run", "offline-run"]: + return "run" if not self._offline else "offline-run" + + @computed_field # type: ignore[prop-decorator] + @property + def run_url(self) -> str: + """The W&B URL where the run can be viewed.""" + project_url = self._project_url_base() + if not all([project_url, self.run_id]): + return "" + + query = self._get_url_query_string() + # Exclude specific safe characters from URL encoding to prevent 404 errors + safe_chars = "=+&$@" + return f"{project_url}/runs/{quote(self.run_id or '', safe=safe_chars)}{query}" + + @computed_field # type: ignore[prop-decorator] + @property + def settings_workspace(self) -> str: + """The path to the workspace settings file.""" + return _path_convert(self.wandb_dir, "settings") + + @computed_field # type: ignore[prop-decorator] + @property + def sweep_url(self) -> str: + """The W&B URL where the sweep can be viewed.""" + project_url = self._project_url_base() + if not all([project_url, self.sweep_id]): + return "" + + query = self._get_url_query_string() + return f"{project_url}/sweeps/{quote(self.sweep_id or '')}{query}" + + @computed_field # type: ignore[prop-decorator] + @property + def sync_dir(self) -> str: + return _path_convert( + self.wandb_dir, + f"{self.run_mode}-{self.timespec}-{self.run_id}", + ) + + @computed_field # type: ignore[prop-decorator] + @property + def sync_file(self) -> str: + """Path to the append-only binary transaction log file.""" + return _path_convert(self.sync_dir, f"run-{self.run_id}.wandb") + + @computed_field # type: ignore[prop-decorator] + @property + def sync_symlink_latest(self) -> str: + return _path_convert(self.wandb_dir, "latest-run") + + @computed_field # type: ignore[prop-decorator] + @property + def timespec(self) -> str: + return self._start_datetime + + @computed_field # type: ignore[prop-decorator] + @property + def wandb_dir(self) -> str: + """Full path to the wandb directory.""" + stage_dir = ( + ".wandb" + os.sep + if os.path.exists(os.path.join(self.root_dir, ".wandb")) + else "wandb" + os.sep + ) + path = os.path.join(self.root_dir, stage_dir) + return os.path.expanduser(path) + + # Methods to collect and update settings from different sources. + # + # The Settings class does not track the source of the settings, + # so it is up to the developer to ensure that the settings are applied + # in the correct order. Most of the updates are done in + # wandb/sdk/wandb_setup.py::_WandbSetup._settings_setup. + + def update_from_system_config_file(self): + """Update settings from the system config file.""" + if not self.settings_system or not os.path.exists(self.settings_system): + return + for key, value in self._load_config_file(self.settings_system).items(): + if value is not None: + setattr(self, key, value) + + def update_from_workspace_config_file(self): + """Update settings from the workspace config file.""" + if not self.settings_workspace or not os.path.exists(self.settings_workspace): + return + for key, value in self._load_config_file(self.settings_workspace).items(): + if value is not None: + setattr(self, key, value) + + def update_from_env_vars(self, environ: Dict[str, Any]): + """Update settings from environment variables.""" + env_prefix: str = "WANDB_" + private_env_prefix: str = env_prefix + "_" + special_env_var_names = { + "WANDB_SERVICE_TRANSPORT": "x_service_transport", + "WANDB_DIR": "root_dir", + "WANDB_NAME": "run_name", + "WANDB_NOTES": "run_notes", + "WANDB_TAGS": "run_tags", + "WANDB_JOB_TYPE": "run_job_type", + "WANDB_HTTP_TIMEOUT": "x_graphql_timeout_seconds", + "WANDB_FILE_PUSHER_TIMEOUT": "x_file_transfer_timeout_seconds", + "WANDB_USER_EMAIL": "email", + } + env = dict() + for setting, value in environ.items(): + if not setting.startswith(env_prefix): + continue + + if setting in special_env_var_names: + key = special_env_var_names[setting] + elif setting.startswith(private_env_prefix): + key = "x_" + setting[len(private_env_prefix) :].lower() + else: + # otherwise, strip the prefix and convert to lowercase + key = setting[len(env_prefix) :].lower() + + if key in self.__dict__: + if key in ("ignore_globs", "run_tags"): + value = value.split(",") + env[key] = value + + for key, value in env.items(): + if value is not None: + setattr(self, key, value) + + def update_from_system_environment(self): + """Update settings from the system environment.""" + # For code saving, only allow env var override if value from server is true, or + # if no preference was specified. + if (self.save_code is True or self.save_code is None) and ( + os.getenv(env.SAVE_CODE) is not None + or os.getenv(env.DISABLE_CODE) is not None + ): + self.save_code = env.should_save_code() + + if os.getenv(env.DISABLE_GIT) is not None: + self.disable_git = env.disable_git() + + # Attempt to get notebook information if not already set by the user + if self._jupyter and (self.notebook_name is None or self.notebook_name == ""): + meta = wandb.jupyter.notebook_metadata(self.silent) # type: ignore + self.x_jupyter_path = meta.get("path") + self.x_jupyter_name = meta.get("name") + self.x_jupyter_root = meta.get("root") + elif ( + self._jupyter + and self.notebook_name is not None + and os.path.exists(self.notebook_name) + ): + self.x_jupyter_path = self.notebook_name + self.x_jupyter_name = self.notebook_name + self.x_jupyter_root = os.getcwd() + elif self._jupyter: + wandb.termwarn( + "WANDB_NOTEBOOK_NAME should be a path to a notebook file, " + f"couldn't find {self.notebook_name}.", + ) + + # host is populated by update_from_env_vars if the corresponding env + # vars exist -- but if they don't, we'll fill them in here. + if self.host is None: + self.host = socket.gethostname() # type: ignore + + _executable = ( + self.x_executable + or os.environ.get(env._EXECUTABLE) + or sys.executable + or shutil.which("python3") + or "python3" + ) + self.x_executable = _executable + + if self.docker is None: + self.docker = env.get_docker(util.image_id_from_k8s()) + + # proceed if not in CLI mode + if self.x_cli_only_mode: + return + + program = self.program or self._get_program() + + if program is not None: + try: + root = ( + GitRepo().root or os.getcwd() + if not self.disable_git + else os.getcwd() + ) + except Exception: + # if the git command fails, fall back to the current working directory + root = os.getcwd() + + self.program_relpath = self.program_relpath or self._get_program_relpath( + program, root + ) + program_abspath = os.path.abspath( + os.path.join(root, os.path.relpath(os.getcwd(), root), program) + ) + if os.path.exists(program_abspath): + self.program_abspath = program_abspath + else: + program = "" + + self.program = program + + def update_from_dict(self, settings: Dict[str, Any]) -> None: + """Update settings from a dictionary.""" + for key, value in dict(settings).items(): + if value is not None: + setattr(self, key, value) + + def update_from_settings(self, settings: Settings) -> None: + """Update settings from another instance of `Settings`.""" + d = {field: getattr(settings, field) for field in settings.model_fields_set} + if d: + self.update_from_dict(d) + + # Helper methods. + + def to_proto(self) -> wandb_settings_pb2.Settings: + """Generate a protobuf representation of the settings.""" + settings_proto = wandb_settings_pb2.Settings() + for k, v in self.model_dump(exclude_none=True).items(): + # Client-only settings that don't exist on the protobuf. + if k in ("reinit",): + continue + + # Special case for x_stats_open_metrics_filters. + if k == "x_stats_open_metrics_filters": + if isinstance(v, (list, set, tuple)): + setting = getattr(settings_proto, k) + setting.sequence.value.extend(v) + elif isinstance(v, dict): + setting = getattr(settings_proto, k) + for key, value in v.items(): + for kk, vv in value.items(): + setting.mapping.value[key].value[kk] = vv + else: + raise TypeError(f"Unsupported type {type(v)} for setting {k}") + continue + + # Special case for RunMoment fields. + if k in ("fork_from", "resume_from"): + run_moment = ( + v + if isinstance(v, RunMoment) + else RunMoment( + run=v.get("run"), + value=v.get("value"), + metric=v.get("metric"), + ) + ) + getattr(settings_proto, k).CopyFrom( + wandb_settings_pb2.RunMoment( + run=run_moment.run, + value=run_moment.value, + metric=run_moment.metric, + ) + ) + continue + + if isinstance(v, bool): + getattr(settings_proto, k).CopyFrom(BoolValue(value=v)) + elif isinstance(v, int): + getattr(settings_proto, k).CopyFrom(Int32Value(value=v)) + elif isinstance(v, float): + getattr(settings_proto, k).CopyFrom(DoubleValue(value=v)) + elif isinstance(v, str): + getattr(settings_proto, k).CopyFrom(StringValue(value=v)) + elif isinstance(v, (list, set, tuple)): + # we only support sequences of strings for now + sequence = getattr(settings_proto, k) + sequence.value.extend(v) + elif isinstance(v, dict): + mapping = getattr(settings_proto, k) + for key, value in v.items(): + # we only support dicts with string values for now + mapping.value[key] = value + elif v is None: + # None means that the setting value was not set. + pass + else: + raise TypeError(f"Unsupported type {type(v)} for setting {k}") + + return settings_proto + + def _get_program(self) -> Optional[str]: + """Get the program that started the current process.""" + if self._jupyter: + # If in a notebook, try to get the program from the notebook metadata. + if self.notebook_name: + return self.notebook_name + + if not self.x_jupyter_path: + return self.program + + if self.x_jupyter_path.startswith("fileId="): + return self.x_jupyter_name + + return self.x_jupyter_path + + # If not in a notebook, try to get the program from the environment + # or the __main__ module for scripts run as `python -m ...`. + program = os.getenv(env.PROGRAM) + if program is not None: + return program + + try: + import __main__ + except ImportError: + return None + + try: + if __main__.__spec__ is None: + python_args = __main__.__file__ + else: + python_args = f"-m {__main__.__spec__.name}" + except AttributeError: + return None + + return python_args + + @staticmethod + def _get_program_relpath(program: str, root: Optional[str] = None) -> Optional[str]: + """Get the relative path to the program from the root directory.""" + if not program: + return None + + root = root or os.getcwd() + if not root: + return None + + # For windows if the root and program are on different drives, + # os.path.relpath will raise a ValueError. + if not util.are_paths_on_same_drive(root, program): + return None + + full_path_to_program = os.path.join( + root, os.path.relpath(os.getcwd(), root), program + ) + if os.path.exists(full_path_to_program): + relative_path = os.path.relpath(full_path_to_program, start=root) + if "../" in relative_path: + return None + return relative_path + + return None + + @staticmethod + def _load_config_file(file_name: str, section: str = "default") -> dict: + """Load a config file and return the settings for a given section.""" + parser = configparser.ConfigParser() + parser.add_section(section) + parser.read(file_name) + config: Dict[str, Any] = dict() + for k in parser[section]: + config[k] = parser[section][k] + if k == "ignore_globs": + config[k] = config[k].split(",") + return config + + def _project_url_base(self) -> str: + """Construct the base URL for the project.""" + if not all([self.entity, self.project]): + return "" + + app_url = util.app_url(self.base_url) + return f"{app_url}/{quote(self.entity or '')}/{quote(self.project or '')}" + + def _get_url_query_string(self) -> str: + """Construct the query string for project, run, and sweep URLs.""" + # TODO: remove dependency on Api() + if self.anonymous not in ["allow", "must"]: + return "" + + api_key = apikey.api_key(settings=self) + + return f"?{urlencode({'apiKey': api_key})}" + + @staticmethod + def _runmoment_preprocessor( + val: Union[RunMoment, str, None], + ) -> Optional[RunMoment]: + """Preprocess the setting for forking or resuming a run.""" + if isinstance(val, RunMoment) or val is None: + return val + elif isinstance(val, str): + return RunMoment.from_uri(val) + + if not IS_PYDANTIC_V2: + + def model_copy(self, *args, **kwargs): + return self.copy(*args, **kwargs) + + def model_dump(self, **kwargs): + """Compatibility method for Pydantic v1 to mimic v2's model_dump. + + In v1, this is equivalent to dict() but also includes computed properties. + + Args: + **kwargs: Options passed to the dict method + - exclude_none: Whether to exclude fields with None values + + Returns: + A dictionary of the model's fields and computed properties + """ + # Handle exclude_none separately since it's named differently in v1 + exclude_none = kwargs.pop("exclude_none", False) + + # Start with regular fields from dict() + result = self.dict(**kwargs) + + # Get all computed properties + for name in dir(self.__class__): + attr = getattr(self.__class__, name, None) + if isinstance(attr, property): + try: + # Only include properties that don't raise errors + value = getattr(self, name) + result[name] = value + except (AttributeError, NotImplementedError, TypeError, ValueError): + # Skip properties that can't be accessed or raise errors + pass + elif isinstance(attr, RunMoment): + value = getattr(self, name) + result[name] = value + + # Special Pydantic attributes that should always be excluded + exclude_fields = { + "model_config", + "model_fields", + "model_fields_set", + "__fields__", + "__model_fields_set", + "__pydantic_self__", + "__pydantic_initialised__", + } + + # Remove special Pydantic attributes + for field in exclude_fields: + if field in result: + del result[field] + + if exclude_none: + # Remove None values from the result + return {k: v for k, v in result.items() if v is not None} + + return result + + @property + def model_fields_set(self) -> set: + """Return a set of fields that have been explicitly set. + + This is a compatibility property for Pydantic v1 to mimic v2's model_fields_set. + """ + return getattr(self, "__fields_set__", set()) diff --git a/lib/python3.12/site-packages/wandb/sdk/wandb_setup.py b/lib/python3.12/site-packages/wandb/sdk/wandb_setup.py new file mode 100644 index 0000000000000000000000000000000000000000..36bcf42ec97a878f184b324c401e1c338ac7e32d --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/wandb_setup.py @@ -0,0 +1,480 @@ +"""Global W&B library state. + +This module manages global state, which for wandb includes: + +- Settings configured through `wandb.setup()` +- The list of active runs +- A subprocess ("the internal service") that asynchronously uploads metrics + +This module is fork-aware: in a forked process such as that spawned by the +`multiprocessing` module, `wandb.singleton()` returns a new object, not the +one inherited from the parent process. This requirement comes from backward +compatibility with old design choices: the hardest one to fix is that wandb +was originally designed to have a single run for the entire process that +`wandb.init()` was meant to return. Back then, the only way to create +multiple simultaneous runs in a single script was to run subprocesses, and since +the built-in `multiprocessing` module forks by default, this required a PID +check to make `wandb.init()` ignore the inherited global run. + +Another reason for fork-awareness is that the process that starts up +the internal service owns it and is responsible for shutting it down, +and child processes shouldn't also try to do that. This is easier to +redesign. +""" + +from __future__ import annotations + +import logging +import os +import pathlib +import sys +from typing import TYPE_CHECKING, Any, Union + +import wandb +import wandb.integration.sagemaker as sagemaker +from wandb.env import CONFIG_DIR +from wandb.sdk.lib import import_hooks, wb_logging + +from . import wandb_settings +from .lib import config_util, server + +if TYPE_CHECKING: + from wandb.sdk import wandb_run + from wandb.sdk.lib.service.service_connection import ServiceConnection + from wandb.sdk.wandb_settings import Settings + + +class _EarlyLogger: + """Early logger which captures logs in memory until logging can be configured.""" + + def __init__(self) -> None: + self._log: list[tuple] = [] + self._exception: list[tuple] = [] + # support old warn() as alias of warning() + self.warn = self.warning + + def debug(self, msg: str, *args: Any, **kwargs: Any) -> None: + self._log.append((logging.DEBUG, msg, args, kwargs)) + + def info(self, msg: str, *args: Any, **kwargs: Any) -> None: + self._log.append((logging.INFO, msg, args, kwargs)) + + def warning(self, msg: str, *args: Any, **kwargs: Any) -> None: + self._log.append((logging.WARNING, msg, args, kwargs)) + + def error(self, msg: str, *args: Any, **kwargs: Any) -> None: + self._log.append((logging.ERROR, msg, args, kwargs)) + + def critical(self, msg: str, *args: Any, **kwargs: Any) -> None: + self._log.append((logging.CRITICAL, msg, args, kwargs)) + + def exception(self, msg: str, *args: Any, **kwargs: Any) -> None: + self._exception.append((msg, args, kwargs)) + + def log(self, level: str, msg: str, *args: Any, **kwargs: Any) -> None: + self._log.append((level, msg, args, kwargs)) + + def _flush(self, new_logger: Logger) -> None: + assert self is not new_logger + for level, msg, args, kwargs in self._log: + new_logger.log(level, msg, *args, **kwargs) + for msg, args, kwargs in self._exception: + new_logger.exception(msg, *args, **kwargs) + + +Logger = Union[logging.Logger, _EarlyLogger] + + +class _WandbSetup: + """W&B library singleton.""" + + def __init__( + self, + pid: int, + settings: Settings | None = None, + environ: dict | None = None, + ) -> None: + self._connection: ServiceConnection | None = None + + self._active_runs: list[wandb_run.Run] = [] + + self._environ = environ or dict(os.environ) + self._sweep_config: dict | None = None + self._config: dict | None = None + self._server: server.Server | None = None + self._pid = pid + + # TODO(jhr): defer strict checks until settings are fully initialized + # and logging is ready + self._logger: Logger = _EarlyLogger() + + self._settings = self._settings_setup(settings) + + wandb.termsetup(self._settings, None) + + self._setup() + + def add_active_run(self, run: wandb_run.Run) -> None: + """Append a run to the active runs list. + + This must be called when a run is initialized. + + Args: + run: A newly initialized run. + """ + if run not in self._active_runs: + self._active_runs.append(run) + + def remove_active_run(self, run: wandb_run.Run) -> None: + """Remove the run from the active runs list. + + This must be called when a run is finished. + + Args: + run: A run that is finished or crashed. + """ + try: + self._active_runs.remove(run) + except ValueError: + pass # Removing a run multiple times is not an error. + + @property + def most_recent_active_run(self) -> wandb_run.Run | None: + """The most recently initialized run that is not yet finished.""" + if not self._active_runs: + return None + + return self._active_runs[-1] + + def finish_all_active_runs(self) -> None: + """Finish all unfinished runs. + + NOTE: This is slightly inefficient as it finishes runs one at a time. + This only exists to support using the `reinit="finish_previous"` + setting together with `reinit="create_new"` which does not seem to be a + useful pattern. Since `"create_new"` should eventually become the + default and only behavior, it does not seem worth optimizing. + """ + # Take a snapshot as each call to `finish()` modifies `_active_runs`. + runs_copy = list(self._active_runs) + for run in runs_copy: + run.finish() + + def _settings_setup( + self, + settings: Settings | None, + ) -> wandb_settings.Settings: + s = wandb_settings.Settings() + + # the pid of the process to monitor for system stats + pid = os.getpid() + self._logger.info(f"Current SDK version is {wandb.__version__}") + self._logger.info(f"Configure stats pid to {pid}") + s.x_stats_pid = pid + + if settings and settings.settings_system: + s.settings_system = settings.settings_system + elif config_dir_str := os.getenv(CONFIG_DIR, None): + config_dir = pathlib.Path(config_dir_str).expanduser() + s.settings_system = str(config_dir / "settings") + else: + s.settings_system = str( + pathlib.Path("~", ".config", "wandb", "settings").expanduser() + ) + + # load settings from the system config + if s.settings_system: + self._logger.info(f"Loading settings from {s.settings_system}") + s.update_from_system_config_file() + + # load settings from the workspace config + if s.settings_workspace: + self._logger.info(f"Loading settings from {s.settings_workspace}") + s.update_from_workspace_config_file() + + # load settings from the environment variables + self._logger.info("Loading settings from environment variables") + s.update_from_env_vars(self._environ) + + # infer settings from the system environment + s.update_from_system_environment() + + # load SageMaker settings + check_sagemaker_env = not s.sagemaker_disable + if settings and settings.sagemaker_disable: + check_sagemaker_env = False + if check_sagemaker_env and sagemaker.is_using_sagemaker(): + self._logger.info("Loading SageMaker settings") + sagemaker.set_global_settings(s) + + # load settings from the passed init/setup settings + if settings: + s.update_from_settings(settings) + + return s + + def _update(self, settings: Settings | None = None) -> None: + if not settings: + return + self._settings.update_from_settings(settings) + + def _update_user_settings(self) -> None: + # Get rid of cached results to force a refresh. + self._server = None + user_settings = self._load_user_settings() + if user_settings is not None: + self._settings.update_from_dict(user_settings) + + def _early_logger_flush(self, new_logger: Logger) -> None: + if self._logger is new_logger: + return + + if isinstance(self._logger, _EarlyLogger): + self._logger._flush(new_logger) + self._logger = new_logger + + def _get_logger(self) -> Logger: + return self._logger + + @property + def settings(self) -> wandb_settings.Settings: + return self._settings + + def _get_entity(self) -> str | None: + if self._settings and self._settings._offline: + return None + entity = self.viewer.get("entity") + return entity + + def _get_username(self) -> str | None: + if self._settings and self._settings._offline: + return None + return self.viewer.get("username") + + def _get_teams(self) -> list[str]: + if self._settings and self._settings._offline: + return [] + teams = self.viewer.get("teams") + if teams: + teams = [team["node"]["name"] for team in teams["edges"]] + return teams or [] + + @property + def viewer(self) -> dict[str, Any]: + if self._server is None: + self._server = server.Server(settings=self._settings) + + return self._server.viewer + + def _load_user_settings(self) -> dict[str, Any] | None: + # offline? + if self._server is None: + return None + + flags = self._server._flags + user_settings = dict() + if "code_saving_enabled" in flags: + user_settings["save_code"] = flags["code_saving_enabled"] + + email = self.viewer.get("email", None) + if email: + user_settings["email"] = email + + return user_settings + + def _setup(self) -> None: + sweep_path = self._settings.sweep_param_path + if sweep_path: + self._sweep_config = config_util.dict_from_config_file( + sweep_path, must_exist=True + ) + + # if config_paths was set, read in config dict + if self._settings.config_paths: + # TODO(jhr): handle load errors, handle list of files + for config_path in self._settings.config_paths: + config_dict = config_util.dict_from_config_file(config_path) + if config_dict is None: + continue + if self._config is not None: + self._config.update(config_dict) + else: + self._config = config_dict + + def _teardown(self, exit_code: int | None = None) -> None: + import_hooks.unregister_all_post_import_hooks() + + if not self._connection: + return + + # Reset to None so that setup() creates a new connection. + connection = self._connection + self._connection = None + + internal_exit_code = connection.teardown(exit_code or 0) + if internal_exit_code not in (None, 0): + sys.exit(internal_exit_code) + + def ensure_service(self) -> ServiceConnection: + """Returns a connection to the service process creating it if needed.""" + if self._connection: + return self._connection + + from wandb.sdk.lib.service import service_connection + + self._connection = service_connection.connect_to_service(self._settings) + return self._connection + + def assert_service(self) -> ServiceConnection: + """Returns a connection to the service process, asserting it exists. + + Unlike ensure_service(), this will not start up a service process + if it didn't already exist. + """ + if not self._connection: + raise AssertionError("Expected service process to exist.") + + return self._connection + + +_singleton: _WandbSetup | None = None +"""The W&B library singleton, or None if not yet set up. + +The value is invalid and must not be used if `os.getpid() != _singleton._pid`. +""" + + +def singleton() -> _WandbSetup: + """The W&B singleton for the current process. + + The first call to this in this process (which may be a fork of another + process) creates the singleton, and all subsequent calls return it + until teardown(). This does not start the service process. + """ + return _setup(start_service=False) + + +def singleton_if_setup() -> _WandbSetup | None: + """The W&B singleton for the current process or None if it isn't set up. + + Always prefer singleton() over this function. + + Unlike singleton(), this never creates the singleton and therefore never + initializes global settings from the environment. This is useful only + during tests, which may modify the environment after having imported wandb + and called certain functions. + """ + if _singleton and _singleton._pid == os.getpid(): + return _singleton + else: + return None + + +@wb_logging.log_to_all_runs() +def _setup( + settings: Settings | None = None, + start_service: bool = True, +) -> _WandbSetup: + """Set up library context. + + Args: + settings: Global settings to set, or updates to the global settings + if the singleton has already been initialized. + start_service: Whether to start up the service process. + NOTE: A service process will only be started if allowed by the + global settings (after the given updates). The service will not + start up if the mode resolves to "disabled". + """ + global _singleton + + pid = os.getpid() + + if _singleton and _singleton._pid == pid: + _singleton._update(settings=settings) + else: + _singleton = _WandbSetup(settings=settings, pid=pid) + + if start_service and not _singleton.settings._noop: + _singleton.ensure_service() + + return _singleton + + +def setup(settings: Settings | None = None) -> _WandbSetup: + """Prepares W&B for use in the current process and its children. + + You can usually ignore this as it is implicitly called by `wandb.init()`. + + When using wandb in multiple processes, calling `wandb.setup()` + in the parent process before starting child processes may improve + performance and resource utilization. + + Note that `wandb.setup()` modifies `os.environ`, and it is important + that child processes inherit the modified environment variables. + + See also `wandb.teardown()`. + + Args: + settings: Configuration settings to apply globally. These can be + overridden by subsequent `wandb.init()` calls. + + Example: + ```python + import multiprocessing + + import wandb + + + def run_experiment(params): + with wandb.init(config=params): + # Run experiment + pass + + + if __name__ == "__main__": + # Start backend and set global config + wandb.setup(settings={"project": "my_project"}) + + # Define experiment parameters + experiment_params = [ + {"learning_rate": 0.01, "epochs": 10}, + {"learning_rate": 0.001, "epochs": 20}, + ] + + # Start multiple processes, each running a separate experiment + processes = [] + for params in experiment_params: + p = multiprocessing.Process(target=run_experiment, args=(params,)) + p.start() + processes.append(p) + + # Wait for all processes to complete + for p in processes: + p.join() + + # Optional: Explicitly shut down the backend + wandb.teardown() + ``` + """ + return _setup(settings=settings) + + +@wb_logging.log_to_all_runs() +def teardown(exit_code: int | None = None) -> None: + """Waits for wandb to finish and frees resources. + + Completes any runs that were not explicitly finished + using `run.finish()` and waits for all data to be uploaded. + + It is recommended to call this at the end of a session + that used `wandb.setup()`. It is invoked automatically + in an `atexit` hook, but this is not reliable in certain setups + such as when using Python's `multiprocessing` module. + """ + global _singleton + + orig_singleton = _singleton + _singleton = None + + if orig_singleton: + orig_singleton._teardown(exit_code=exit_code) diff --git a/lib/python3.12/site-packages/wandb/sdk/wandb_summary.py b/lib/python3.12/site-packages/wandb/sdk/wandb_summary.py new file mode 100644 index 0000000000000000000000000000000000000000..f15c7cc041959115b0cfc0024f5ec2b690209558 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/wandb_summary.py @@ -0,0 +1,150 @@ +import abc +import typing as t + +from .interface.summary_record import SummaryItem, SummaryRecord + + +def _get_dict(d): + if isinstance(d, dict): + return d + # assume argparse Namespace + return vars(d) + + +class SummaryDict(metaclass=abc.ABCMeta): + """dict-like wrapper for the nested dictionaries in a SummarySubDict. + + Triggers self._root._callback on property changes. + """ + + @abc.abstractmethod + def _as_dict(self): + raise NotImplementedError + + @abc.abstractmethod + def _update(self, record: SummaryRecord): + raise NotImplementedError + + def keys(self): + return [k for k in self._as_dict().keys() if k != "_wandb"] + + def get(self, key, default=None): + return self._as_dict().get(key, default) + + def __getitem__(self, key): + item = self._as_dict()[key] + + if isinstance(item, dict): + # this nested dict needs to be wrapped: + wrapped_item = SummarySubDict() + object.__setattr__(wrapped_item, "_items", item) + object.__setattr__(wrapped_item, "_parent", self) + object.__setattr__(wrapped_item, "_parent_key", key) + + return wrapped_item + + # this item isn't a nested dict + return item + + __getattr__ = __getitem__ + + def __setitem__(self, key, val): + self.update({key: val}) + + __setattr__ = __setitem__ + + def __delattr__(self, key): + record = SummaryRecord() + item = SummaryItem() + item.key = (key,) + record.remove = (item,) + self._update(record) + + __delitem__ = __delattr__ + + def update(self, d: t.Dict): + # import ipdb; ipdb.set_trace() + record = SummaryRecord() + for key, value in d.items(): + item = SummaryItem() + item.key = (key,) + item.value = value + record.update.append(item) + + self._update(record) + + +class Summary(SummaryDict): + """Track single values for each metric for each run. + + By default, a metric's summary is the last value of its History. + + For example, `wandb.log({'accuracy': 0.9})` will add a new step to History and + update Summary to the latest value. In some cases, it's more useful to have + the maximum or minimum of a metric instead of the final value. You can set + history manually `(wandb.summary['accuracy'] = best_acc)`. + + In the UI, summary metrics appear in the table to compare across runs. + Summary metrics are also used in visualizations like the scatter plot and + parallel coordinates chart. + + After training has completed, you may want to save evaluation metrics to a + run. Summary can handle numpy arrays and PyTorch/TensorFlow tensors. When + you save one of these types to Summary, we persist the entire tensor in a + binary file and store high level metrics in the summary object, such as min, + mean, variance, and 95th percentile. + + Examples: + ```python + wandb.init(config=args) + + best_accuracy = 0 + for epoch in range(1, args.epochs + 1): + test_loss, test_accuracy = test() + if test_accuracy > best_accuracy: + wandb.run.summary["best_accuracy"] = test_accuracy + best_accuracy = test_accuracy + ``` + """ + + _update_callback: t.Callable + _get_current_summary_callback: t.Callable + + def __init__(self, get_current_summary_callback: t.Callable): + super().__init__() + object.__setattr__(self, "_update_callback", None) + object.__setattr__( + self, "_get_current_summary_callback", get_current_summary_callback + ) + + def _set_update_callback(self, update_callback: t.Callable): + object.__setattr__(self, "_update_callback", update_callback) + + def _as_dict(self): + return self._get_current_summary_callback() + + def _update(self, record: SummaryRecord): + if self._update_callback: # type: ignore + self._update_callback(record) + + +class SummarySubDict(SummaryDict): + """Non-root node of the summary data structure. + + Contains a path to itself from the root. + """ + + _items: t.Dict + _parent: SummaryDict + _parent_key: str + + def __init__(self): + object.__setattr__(self, "_items", dict()) + object.__setattr__(self, "_parent", None) + object.__setattr__(self, "_parent_key", None) + + def _as_dict(self): + return self._items + + def _update(self, record: SummaryRecord): + return self._parent._update(record._add_next_parent(self._parent_key)) diff --git a/lib/python3.12/site-packages/wandb/sdk/wandb_sweep.py b/lib/python3.12/site-packages/wandb/sdk/wandb_sweep.py new file mode 100644 index 0000000000000000000000000000000000000000..46aeecbc88c138372419f7a59fdada25c7062c1f --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/wandb_sweep.py @@ -0,0 +1,119 @@ +import urllib.parse +from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Union + +import wandb +from wandb import env +from wandb.apis import InternalApi +from wandb.sdk.launch.sweeps.utils import handle_sweep_config_violations + +from . import wandb_login + +if TYPE_CHECKING: + from wandb.wandb_controller import _WandbController + + +def _get_sweep_url(api, sweep_id): + """Return sweep url if we can figure it out.""" + if api.api_key: + if api.settings("entity") is None: + viewer = api.viewer() + if viewer.get("entity"): + api.set_setting("entity", viewer["entity"]) + project = api.settings("project") + if not project: + return + if api.settings("entity"): + return "{base}/{entity}/{project}/sweeps/{sweepid}".format( + base=api.app_url, + entity=urllib.parse.quote(api.settings("entity")), + project=urllib.parse.quote(project), + sweepid=urllib.parse.quote(sweep_id), + ) + + +def sweep( + sweep: Union[dict, Callable], + entity: Optional[str] = None, + project: Optional[str] = None, + prior_runs: Optional[List[str]] = None, +) -> str: + """Initialize a hyperparameter sweep. + + Search for hyperparameters that optimizes a cost function + of a machine learning model by testing various combinations. + + Make note the unique identifier, `sweep_id`, that is returned. + At a later step provide the `sweep_id` to a sweep agent. + + Args: + sweep: The configuration of a hyperparameter search. + (or configuration generator). See + [Sweep configuration structure](https://docs.wandb.ai/guides/sweeps/define-sweep-configuration) + for information on how to define your sweep. + If you provide a callable, ensure that the callable does + not take arguments and that it returns a dictionary that + conforms to the W&B sweep config spec. + entity: The username or team name where you want to send W&B + runs created by the sweep to. Ensure that the entity you + specify already exists. If you don't specify an entity, + the run will be sent to your default entity, + which is usually your username. + project: The name of the project where W&B runs created from + the sweep are sent to. If the project is not specified, the + run is sent to a project labeled 'Uncategorized'. + prior_runs: The run IDs of existing runs to add to this sweep. + + Returns: + sweep_id: str. A unique identifier for the sweep. + """ + if callable(sweep): + sweep = sweep() + """Sweep create for controller api and jupyter (eventually for cli).""" + + # Project may be only found in the sweep config. + if project is None and isinstance(sweep, dict): + project = sweep.get("project", None) + + if entity: + env.set_entity(entity) + if project: + env.set_project(project) + + # Make sure we are logged in + if wandb.run is None: + wandb_login._login(_silent=True) + api = InternalApi() + sweep_id, warnings = api.upsert_sweep(sweep, prior_runs=prior_runs) + handle_sweep_config_violations(warnings) + print("Create sweep with ID:", sweep_id) # noqa: T201 + sweep_url = _get_sweep_url(api, sweep_id) + if sweep_url: + print("Sweep URL:", sweep_url) # noqa: T201 + return sweep_id + + +def controller( + sweep_id_or_config: Optional[Union[str, Dict]] = None, + entity: Optional[str] = None, + project: Optional[str] = None, +) -> "_WandbController": + """Public sweep controller constructor. + + Usage: + ```python + import wandb + + tuner = wandb.controller(...) + print(tuner.sweep_config) + print(tuner.sweep_id) + tuner.configure_search(...) + tuner.configure_stopping(...) + ``` + + """ + from ..wandb_controller import _WandbController + + c = _WandbController( + sweep_id_or_config=sweep_id_or_config, entity=entity, project=project + ) + return c diff --git a/lib/python3.12/site-packages/wandb/sdk/wandb_sync.py b/lib/python3.12/site-packages/wandb/sdk/wandb_sync.py new file mode 100644 index 0000000000000000000000000000000000000000..9c5a3fde47bdda06ba979f912eae0ffc31502548 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/wandb_sync.py @@ -0,0 +1,71 @@ +import pathlib +from typing import TYPE_CHECKING, Optional + +from wandb import util +from wandb.errors.term import termerror, termlog + +from . import wandb_setup +from .backend.backend import Backend +from .lib.runid import generate_id + +if TYPE_CHECKING: + from wandb.proto import wandb_internal_pb2 + + +def _sync( + path: str, + run_id: Optional[str] = None, + project: Optional[str] = None, + entity: Optional[str] = None, + mark_synced: Optional[bool] = None, + append: Optional[bool] = None, + skip_console: Optional[bool] = None, +) -> "wandb_internal_pb2.SyncResponse": + wl = wandb_setup.setup() + assert wl is not None + + stream_id = generate_id() + + settings = wl.settings.to_proto() + p = pathlib.Path(path) + + # update sync_file setting to point to the passed path + settings.sync_file.value = str(p.absolute()) + settings.sync_dir.value = str(p.parent.absolute()) + settings.files_dir.value = str(p.parent.absolute() / "files") + settings.x_sync.value = True + if run_id: + settings.run_id.value = run_id + if entity: + settings.entity.value = entity + if project: + settings.project.value = project + if skip_console: + settings.console.value = "off" + if append: + settings.resume.value = "allow" + + service = wl.ensure_service() + service.inform_init(settings=settings, run_id=stream_id) + + backend = Backend(settings=wl.settings, service=service) + backend.ensure_launched() + + assert backend.interface + backend.interface._stream_id = stream_id # type: ignore + + handle = backend.interface.deliver_finish_sync() + result = handle.wait_or(timeout=None) + response = result.response.sync_response + if response.url: + termlog(f"Synced {p} to {util.app_url(response.url)}") + # create a .synced file in the directory if mark_synced is true + if mark_synced: + with open(f"{p}.synced", "w"): + pass + else: + termerror(f"Failed to sync {p}") + if response.error and response.error.message: + termerror(response.error.message) + + return response diff --git a/lib/python3.12/site-packages/wandb/sdk/wandb_watch.py b/lib/python3.12/site-packages/wandb/sdk/wandb_watch.py new file mode 100644 index 0000000000000000000000000000000000000000..4a29cf1b6ccea8074d57ee414b961326bdc1cf2e --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sdk/wandb_watch.py @@ -0,0 +1,148 @@ +"""watch.""" + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Sequence + +try: + from typing import Literal +except ImportError: + from typing_extensions import Literal # type: ignore + +import wandb + +from .lib import telemetry + +if TYPE_CHECKING: + import torch # type: ignore [import-not-found] + + from wandb.sdk.wandb_run import Run + +logger = logging.getLogger("wandb") + +_global_watch_idx = 0 + + +def _watch( + run: Run, + models: torch.nn.Module | Sequence[torch.nn.Module], + criterion: torch.F | None = None, + log: Literal["gradients", "parameters", "all"] | None = "gradients", + log_freq: int = 1000, + idx: int | None = None, + log_graph: bool = False, +): + """Hooks into the given PyTorch model(s) to monitor gradients and the model's computational graph. + + This function can track parameters, gradients, or both during training. It should be + extended to support arbitrary machine learning models in the future. + + Args: + run (wandb.sdk.wandb_run.Run): The run object to log to. + models (Union[torch.nn.Module, Sequence[torch.nn.Module]]): + A single model or a sequence of models to be monitored. + criterion (Optional[torch.F]): + The loss function being optimized (optional). + log (Optional[Literal["gradients", "parameters", "all"]]): + Specifies whether to log "gradients", "parameters", or "all". + Set to None to disable logging. (default="gradients") + log_freq (int): + Frequency (in batches) to log gradients and parameters. (default=1000) + idx (Optional[int]): + Index used when tracking multiple models with `wandb.watch`. (default=None) + log_graph (bool): + Whether to log the model's computational graph. (default=False) + + Returns: + wandb.Graph: + The graph object, which will be populated after the first backward pass. + + Raises: + ValueError: If `wandb.init` has not been called. + TypeError: If any of the models are not instances of `torch.nn.Module`. + """ + global _global_watch_idx + + with telemetry.context() as tel: + tel.feature.watch = True + + logger.info("Watching") + + if log not in {"gradients", "parameters", "all", None}: + raise ValueError("log must be one of 'gradients', 'parameters', 'all', or None") + + log_parameters = log in {"parameters", "all"} + log_gradients = log in {"gradients", "all"} + + if not isinstance(models, (tuple, list)): + models = (models,) + + torch = wandb.util.get_module( + "torch", required="wandb.watch only works with pytorch, couldn't import torch." + ) + + for model in models: + if not isinstance(model, torch.nn.Module): + raise TypeError( + f"Expected a pytorch model (torch.nn.Module). Received {type(model)}" + ) + + graphs = [] + prefix = "" + + if idx is None: + idx = _global_watch_idx + for local_idx, model in enumerate(models): + global_idx = idx + local_idx + _global_watch_idx += 1 + if global_idx > 0: + # TODO: this makes ugly chart names like gradients/graph_1conv1d.bias + prefix = f"graph_{global_idx}" + + if log_parameters: + run._torch.add_log_parameters_hook( + model, + prefix=prefix, + log_freq=log_freq, + ) + + if log_gradients: + run._torch.add_log_gradients_hook( + model, + prefix=prefix, + log_freq=log_freq, + ) + + if log_graph: + graph = run._torch.hook_torch(model, criterion, graph_idx=global_idx) + graphs.append(graph) + # NOTE: the graph is set in run.summary by hook_torch on the backward pass + return graphs + + +def _unwatch( + run: Run, models: torch.nn.Module | Sequence[torch.nn.Module] | None = None +) -> None: + """Remove pytorch model topology, gradient and parameter hooks. + + Args: + run (wandb.sdk.wandb_run.Run): + The run object to log to. + models (torch.nn.Module | Sequence[torch.nn.Module]): + Optional list of pytorch models that have had watch called on them + """ + if models: + if not isinstance(models, (tuple, list)): + models = (models,) + for model in models: + if not hasattr(model, "_wandb_hook_names"): + wandb.termwarn(f"{model} model has not been watched") + else: + for name in model._wandb_hook_names: + run._torch.unhook(name) + delattr(model, "_wandb_hook_names") + # TODO: we should also remove recursively model._wandb_watch_called + + else: + run._torch.unhook_all()