Skip to content

Commit

Permalink
fix: WEAVE_DISABLE now disables all weave tracing, including the init…
Browse files Browse the repository at this point in the history
…ialization requests (#2368)

* init

* init

* working away

* init

* working

* style

* reset docs

* reset docs

* docs

* docs

* init
  • Loading branch information
tssweeney authored Sep 16, 2024
1 parent 3549374 commit 412efc3
Show file tree
Hide file tree
Showing 4 changed files with 78 additions and 9 deletions.
21 changes: 21 additions & 0 deletions weave/tests/trace/test_trace_settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
10 changes: 5 additions & 5 deletions weave/trace/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand All @@ -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


Expand Down
28 changes: 25 additions & 3 deletions weave/trace/settings.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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.
Expand Down Expand Up @@ -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]
28 changes: 27 additions & 1 deletion weave/trace/weave_init.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 412efc3

Please sign in to comment.