wandb 0.17.0__py3-none-win32.whl

wandb/agents/pyagent.py

@@ -0,0 +1,364 @@
+"""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 import wandb_sdk
+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):  # noqa:C901
+        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 = "Detected {} failed runs in the first {} seconds, killing sweep.".format(
+                                self.FLAPPING_MAX_FAILURES, self.FLAPPING_MAX_SECONDS
+                            )
+                            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 = "Detected {} failed runs in a row at start, killing sweep.".format(
+                                self._max_initial_failures
+                            )
+                            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 as e:
+                    if self._exit_flag:
+                        logger.debug("Exiting main loop due to exit flag.")
+                        wandb.termlog("Sweep Agent: Killed.")
+                        return
+                    else:
+                        raise e
+        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_sdk.wandb_setup._setup(_reset=True)
+
+            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 as ki:
+            raise ki
+        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(
+            "Starting sweep agent: entity={}, project={}, count={}".format(
+                self._entity, self._project, 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.
+
+    Arguments:
+        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 Exception("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)

wandb/analytics/__init__.py

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

wandb/analytics/sentry.py

@@ -0,0 +1,266 @@
+__all__ = ("Sentry",)
+
+
+import atexit
+import functools
+import os
+import pathlib
+import sys
+from types import TracebackType
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple, Type, Union
+from urllib.parse import quote
+
+if sys.version_info >= (3, 8):
+    from typing import Literal
+else:
+    from typing_extensions import Literal
+
+import sentry_sdk  # type: ignore
+import sentry_sdk.utils  # type: ignore
+
+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.hub: Optional[sentry_sdk.hub.Hub] = 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.hub = sentry_sdk.Hub(client)
+
+    @_safe_noop
+    def message(self, message: str, repeat: bool = True) -> None:
+        """Send a message to Sentry."""
+        if not repeat and message in self._sent_messages:
+            return
+        self._sent_messages.add(message)
+        self.hub.capture_message(message)  # type: ignore
+
+    @_safe_noop
+    def exception(
+        self,
+        exc: Union[
+            str,
+            BaseException,
+            Tuple[
+                Optional[Type[BaseException]],
+                Optional[BaseException],
+                Optional[TracebackType],
+            ],
+            None,
+        ],
+        handled: bool = False,
+        status: Optional["SessionStatus"] = None,
+    ) -> None:
+        """Log an exception to Sentry."""
+        error = Exception(exc) if isinstance(exc, str) else exc
+        # based on self.hub.capture_exception(_exc)
+        if error is not None:
+            exc_info = sentry_sdk.utils.exc_info_from_error(error)
+        else:
+            exc_info = sys.exc_info()
+
+        event, hint = sentry_sdk.utils.event_from_exception(
+            exc_info,
+            client_options=self.hub.client.options,  # type: ignore
+            mechanism={"type": "generic", "handled": handled},
+        )
+        try:
+            self.hub.capture_event(event, hint=hint)  # 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.hub._stack[-1]  # type: ignore
+        if client is not None:
+            client.flush()
+
+        return None
+
+    def reraise(self, exc: Any) -> None:
+        """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.hub is not None
+        # get the current client and scope
+        _, scope = self.hub._stack[-1]
+        session = scope._session
+
+        # if there's no session, start one
+        if session is None:
+            self.hub.start_session()
+
+    @_safe_noop
+    def end_session(self) -> None:
+        """End the current session."""
+        assert self.hub is not None
+        # get the current client and scope
+        client, scope = self.hub._stack[-1]
+        session = scope._session
+
+        if session is not None and client is not None:
+            self.hub.end_session()
+            client.flush()
+
+    @_safe_noop
+    def mark_session(self, status: Optional["SessionStatus"] = None) -> None:
+        """Mark the current session with a status."""
+        assert self.hub is not None
+        _, scope = self.hub._stack[-1]
+        session = scope._session
+
+        if session is not None:
+            session.update(status=status)
+
+    @_safe_noop
+    def configure_scope(
+        self,
+        tags: Optional[Dict[str, Any]] = None,
+        process_context: Optional[str] = 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.hub is not None
+        settings_tags = (
+            "entity",
+            "project",
+            "run_id",
+            "run_url",
+            "sweep_url",
+            "sweep_id",
+            "deployment",
+            "_disable_service",
+            "_require_core",
+            "launch",
+        )
+
+        with self.hub.configure_scope() as scope:
+            scope.set_tag("platform", wandb.util.get_platform_name())
+
+            # set context
+            if process_context:
+                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, ""):
+                    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"
+            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
+                    scope.set_tag(
+                        obj_url,
+                        f"{app_url}/{entity}/{project}/{obj}s/{tags[obj_id]}",
+                    )
+                except Exception:
+                    pass
+
+            email = tags.get("email")
+            if email:
+                scope.user = {"email": email}  # noqa
+
+        # 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"]