wandb 0.21.2__py3-none-macosx_12_0_x86_64.whl

wandb/agents/pyagent.py

@@ -0,0 +1,386 @@
+"""Agent - Agent object.
+
+Manage wandb agent.
+
+"""
+
+import ctypes
+import logging
+import os
+import queue
+import socket
+import sys
+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]
+                        # Extract to reduce a decision point to avoid ruff c901
+                        log_str, term_str = _get_exception_logger_and_term_strs(exc)
+                        logger.error(f"Run {run_id} errored:\n{log_str}")
+                        wandb.termerror(f"Run {run_id} errored:{term_str}")
+                        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"]))
+
+            try:
+                self._function()
+            except KeyboardInterrupt:
+                raise
+            except Exception as e:
+                # Log the run's exceptions directly to stderr to match CLI case, and wrap so we
+                # can identify it as coming from the job later later. This will get automatically
+                # logged by console_capture.py. Exception handler below will also handle exceptions
+                # in setup code.
+                exc_repr = _format_exception_traceback(e)
+                print(exc_repr, file=sys.stderr)  # noqa: T201
+                raise _JobError(f"Run threw exception: {str(e)}") from e
+            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()
+
+
+def _format_exception_traceback(exc):
+    return "".join(traceback.format_exception(type(exc), exc, exc.__traceback__))
+
+
+class _JobError(Exception):
+    """Exception raised when a job fails during execution."""
+
+    pass
+
+
+def _get_exception_logger_and_term_strs(exc):
+    if isinstance(exc, _JobError) and exc.__cause__:
+        # If it's a JobException, get the original exception for display
+        job_exc = exc.__cause__
+        log_str = _format_exception_traceback(job_exc)
+        # Don't long full stacktrace to terminal again because we already
+        # printed it to stderr.
+        term_str = " " + str(job_exc)
+    else:
+        log_str = _format_exception_traceback(exc)
+        term_str = "\n" + log_str
+    return log_str, term_str
+
+
+_INSTANCES = 0
+
+
+def is_running():
+    return bool(_INSTANCES)

wandb/analytics/__init__.py

@@ -0,0 +1,3 @@
+__all__ = ("Sentry",)
+
+from .sentry import Sentry

wandb/analytics/sentry.py

@@ -0,0 +1,267 @@
+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,
+        level: str = "info",
+    ) -> 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, level=level)  # 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()

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"]