From a9bee1fa032eb35ade5ec26ef9d9dae7a35e7a60 Mon Sep 17 00:00:00 2001 From: Andrew Truong Date: Fri, 23 Aug 2024 01:24:51 -0400 Subject: [PATCH] move to legacy: engine_trace --- weave/filesystem.py | 4 +- weave/legacy/arrow/concat.py | 2 +- weave/legacy/artifact_wandb.py | 2 +- weave/legacy/cache.py | 4 +- weave/legacy/compile.py | 2 +- weave/{ => legacy}/engine_trace.py | 110 +++++++++--------- weave/legacy/execute.py | 2 +- weave/legacy/execute_fast.py | 2 +- weave/legacy/io_service.py | 2 +- .../language_features/tagging/tag_store.py | 4 +- weave/legacy/memo.py | 2 +- weave/legacy/op_def.py | 2 +- weave/legacy/ops_arrow/list_join.py | 2 +- .../run_history/history_op_common.py | 2 +- .../run_history/run_history_v1_legacy_ops.py | 2 +- .../run_history_v2_parquet_media.py | 2 +- ...run_history_v3_parquet_stream_optimized.py | 3 +- weave/legacy/ops_domain/run_ops.py | 2 +- weave/legacy/ops_domain/table.py | 3 +- weave/legacy/ops_domain/wbgqlquery_op.py | 4 +- weave/legacy/ops_domain/wbmedia.py | 4 +- weave/legacy/wandb_api.py | 4 +- weave/legacy/wandb_file_manager.py | 4 +- .../wandb_interface/wandb_artifact_pusher.py | 2 +- .../legacy/wandb_interface/wandb_lite_run.py | 4 +- weave/server.py | 3 +- weave/test_scripts/wandb_artifact_perf.py | 3 +- weave/weave_http.py | 4 +- weave/weave_server.py | 2 +- 29 files changed, 96 insertions(+), 92 deletions(-) rename weave/{ => legacy}/engine_trace.py (77%) diff --git a/weave/filesystem.py b/weave/filesystem.py index a3710b35165..8af0aac9bbb 100644 --- a/weave/filesystem.py +++ b/weave/filesystem.py @@ -20,9 +20,9 @@ from aiofiles.threadpool import binary as aiofiles_binary from aiofiles.threadpool import text as aiofiles_text -from weave.legacy import cache +from weave.legacy import cache, engine_trace -from . import engine_trace, environment, errors, util +from . import environment, errors, util tracer = engine_trace.tracer() # type: ignore async_utime = aiofiles_os.wrap(os.utime) # type: ignore diff --git a/weave/legacy/arrow/concat.py b/weave/legacy/arrow/concat.py index a22c5daf766..28198c3d339 100644 --- a/weave/legacy/arrow/concat.py +++ b/weave/legacy/arrow/concat.py @@ -30,7 +30,7 @@ DEBUG = False -from weave import engine_trace +from weave.legacy import engine_trace tracer = engine_trace.tracer() diff --git a/weave/legacy/artifact_wandb.py b/weave/legacy/artifact_wandb.py index 527d3c3c3ab..5a1378b8159 100644 --- a/weave/legacy/artifact_wandb.py +++ b/weave/legacy/artifact_wandb.py @@ -14,7 +14,6 @@ from wandb.sdk.lib.hashutil import b64_to_hex_id, hex_to_b64_id from weave import ( - engine_trace, errors, filesystem, urls, @@ -22,6 +21,7 @@ ) from weave import weave_types as types from weave.legacy import ( + engine_trace, artifact_fs, eager, file_base, diff --git a/weave/legacy/cache.py b/weave/legacy/cache.py index 2faefdde5bb..0bee0170fd3 100644 --- a/weave/legacy/cache.py +++ b/weave/legacy/cache.py @@ -6,8 +6,8 @@ import time import typing -from weave import engine_trace, environment, errors -from weave.legacy import context_state, wandb_api +from weave import environment, errors +from weave.legacy import context_state, wandb_api, engine_trace statsd = engine_trace.statsd() # type: ignore logger = logging.getLogger("root") diff --git a/weave/legacy/compile.py b/weave/legacy/compile.py index 30cbdf4d98f..eeec391beed 100644 --- a/weave/legacy/compile.py +++ b/weave/legacy/compile.py @@ -6,12 +6,12 @@ import typing from weave import ( - engine_trace, errors, weave_internal, ) from weave import weave_types as types from weave.legacy import ( + engine_trace, box, compile_domain, compile_table, diff --git a/weave/engine_trace.py b/weave/legacy/engine_trace.py similarity index 77% rename from weave/engine_trace.py rename to weave/legacy/engine_trace.py index 7423b8048ca..6a9284a2932 100644 --- a/weave/engine_trace.py +++ b/weave/legacy/engine_trace.py @@ -18,66 +18,66 @@ import time import typing -from . import environment, logs -from .legacy import stream_data_interfaces +from .. import environment, logs +from . import stream_data_interfaces # Thanks co-pilot! class DummySpan: - def __init__(self, *args, **kwargs): + def __init__(self, *args, **kwargs): # type: ignore self.args = args self.log_indent_token = None - def __enter__(self): + def __enter__(self): # type: ignore logging.debug("-> %s", self.args[0]) self.log_indent_token = logs.increment_indent() return self - def __exit__(self, *args, **kwargs): + def __exit__(self, *args, **kwargs): # type: ignore logs.reset_indent(self.log_indent_token) logging.debug("<- %s", self.args[0]) - def set_tag(self, key, unredacted_val, redacted_val=None): + def set_tag(self, key, unredacted_val, redacted_val=None): # type: ignore pass - def set_meta(self, *args, **kwargs): + def set_meta(self, *args, **kwargs): # type: ignore pass - def set_metric(self, *args, **kwargs): + def set_metric(self, *args, **kwargs): # type: ignore pass - def finish(self, *args, **kwargs): + def finish(self, *args, **kwargs): # type: ignore pass class TraceContext: - def __getstate__(self) -> dict: + def __getstate__(self) -> dict: # type: ignore return {} - def __setstate__(self, state: dict) -> None: + def __setstate__(self, state: dict) -> None: # type: ignore pass class ContextProvider: - def activate(self, context: TraceContext) -> None: + def activate(self, context: TraceContext) -> None: # type: ignore pass class DummyTrace: - def trace(self, *args, **kwargs): + def trace(self, *args, **kwargs): # type: ignore return DummySpan(*args) @property - def context_provider(self) -> ContextProvider: + def context_provider(self) -> ContextProvider: # type: ignore return ContextProvider() - def current_trace_context(self) -> typing.Optional[TraceContext]: + def current_trace_context(self) -> typing.Optional[TraceContext]: # type: ignore return None - def current_span(self): - return None + def current_span(self): # type: ignore + return None - def current_root_span(self): + def current_root_span(self): # type: ignore return None @@ -89,7 +89,7 @@ def current_root_span(self): _weave_trace_stream = None -def weave_trace_stream(): +def weave_trace_stream(): # type: ignore global _weave_trace_stream if _weave_trace_stream is None: from weave.legacy.wandb_interface.wandb_stream_table import StreamTable @@ -98,25 +98,25 @@ def weave_trace_stream(): return _weave_trace_stream -def span_count(span): +def span_count(span): # type: ignore if span.child_spans is None: return 1 return 1 + sum([span_count(child) for child in span.child_spans]) class WeaveTraceSpan: - def __init__(self, name): - from .legacy.ops_domain import trace_tree + def __init__(self, name): # type: ignore + from .ops_domain import trace_tree self.log_indent_token = None self._token = None self.span = trace_tree.Span(_name=name) self._parent = None - def total_span_count(self): + def total_span_count(self): # type: ignore return span_count(self.span) - def __enter__(self): + def __enter__(self): # type: ignore parent = _weave_trace_span.get() if parent is not None: if parent.span.child_spans is None: @@ -128,11 +128,11 @@ def __enter__(self): self.span.start_time_ms = int(time.time() * 1000) return self - def __exit__(self, *args, **kwargs): + def __exit__(self, *args, **kwargs): # type: ignore self.span.end_time_ms = int(time.time() * 1000) _weave_trace_span.reset(self._token) - from .legacy.ops_domain import trace_tree + from .ops_domain import trace_tree tt = trace_tree.WBTraceTree(json.dumps(dataclasses.asdict(self.span))) tags = self.attributes.get("tags", {}) @@ -147,12 +147,12 @@ def __exit__(self, *args, **kwargs): ) @property - def attributes(self): + def attributes(self): # type: ignore if self.span.attributes is None: self.span.attributes = {} return self.span.attributes - def set_tag(self, key, unredacted_val, redacted_val=None): + def set_tag(self, key, unredacted_val, redacted_val=None): # type: ignore if "tags" not in self.attributes: self.attributes["tags"] = {} if not environment.disable_weave_pii(): @@ -160,7 +160,7 @@ def set_tag(self, key, unredacted_val, redacted_val=None): elif redacted_val is not None: self.attributes["tags"][key] = redacted_val - def set_meta(self, key, unredacted_val, redacted_val=None): + def set_meta(self, key, unredacted_val, redacted_val=None): # type: ignore if "metadata" not in self.attributes: self.attributes["metadata"] = {} if not environment.disable_weave_pii(): @@ -168,7 +168,7 @@ def set_meta(self, key, unredacted_val, redacted_val=None): elif redacted_val is not None: self.attributes["metadata"][key] = redacted_val - def set_metric(self, key, unredacted_val, redacted_val=None): + def set_metric(self, key, unredacted_val, redacted_val=None): # type: ignore if "metrics" not in self.attributes: self.attributes["metrics"] = {} if not environment.disable_weave_pii(): @@ -176,15 +176,15 @@ def set_metric(self, key, unredacted_val, redacted_val=None): elif redacted_val is not None: self.attributes["metrics"][key] = redacted_val - def finish(self, *args, **kwargs): + def finish(self, *args, **kwargs): # type: ignore pass class WeaveTrace: - def __init__(self): + def __init__(self): # type: ignore self._stream_table = None - def trace(self, name, *args, **kwargs): + def trace(self, name, *args, **kwargs): # type: ignore return WeaveTraceSpan(name, *args) @property @@ -194,17 +194,17 @@ def context_provider(self) -> ContextProvider: def current_trace_context(self) -> typing.Optional[TraceContext]: return None - def current_span(self): + def current_span(self): # type: ignore return _weave_trace_span.get() - def current_root_span(self): + def current_root_span(self): # type: ignore cur_span = _weave_trace_span.get() while cur_span is not None and cur_span._parent is not None: cur_span = cur_span._parent return cur_span -def dd_span_to_weave_span(dd_span) -> stream_data_interfaces.TraceSpanDict: +def dd_span_to_weave_span(dd_span) -> stream_data_interfaces.TraceSpanDict: # type: ignore # Use '' for None, currently history2 doesn't read None columns from # the liveset correctly. parent_id = "" @@ -227,7 +227,7 @@ def dd_span_to_weave_span(dd_span) -> stream_data_interfaces.TraceSpanDict: } -def send_proc(queue): +def send_proc(queue): # type: ignore while True: spans = queue.get() if spans is None: @@ -244,35 +244,35 @@ def send_proc(queue): # by datadog, so we have some kind of re-entrancy/deadlock. Putting in a separate # process fixes. class WeaveWriter: - def __init__(self, orig_writer): + def __init__(self, orig_writer): # type: ignore self._orig_writer = orig_writer self._queue = multiprocessing.Queue() self._proc = multiprocessing.Process( target=send_proc, args=(self._queue,), daemon=True ) - def recreate(self): + def recreate(self): # type: ignore return WeaveWriter(self._orig_writer.recreate()) - def stop(self, timeout=None): + def stop(self, timeout=None): # type: ignore self._orig_writer.stop(timeout) - def _ensure_started(self): + def _ensure_started(self): # type: ignore if not self._proc.is_alive(): self._proc.start() - def write(self, spans): + def write(self, spans): # type: ignore if len(spans) == 1: return self._ensure_started() self._queue.put([dd_span_to_weave_span(s) for s in spans]) self._orig_writer.write(spans) - def flush_queue(self): + def flush_queue(self): # type: ignore self._orig_writer.flush_queue() -def patch_ddtrace_set_tag(): +def patch_ddtrace_set_tag(): # type: ignore from inspect import signature from ddtrace import span as ddtrace_span @@ -285,7 +285,7 @@ def patch_ddtrace_set_tag(): old_set_metric = ddtrace_span.Span.set_metric # Only logged redacted values if flag is on - def set_tag(self, key, unredacted_val=None, redacted_val=None): + def set_tag(self, key, unredacted_val=None, redacted_val=None): # type: ignore if redacted_val is not None and environment.disable_weave_pii(): old_set_tag(self, key, redacted_val) elif ( @@ -296,7 +296,7 @@ def set_tag(self, key, unredacted_val=None, redacted_val=None): old_set_tag(self, key, unredacted_val) # Log metric if flag is off or if flag is on and redacted - def set_metric(self, key, val, is_pii_redacted=False): + def set_metric(self, key, val, is_pii_redacted=False): # type: ignore if not environment.disable_weave_pii() or is_pii_redacted or "_dd." in key: old_set_metric(self, key, val) @@ -304,7 +304,7 @@ def set_metric(self, key, val, is_pii_redacted=False): ddtrace_span.Span.set_tag = set_tag -def tracer(): +def tracer(): # type: ignore if os.getenv("DD_ENV"): from ddtrace import tracer as ddtrace_tracer @@ -340,26 +340,26 @@ def new_trace_context() -> typing.Optional[TraceContext]: class DummyStatsd: - def __init__(self, *args, **kwargs): + def __init__(self, *args, **kwargs): # type: ignore pass - def increment(self, *args, **kwargs): + def increment(self, *args, **kwargs): # type: ignore pass - def decrement(self, *args, **kwargs): + def decrement(self, *args, **kwargs): # type: ignore pass - def gauge(self, *args, **kwargs): + def gauge(self, *args, **kwargs): # type: ignore pass - def flush(self, *args, **kwargs): + def flush(self, *args, **kwargs): # type: ignore pass _STATSD = None -def _initialize_statsd(): +def _initialize_statsd(): # type: ignore if os.getenv("DD_ENV"): from datadog import initialize, statsd @@ -369,7 +369,7 @@ def _initialize_statsd(): return DummyStatsd() -def statsd(): +def statsd(): # type: ignore global _STATSD if _STATSD is None: _STATSD = _initialize_statsd() @@ -377,5 +377,5 @@ def statsd(): return _STATSD -def datadog_is_enabled(): +def datadog_is_enabled(): # type: ignore return os.getenv("DD_ENV") diff --git a/weave/legacy/execute.py b/weave/legacy/execute.py index 20937630e3e..23176f1aaeb 100644 --- a/weave/legacy/execute.py +++ b/weave/legacy/execute.py @@ -10,7 +10,6 @@ from collections.abc import Mapping from weave import ( - engine_trace, environment, errors, ref_base, @@ -25,6 +24,7 @@ # Trace / cache # Language Features from weave.legacy import ( + engine_trace, parallelism, registry_mem, box, diff --git a/weave/legacy/execute_fast.py b/weave/legacy/execute_fast.py index dbf15feeadd..91a8e640b5b 100644 --- a/weave/legacy/execute_fast.py +++ b/weave/legacy/execute_fast.py @@ -1,13 +1,13 @@ import logging from weave import ( - engine_trace, errors, ref_base, weave_internal, ) from weave import weave_types as types from weave.legacy import ( + engine_trace, registry_mem, box, compile, diff --git a/weave/legacy/io_service.py b/weave/legacy/io_service.py index a5e0c37f554..4ef61c6218c 100644 --- a/weave/legacy/io_service.py +++ b/weave/legacy/io_service.py @@ -23,12 +23,12 @@ import aioprocessing from weave import ( - engine_trace, errors, filesystem, weave_http, ) from weave.legacy import ( + engine_trace, server_error_handling, artifact_wandb, async_queue, diff --git a/weave/legacy/language_features/tagging/tag_store.py b/weave/legacy/language_features/tagging/tag_store.py index 9a497cb5b4c..85a4c98de9f 100644 --- a/weave/legacy/language_features/tagging/tag_store.py +++ b/weave/legacy/language_features/tagging/tag_store.py @@ -22,9 +22,9 @@ from collections import defaultdict from contextlib import contextmanager -from weave import engine_trace, errors +from weave import errors from weave import weave_types as types -from weave.legacy import box +from weave.legacy import engine_trace, box statsd = engine_trace.statsd() # type: ignore diff --git a/weave/legacy/memo.py b/weave/legacy/memo.py index d18b3f2efd7..4928bb9b062 100644 --- a/weave/legacy/memo.py +++ b/weave/legacy/memo.py @@ -2,7 +2,7 @@ import contextvars import typing -from .. import engine_trace +from weave.legacy import engine_trace statsd = engine_trace.statsd() # type: ignore diff --git a/weave/legacy/op_def.py b/weave/legacy/op_def.py index 78d697e4e3c..0dba3ed28e3 100644 --- a/weave/legacy/op_def.py +++ b/weave/legacy/op_def.py @@ -7,13 +7,13 @@ from typing import Iterator, Sequence from weave import ( - engine_trace, errors, weave_internal, ) from weave import weave_types as types from weave.legacy import ( context_state, + engine_trace, graph, language_autocall, memo, diff --git a/weave/legacy/ops_arrow/list_join.py b/weave/legacy/ops_arrow/list_join.py index 270ca1c0a1d..6c7c34f7f1d 100644 --- a/weave/legacy/ops_arrow/list_join.py +++ b/weave/legacy/ops_arrow/list_join.py @@ -4,7 +4,7 @@ import pyarrow as pa import pyarrow.compute as pc -from weave import engine_trace +from weave.legacy import engine_trace from weave import weave_types as types from weave.query_api import op from weave.legacy import graph diff --git a/weave/legacy/ops_domain/run_history/history_op_common.py b/weave/legacy/ops_domain/run_history/history_op_common.py index 72d945d39f0..00821ae4d6f 100644 --- a/weave/legacy/ops_domain/run_history/history_op_common.py +++ b/weave/legacy/ops_domain/run_history/history_op_common.py @@ -5,13 +5,13 @@ from pyarrow import parquet as pq from weave import ( - engine_trace, errors, util, ) from weave import weave_types as types from weave.query_api import use from weave.legacy import ( + engine_trace, registry_mem, _dict_utils, artifact_base, diff --git a/weave/legacy/ops_domain/run_history/run_history_v1_legacy_ops.py b/weave/legacy/ops_domain/run_history/run_history_v1_legacy_ops.py index 14615652067..0b8e97f8787 100644 --- a/weave/legacy/ops_domain/run_history/run_history_v1_legacy_ops.py +++ b/weave/legacy/ops_domain/run_history/run_history_v1_legacy_ops.py @@ -2,7 +2,7 @@ import pyarrow as pa -from weave import engine_trace +from weave.legacy import engine_trace from weave import weave_types as types from weave.query_api import op from weave.legacy import gql_json_cache diff --git a/weave/legacy/ops_domain/run_history/run_history_v2_parquet_media.py b/weave/legacy/ops_domain/run_history/run_history_v2_parquet_media.py index f7f4c21db26..24afe60b728 100644 --- a/weave/legacy/ops_domain/run_history/run_history_v2_parquet_media.py +++ b/weave/legacy/ops_domain/run_history/run_history_v2_parquet_media.py @@ -6,7 +6,7 @@ import pyarrow as pa -from weave import engine_trace +from weave.legacy import engine_trace from weave import weave_types as types from weave.query_api import op, use from weave.legacy import artifact_mem, gql_json_cache diff --git a/weave/legacy/ops_domain/run_history/run_history_v3_parquet_stream_optimized.py b/weave/legacy/ops_domain/run_history/run_history_v3_parquet_stream_optimized.py index 3a767cce087..0412c764fdd 100644 --- a/weave/legacy/ops_domain/run_history/run_history_v3_parquet_stream_optimized.py +++ b/weave/legacy/ops_domain/run_history/run_history_v3_parquet_stream_optimized.py @@ -5,10 +5,11 @@ import pyarrow as pa -from weave import engine_trace, errors +from weave import errors from weave import weave_types as types from weave.query_api import op from weave.legacy import ( + engine_trace, artifact_base, artifact_fs, artifact_mem, diff --git a/weave/legacy/ops_domain/run_ops.py b/weave/legacy/ops_domain/run_ops.py index a9cd9df908c..1699c8d4e0b 100644 --- a/weave/legacy/ops_domain/run_ops.py +++ b/weave/legacy/ops_domain/run_ops.py @@ -37,7 +37,7 @@ import json import typing -from weave import engine_trace +from weave.legacy import engine_trace from weave import weave_types as types from weave.query_api import op from weave.legacy import compile_table diff --git a/weave/legacy/ops_domain/table.py b/weave/legacy/ops_domain/table.py index 651929a7f2e..d7c6b72e146 100644 --- a/weave/legacy/ops_domain/table.py +++ b/weave/legacy/ops_domain/table.py @@ -5,7 +5,7 @@ import logging import typing -from weave import engine_trace, errors, util, weave_internal +from weave import errors, util, weave_internal from weave import weave_types as types from weave.query_api import op, weave_class from weave.legacy import ( @@ -14,6 +14,7 @@ io_service, ops_arrow, wandb_util, + engine_trace, ) from weave.legacy import timestamp as weave_timestamp from weave.legacy.ops_domain import trace_tree, wbmedia diff --git a/weave/legacy/ops_domain/wbgqlquery_op.py b/weave/legacy/ops_domain/wbgqlquery_op.py index 1b78af45ae3..24dc73e3f8b 100644 --- a/weave/legacy/ops_domain/wbgqlquery_op.py +++ b/weave/legacy/ops_domain/wbgqlquery_op.py @@ -1,10 +1,10 @@ import logging import typing -from weave import engine_trace, environment, errors +from weave import environment, errors from weave import weave_types as types from weave.query_api import op -from weave.legacy import mappers_gql, partial_object +from weave.legacy import engine_trace, mappers_gql, partial_object from weave.legacy.language_features.tagging import tagged_value_type from weave.legacy.ops_domain import wb_domain_types as wdt from weave.legacy.wandb_client_api import wandb_gql_query diff --git a/weave/legacy/ops_domain/wbmedia.py b/weave/legacy/ops_domain/wbmedia.py index 20cccd808bd..34674731e97 100644 --- a/weave/legacy/ops_domain/wbmedia.py +++ b/weave/legacy/ops_domain/wbmedia.py @@ -5,8 +5,8 @@ import typing from weave import query_api as weave -from weave import engine_trace, errors, types -from weave.legacy import artifact_fs, file_base +from weave import errors, types +from weave.legacy import engine_trace, artifact_fs, file_base from weave.legacy.language_features.tagging.tag_store import isolated_tagging_context from weave.legacy.ops_primitives import html, markdown diff --git a/weave/legacy/wandb_api.py b/weave/legacy/wandb_api.py index cc26c0f7e66..b551538a6ae 100644 --- a/weave/legacy/wandb_api.py +++ b/weave/legacy/wandb_api.py @@ -14,9 +14,9 @@ from gql.transport.requests import RequestsHTTPTransport from requests.auth import HTTPBasicAuth -from weave import engine_trace, errors +from weave import errors from weave import environment as weave_env -from weave.legacy import wandb_client_api +from weave.legacy import wandb_client_api, engine_trace # Importing at the top-level namespace so other files can import from here. from weave.legacy.context_state import WandbApiContext, _wandb_api_context diff --git a/weave/legacy/wandb_file_manager.py b/weave/legacy/wandb_file_manager.py index d3868c29e8f..21c29a064af 100644 --- a/weave/legacy/wandb_file_manager.py +++ b/weave/legacy/wandb_file_manager.py @@ -12,9 +12,9 @@ from requests.auth import HTTPBasicAuth from wandb.sdk.lib import hashutil -from weave import engine_trace, errors, filesystem, weave_http +from weave import errors, filesystem, weave_http from weave import environment as weave_env -from weave.legacy import artifact_wandb, cache, wandb_api +from weave.legacy import artifact_wandb, cache, wandb_api, engine_trace tracer = engine_trace.tracer() # type: ignore diff --git a/weave/legacy/wandb_interface/wandb_artifact_pusher.py b/weave/legacy/wandb_interface/wandb_artifact_pusher.py index 8befef2c4e6..fa80c077c16 100644 --- a/weave/legacy/wandb_interface/wandb_artifact_pusher.py +++ b/weave/legacy/wandb_interface/wandb_artifact_pusher.py @@ -5,7 +5,7 @@ from wandb import Artifact from wandb.apis.public import api as wb_public -from weave import engine_trace +from weave.legacy import engine_trace from weave.legacy import wandb_client_api from weave.legacy.wandb_interface.wandb_lite_run import InMemoryLazyLiteRun diff --git a/weave/legacy/wandb_interface/wandb_lite_run.py b/weave/legacy/wandb_interface/wandb_lite_run.py index 4f565c527c2..95fa4c8e3f0 100644 --- a/weave/legacy/wandb_interface/wandb_lite_run.py +++ b/weave/legacy/wandb_interface/wandb_lite_run.py @@ -15,8 +15,8 @@ from wandb.sdk.internal.sender import _manifest_json_from_proto from wandb.sdk.lib import runid -from weave import engine_trace, errors -from weave.legacy import wandb_client_api +from weave import errors +from weave.legacy import engine_trace, wandb_client_api logger = logging.getLogger(__name__) diff --git a/weave/server.py b/weave/server.py index d620ef5167e..0792737639a 100644 --- a/weave/server.py +++ b/weave/server.py @@ -15,6 +15,7 @@ from weave.legacy import ( cache, context, + engine_trace, execute, gql_json_cache, graph, @@ -25,7 +26,7 @@ from weave.legacy.language_features.tagging import tag_store from weave.legacy.language_features.tagging.tag_store import isolated_tagging_context -from . import engine_trace, logs, storage, util, weave_types +from . import logs, storage, util, weave_types # A function to monkeypatch the request post method # def patch_request_post(): diff --git a/weave/test_scripts/wandb_artifact_perf.py b/weave/test_scripts/wandb_artifact_perf.py index 7452832dff3..0ba0b420b0a 100644 --- a/weave/test_scripts/wandb_artifact_perf.py +++ b/weave/test_scripts/wandb_artifact_perf.py @@ -16,11 +16,12 @@ from weave.legacy import ( artifact_wandb, async_map, + engine_trace, wandb_api, wandb_file_manager, ) -from .. import engine_trace, filesystem, weave_http +from .. import filesystem, weave_http tracer = engine_trace.tracer() # type: ignore diff --git a/weave/weave_http.py b/weave/weave_http.py index 8d35a00b5c0..075693b8edf 100644 --- a/weave/weave_http.py +++ b/weave/weave_http.py @@ -13,9 +13,9 @@ import requests.auth import yarl -from weave.legacy import server_error_handling +from weave.legacy import engine_trace, server_error_handling -from . import engine_trace, filesystem +from . import filesystem logging.getLogger("aiohttp.access").setLevel(logging.WARNING) logging.getLogger("aiohttp.client").setLevel(logging.WARNING) diff --git a/weave/weave_server.py b/weave/weave_server.py index 4ec00b16358..b9992991850 100644 --- a/weave/weave_server.py +++ b/weave/weave_server.py @@ -26,7 +26,6 @@ from werkzeug.exceptions import HTTPException from weave import ( - engine_trace, environment, errors, filesystem, @@ -37,6 +36,7 @@ ) from weave.legacy import ( context_state, + engine_trace, graph, registry_mem, value_or_error,