split: further split trace and legacy env

wandb · Aug 27, 2024 · 34958ba · 34958ba
1 parent 11003d5
commit 34958ba
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/weave/trace/env.py b/weave/trace/env.py
@@ -5,3 +5,7 @@
 
 def get_weave_parallelism() -> int:
     return int(os.getenv(WEAVE_PARALLELISM, "20"))
+
+
+def wandb_production() -> bool:
+    return os.getenv("WEAVE_ENV") == "wandb_production"
diff --git a/weave/trace/op_type.py b/weave/trace/op_type.py
@@ -14,6 +14,7 @@
 from typing import Any, Callable, Optional, Union, get_args, get_origin
 
 from weave.legacy.weave import artifact_fs, context_state, errors, storage
+from weave.trace import env as environment
 from weave.trace.ipython import (
     ClassNotFoundError,
     get_class_source,
@@ -22,7 +23,6 @@
 from weave.trace.op import Op
 from weave.trace.refs import ObjectRef
 
-from ..legacy.weave import environment
 from . import serializer
 
 WEAVE_OP_PATTERN = re.compile(r"@weave\.op(\(\))?")