wandb 0.17.6__py3-none-win32.whl → 0.17.8__py3-none-win32.whl

wandb/sdk/wandb_run.py

@@ -42,6 +42,7 @@ from wandb._globals import _datatypes_set_callback
 from wandb.apis import internal, public
 from wandb.apis.internal import Api
 from wandb.apis.public import Api as PublicApi
+from wandb.errors import CommError
 from wandb.proto.wandb_internal_pb2 import (
     MetricRecord,
     PollExitResponse,
@@ -69,7 +70,7 @@ from wandb.viz import CustomChart, Visualize, custom_chart
 
 from . import wandb_config, wandb_metric, wandb_summary
 from .data_types._dtypes import TypeRegistry
-from .interface.interface import GlobStr, InterfaceBase
+from .interface.interface import FilesDict, GlobStr, InterfaceBase, PolicyName
 from .interface.summary_record import SummaryRecord
 from .lib import (
     config_util,
@@ -89,6 +90,7 @@ from .lib.printer import get_printer
 from .lib.proto_util import message_to_dict
 from .lib.reporting import Reporter
 from .lib.wburls import wburls
+from .wandb_alerts import AlertLevel
 from .wandb_settings import Settings
 from .wandb_setup import _WandbSetup
 
@@ -108,9 +110,7 @@ if TYPE_CHECKING:
         SampledHistoryResponse,
     )
 
-    from .interface.interface import FilesDict, PolicyName
     from .lib.printer import PrinterJupyter, PrinterTerm
-    from .wandb_alerts import AlertLevel
 
     class GitSourceDict(TypedDict):
         remote: str
@@ -171,10 +171,12 @@ class RunStatusChecker:
         interface: InterfaceBase,
         stop_polling_interval: int = 15,
         retry_polling_interval: int = 5,
+        internal_messages_polling_interval: int = 10,
     ) -> None:
         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._join_event = threading.Event()
 
@@ -295,7 +297,7 @@ class RunStatusChecker:
             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():
+                if not wandb.agents.pyagent.is_running():  # type: ignore
                     thread.interrupt_main()
                     return
 
@@ -323,7 +325,7 @@ class RunStatusChecker:
             self._loop_check_status(
                 lock=self._internal_messages_lock,
                 set_handle=lambda x: setattr(self, "_internal_messages_handle", x),
-                timeout=1,
+                timeout=self._internal_messages_polling_interval,
                 request=self._interface.deliver_internal_messages,
                 process=_process_internal_messages,
             )
@@ -621,7 +623,7 @@ class Run:
         )
         self.summary._set_update_callback(self._summary_update_callback)
         self._step = 0
-        self._torch_history: Optional[wandb.wandb_torch.TorchHistory] = None
+        self._torch_history: Optional[wandb.wandb_torch.TorchHistory] = None  # type: ignore
 
         # 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,
@@ -796,11 +798,15 @@ class Run:
                     self._unique_launch_artifact_sequence_names[sequence_name] = item
 
     def _telemetry_callback(self, telem_obj: telemetry.TelemetryRecord) -> None:
+        if not hasattr(self, "_telemetry_obj"):
+            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:
@@ -920,9 +926,9 @@ class Run:
         self.__dict__.update(state)
 
     @property
-    def _torch(self) -> "wandb.wandb_torch.TorchHistory":
+    def _torch(self) -> "wandb.wandb_torch.TorchHistory":  # type: ignore
         if self._torch_history is None:
-            self._torch_history = wandb.wandb_torch.TorchHistory()
+            self._torch_history = wandb.wandb_torch.TorchHistory()  # type: ignore
         return self._torch_history
 
     @property
@@ -1089,13 +1095,14 @@ class Run:
     @_run_decorator._attach
     def mode(self) -> str:
         """For compatibility with `0.9.x` and earlier, deprecate eventually."""
-        deprecate.deprecate(
-            field_name=deprecate.Deprecated.run__mode,
-            warning_message=(
-                "The mode property of wandb.run is deprecated "
-                "and will be removed in a future release."
-            ),
-        )
+        if hasattr(self, "_telemetry_obj"):
+            deprecate.deprecate(
+                field_name=deprecate.Deprecated.run__mode,
+                warning_message=(
+                    "The mode property of wandb.run is deprecated "
+                    "and will be removed in a future release."
+                ),
+            )
         return "dryrun" if self._settings._offline else "run"
 
     @property
@@ -1675,15 +1682,15 @@ class Run:
         commit: Optional[bool] = None,
         sync: Optional[bool] = None,
     ) -> None:
-        """Log a dictionary of data to the current run's history.
+        """Upload run data.
 
-        Use `wandb.log` to log data from runs, such as scalars, images, video,
+        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 `wandb.log({"train-loss": 0.5, "accuracy": 0.9})`.
+        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.
 
@@ -1692,48 +1699,91 @@ class Run:
         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).
 
-        In the UI, summary values show up in the run table to compare single values across runs.
-        Summary values can also be set directly with `wandb.run.summary["key"] = value`.
-
         Logged values don't have to be scalars. Logging any wandb object is supported.
-        For example `wandb.log({"example": wandb.Image("myimage.jpg")})` will log an
+        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.
-        `wandb.Table`s can be used to logged structured data. See our
+        You can use `wandb.Table` to log structured data. See our
         [guide to logging tables](https://docs.wandb.ai/guides/data-vis/log-tables)
         for details.
 
-        Logging nested metrics is encouraged and is supported in the W&B UI.
-        If you log with a nested dictionary like `wandb.log({"train":
-        {"acc": 0.9}, "val": {"acc": 0.8}})`, the metrics will be organized into
-        `train` and `val` sections in the W&B UI.
+        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,
+        })
+        ```
 
-        wandb keeps track of a global step, which by default increments with each
-        call to `wandb.log`, so logging related metrics together is encouraged.
-        If it's inconvenient to log related metrics together
-        calling `wandb.log({"train-loss": 0.5}, commit=False)` and then
-        `wandb.log({"accuracy": 0.9})` is equivalent to calling
-        `wandb.log({"train-loss": 0.5, "accuracy": 0.9})`.
+        Only one level of nesting is supported; `run.log({"a/b/c": 1})`
+        produces a section named "a/b".
 
-        `wandb.log` is not intended to be called more than a few times per second.
-        If you want to log more frequently than that it's better to aggregate
-        the data on the client side or you may get degraded performance.
+        `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)
+        ```
 
         Arguments:
-            data: (dict, optional) A dict of serializable python objects i.e `str`,
-                `ints`, `floats`, `Tensors`, `dicts`, or any of the `wandb.data_types`.
-            commit: (boolean, optional) Save the metrics dict to the wandb server
-                and increment the step.  If false `wandb.log` just updates the current
-                metrics dict with the data argument and metrics won't be saved until
-                `wandb.log` is called with `commit=True`.
-            step: (integer, optional) The global step in processing. This persists
-                any non-committed earlier steps but defaults to not committing the
-                specified step.
-            sync: (boolean, True) This argument is deprecated and currently doesn't
-                change the behaviour of `wandb.log`.
+            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`.
+            sync: This argument is deprecated and does nothing.
 
         Examples:
             For more and more detailed examples, see
@@ -1885,7 +1935,7 @@ class Run:
         self,
         glob_str: Optional[Union[str, os.PathLike]] = None,
         base_path: Optional[Union[str, os.PathLike]] = None,
-        policy: "PolicyName" = "live",
+        policy: PolicyName = "live",
     ) -> Union[bool, List[str]]:
         """Sync one or more files to W&B.
 
@@ -2158,12 +2208,13 @@ class Run:
     @_run_decorator._attach
     def join(self, exit_code: Optional[int] = None) -> None:
         """Deprecated alias for `finish()` - use finish instead."""
-        deprecate.deprecate(
-            field_name=deprecate.Deprecated.run__join,
-            warning_message=(
-                "wandb.run.join() is deprecated, please use wandb.run.finish()."
-            ),
-        )
+        if hasattr(self, "_telemetry_obj"):
+            deprecate.deprecate(
+                field_name=deprecate.Deprecated.run__join,
+                warning_message=(
+                    "wandb.run.join() is deprecated, please use wandb.run.finish()."
+                ),
+            )
         self._finish(exit_code=exit_code)
 
     @_run_decorator._noop_on_finish()
@@ -2368,11 +2419,7 @@ class Run:
             return
         self._atexit_cleanup_called = True
 
-        exit_code = (
-            exit_code  #
-            or (self._hooks and self._hooks.exit_code)
-            or 0
-        )
+        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}")
 
@@ -2389,7 +2436,7 @@ class Run:
             self._on_finish()
 
         except KeyboardInterrupt:
-            if not wandb.wandb_agent._is_running():
+            if not wandb.wandb_agent._is_running():  # type: ignore
                 wandb.termerror("Control-C detected -- Run data was not synced")
             raise
 
@@ -2433,18 +2480,6 @@ class Run:
     def _on_init(self) -> None:
         if self._settings._offline:
             return
-        if self._backend and self._backend.interface:
-            if not self._settings._disable_update_check:
-                logger.info("communicating current version")
-                version_handle = self._backend.interface.deliver_check_version(
-                    current_version=wandb.__version__
-                )
-                version_result = version_handle.wait(timeout=30)
-                if not version_result:
-                    version_handle.abandon()
-                else:
-                    self._check_version = version_result.response.check_version_response
-                    logger.info("got version response %s", self._check_version)
 
     def _on_start(self) -> None:
         # would like to move _set_global to _on_ready to unify _on_start and _on_attach
@@ -2453,9 +2488,7 @@ class Run:
         # TODO(jupyter) However _header calls _header_run_info that uses wandb.jupyter that uses
         #               `wandb.run` and hence breaks
         self._set_globals()
-        self._header(
-            self._check_version, settings=self._settings, printer=self._printer
-        )
+        self._header(settings=self._settings, printer=self._printer)
 
         if self._settings.save_code and self._settings.code_dir is not None:
             self.log_code(self._settings.code_dir)
@@ -2640,6 +2673,18 @@ class Run:
 
         assert self._backend and self._backend.interface
 
+        if not self._settings._disable_update_check:
+            logger.info("communicating current version")
+            version_handle = self._backend.interface.deliver_check_version(
+                current_version=wandb.__version__
+            )
+            version_result = version_handle.wait(timeout=10)
+            if not version_result:
+                version_handle.abandon()
+            else:
+                self._check_version = version_result.response.check_version_response
+                logger.info("got version response %s", self._check_version)
+
         # get the server info before starting the defer state machine as
         # it will stop communication with the server
         server_info_handle = self._backend.interface.deliver_request_server_info()
@@ -2839,12 +2884,12 @@ class Run:
         idx=None,
         log_graph=False,
     ) -> None:
-        wandb.watch(models, criterion, log, log_freq, idx, log_graph)
+        wandb.watch(models, criterion, log, log_freq, idx, log_graph)  # type: ignore
 
     # TODO(jhr): annotate this
     @_run_decorator._attach
     def unwatch(self, models=None) -> None:  # type: ignore
-        wandb.unwatch(models=models)
+        wandb.unwatch(models=models)  # type: ignore
 
     # TODO(kdg): remove all artifact swapping logic
     def _swap_artifact_name(self, artifact_name: str, use_as: Optional[str]) -> str:
@@ -3514,7 +3559,7 @@ class Run:
             artifact = self._log_artifact(
                 artifact_or_path=path, name=name, type=artifact.type
             )
-        except (ValueError, wandb.CommError):
+        except (ValueError, CommError):
             artifact = self._log_artifact(
                 artifact_or_path=path, name=name, type="model"
             )
@@ -3534,13 +3579,13 @@ class Run:
         Arguments:
             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 wandb.AlertLevel, optional) The alert level to use, either: `INFO`, `WARN`, or `ERROR`.
+            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 wandb.AlertLevel.INFO
-        level_str: str = level.value if isinstance(level, wandb.AlertLevel) else level
-        if level_str not in {lev.value for lev in wandb.AlertLevel}:
+        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)
@@ -3633,14 +3678,10 @@ class Run:
     # with the service execution path that doesn't have access to the run instance
     @staticmethod
     def _header(
-        check_version: Optional["CheckVersionResponse"] = None,
         *,
         settings: "Settings",
         printer: Union["PrinterTerm", "PrinterJupyter"],
     ) -> None:
-        Run._header_version_check_info(
-            check_version, settings=settings, printer=printer
-        )
         Run._header_wandb_version_info(settings=settings, printer=printer)
         Run._header_sync_info(settings=settings, printer=printer)
         Run._header_run_info(settings=settings, printer=printer)
@@ -3661,7 +3702,9 @@ class Run:
             printer.display(check_version.yank_message, level="warn")
 
         printer.display(
-            check_version.upgrade_message, off=not check_version.upgrade_message
+            check_version.upgrade_message,
+            off=not check_version.upgrade_message,
+            level="warn",
         )
 
     @staticmethod
@@ -3723,12 +3766,12 @@ class Run:
             return
 
         if printer._html:
-            if not wandb.jupyter.maybe_display():
+            if not wandb.jupyter.maybe_display():  # type: ignore
                 run_line = f"<strong>{printer.link(run_url, run_name)}</strong>"
                 project_line, sweep_line = "", ""
 
                 # TODO(settings): make settings the source of truth
-                if not wandb.jupyter.quiet():
+                if not wandb.jupyter.quiet():  # type: ignore
                     doc_html = printer.link(wburls.get("doc_run"), "docs")
 
                     project_html = printer.link(project_url, "Weights & Biases")
@@ -4177,9 +4220,11 @@ class Run:
             printer.display(check_version.yank_message, level="warn")
 
         # only display upgrade message if packages are bad
-        package_problem = check_version.delete_message or check_version.yank_message
-        if package_problem and check_version.upgrade_message:
-            printer.display(check_version.upgrade_message)
+        if check_version.upgrade_message:
+            printer.display(
+                check_version.upgrade_message,
+                level="warn",
+            )
 
     @staticmethod
     def _footer_notify_wandb_core(

wandb/sdk/wandb_settings.py

@@ -356,6 +356,7 @@ class SettingsData:
     _python: str
     _runqueue_item_id: str
     _require_core: bool
+    _require_legacy_service: bool
     _save_requirements: bool
     _service_transport: str
     _service_wait: float
@@ -650,7 +651,7 @@ class Settings(SettingsData):
             },
             _disable_service={
                 "value": False,
-                "preprocessor": _str_as_bool,
+                "preprocessor": self._process_disable_service,
                 "is_policy": True,
             },
             _disable_setproctitle={"value": False, "preprocessor": _str_as_bool},
@@ -718,6 +719,7 @@ class Settings(SettingsData):
                 "preprocessor": _str_as_json,
             },
             _require_core={"value": False, "preprocessor": _str_as_bool},
+            _require_legacy_service={"value": False, "preprocessor": _str_as_bool},
             _save_requirements={"value": True, "preprocessor": _str_as_bool},
             _service_wait={
                 "value": 30,
@@ -1172,6 +1174,16 @@ class Settings(SettingsData):
 
         return True
 
+    @staticmethod
+    def _process_disable_service(value: Union[str, bool]) -> bool:
+        value = _str_as_bool(value)
+        if value:
+            wandb.termwarn(
+                "Disabling the wandb service is deprecated as of version 0.18.0 and will be removed in version 0.19.0.",
+                repeat=False,
+            )
+        return value
+
     @staticmethod
     def _validate__service_wait(value: float) -> bool:
         if value <= 0:
@@ -1715,7 +1727,7 @@ class Settings(SettingsData):
 
         # 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)
+            meta = wandb.jupyter.notebook_metadata(self.silent)  # type: ignore
             settings["_jupyter_path"] = meta.get("path")
             settings["_jupyter_name"] = meta.get("name")
             settings["_jupyter_root"] = meta.get("root")
@@ -1875,9 +1887,21 @@ class Settings(SettingsData):
 
         # update settings
         self.update(init_settings, source=Source.INIT)
+        self._handle_fork_logic()
         self._handle_rewind_logic()
         self._handle_resume_logic()
 
+    def _handle_fork_logic(self) -> None:
+        if self.fork_from is None:
+            return
+
+        if self.run_id is not None and (self.fork_from.run == self.run_id):
+            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."
+            )
+
     def _handle_rewind_logic(self) -> None:
         if self.resume_from is None:
             return

wandb/sdk/wandb_setup.py

@@ -18,6 +18,7 @@ import threading
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
 
 import wandb
+from wandb.sdk.lib import import_hooks
 
 from . import wandb_manager, wandb_settings
 from .lib import config_util, server, tracelog
@@ -268,6 +269,8 @@ class _WandbSetup__WandbSetup:  # noqa: N801
                     self._config = config_dict
 
     def _teardown(self, exit_code: Optional[int] = None) -> None:
+        import_hooks.unregister_all_post_import_hooks()
+
         if not self._manager:
             return
 
@@ -319,14 +322,77 @@ def _setup(
     return wl
 
 
-def setup(
-    settings: Optional[Settings] = None,
-) -> Optional["_WandbSetup"]:
+def setup(settings: Optional[Settings] = None) -> Optional["_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 (Optional[Union[Dict[str, Any], wandb.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()
+        ```
+    """
     ret = _setup(settings=settings)
     return ret
 
 
 def teardown(exit_code: Optional[int] = 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.
+    """
     setup_instance = _WandbSetup._instance
     _WandbSetup._instance = None
 

wandb/sdk/wandb_sweep.py

@@ -1,5 +1,5 @@
 import urllib.parse
-from typing import Callable, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Union
 
 import wandb
 from wandb import env
@@ -8,6 +8,9 @@ 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."""
@@ -93,7 +96,7 @@ 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: