From 412efc3a6d6e52d77fa9b517c8ba6c6999820c3a Mon Sep 17 00:00:00 2001
From: Tim Sweeney <tim@wandb.com>
Date: Mon, 16 Sep 2024 15:14:59 -0700
Subject: [PATCH] fix: WEAVE_DISABLE now disables all weave tracing, including
 the initialization requests (#2368)

* init

* init

* working away

* init

* working

* style

* reset docs

* reset docs

* docs

* docs

* init
---
 weave/tests/trace/test_trace_settings.py | 21 ++++++++++++++++++
 weave/trace/api.py                       | 10 ++++-----
 weave/trace/settings.py                  | 28 +++++++++++++++++++++---
 weave/trace/weave_init.py                | 28 +++++++++++++++++++++++-
 4 files changed, 78 insertions(+), 9 deletions(-)

diff --git a/weave/tests/trace/test_trace_settings.py b/weave/tests/trace/test_trace_settings.py
index bf7f29e92e0b..6686fff6bf32 100644
--- a/weave/tests/trace/test_trace_settings.py
+++ b/weave/tests/trace/test_trace_settings.py
@@ -37,6 +37,27 @@ def test_disabled_env(client):
     ), "Disabled weave should be faster than enabled weave"
 
 
+def test_disabled_env_client():
+    os.environ["WEAVE_DISABLED"] = "true"
+    client = weave.init("entity/project")
+
+    # Verify that the client is disabled
+    # Would be nicer to have a specific property
+    assert client.project == "DISABLED"
+
+    @weave.op
+    def func():
+        return 1
+
+    assert func() == 1
+
+    # No error implies that no calls were sent to the server
+    # since this would require writing to `entity/project`
+    client._flush()
+
+    os.environ["WEAVE_DISABLED"] = "false"
+
+
 def test_print_call_link_setting(client):
     captured_stdout = io.StringIO()
     sys.stdout = captured_stdout
diff --git a/weave/trace/api.py b/weave/trace/api.py
index 3aeb7bc07114..fa3ad9437e11 100644
--- a/weave/trace/api.py
+++ b/weave/trace/api.py
@@ -18,7 +18,7 @@
 from .constants import TRACE_OBJECT_EMOJI
 from .op import Op, op
 from .refs import ObjectRef, parse_uri
-from .settings import UserSettings, parse_and_apply_settings
+from .settings import UserSettings, parse_and_apply_settings, should_disable_weave
 from .table import Table
 
 
@@ -41,11 +41,11 @@ def init(
     Returns:
         A Weave client.
     """
-    # This is the stream-table backend. Disabling it in favor of the new
-    # trace-server backend.
-    # return weave_init.init_wandb(project_name).client
-    # return weave_init.init_trace_remote(project_name).client
     parse_and_apply_settings(settings)
+
+    if should_disable_weave():
+        return weave_init.init_weave_disabled().client
+
     return weave_init.init_weave(project_name).client
 
 
diff --git a/weave/trace/settings.py b/weave/trace/settings.py
index 333e0ae3b867..1d177eda40d2 100644
--- a/weave/trace/settings.py
+++ b/weave/trace/settings.py
@@ -1,8 +1,22 @@
 """Settings for Weave.
 
-To add new settings:
-1. Add a new field to `UserSettings`
-2. Add a new `should_{xyz}` function
+## `disabled`
+
+* Environment Variable: `WEAVE_DISABLED`
+* Settings Key: `disabled`
+* Default: `False`
+* Type: `bool`
+
+If True, all weave ops will behave like regular functions and no network requests will be made.
+
+## `print_call_link`
+
+* Environment Variable: `WEAVE_PRINT_CALL_LINK`
+* Settings Key: `print_call_link`
+* Default: `True`
+* Type: `bool`
+
+If True, prints a link to the Weave UI when calling a weave op.
 """
 
 import os
@@ -13,6 +27,11 @@
 
 SETTINGS_PREFIX = "WEAVE_"
 
+# Attention Devs:
+# To add new settings:
+# 1. Add a new field to `UserSettings`
+# 2. Add a new `should_{xyz}` function
+
 
 class UserSettings(BaseModel):
     """User configuration for Weave.
@@ -85,3 +104,6 @@ def _should(name: str) -> bool:
     if env := os.getenv(f"{SETTINGS_PREFIX}{name.upper()}"):
         return _str2bool_truthy(env)
     return _context_vars[name].get()
+
+
+__doc_spec__ = [UserSettings]
diff --git a/weave/trace/weave_init.py b/weave/trace/weave_init.py
index 6d35928bb44b..053a5f18324e 100644
--- a/weave/trace/weave_init.py
+++ b/weave/trace/weave_init.py
@@ -148,10 +148,36 @@ def init_weave(
     return _current_inited_client
 
 
+def init_weave_disabled() -> InitializedClient:
+    """Initialize a dummy client that does nothing.
+
+    This is used when the program is execuring with Weave disabled.
+
+    Note: as currently implemented, any explicit calls to client.{X} will
+    likely fail, since the user is not authenticated. The purpose of
+    disabling weave is to disable _tracing_. Programs that attempt to
+    make requests (eg. publishing, fetching, querying) while disabled
+    will fail.
+    """
+    global _current_inited_client
+    if _current_inited_client is not None:
+        _current_inited_client.reset()
+
+    client = weave_client.WeaveClient(
+        "DISABLED",
+        "DISABLED",
+        init_weave_get_server("DISABLED", should_batch=False),
+        ensure_project_exists=False,
+    )
+
+    return InitializedClient(client)
+
+
 def init_weave_get_server(
     api_key: typing.Optional[str] = None,
+    should_batch: bool = True,
 ) -> remote_http_trace_server.RemoteHTTPTraceServer:
-    res = remote_http_trace_server.RemoteHTTPTraceServer.from_env(True)
+    res = remote_http_trace_server.RemoteHTTPTraceServer.from_env(should_batch)
     if api_key is not None:
         res.set_auth(("api", api_key))
     return res