wip

weave/trace/concurrent/futures.py

@@ -34,6 +34,7 @@
 import logging
 from concurrent.futures import Future, wait
 from contextvars import ContextVar
+import random
 from threading import Lock
 from typing import Any, Callable, TypeVar
 
@@ -68,6 +69,7 @@ def __init__(
         max_workers: int | None = None,
         thread_name_prefix: str = THREAD_NAME_PREFIX,
     ):
+        self._id = random.randint(0, 1000000)
         self._max_workers = max_workers
         self._executor: ContextAwareThreadPoolExecutor | None = None
         if max_workers != 0:

weave/trace/weave_client.py

@@ -10,6 +10,7 @@
 from collections.abc import Iterator, Sequence
 from concurrent.futures import Future
 from functools import lru_cache
+import threading
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -670,6 +671,9 @@ def to_dict(self) -> CallDict:
         )
 
 
+images = set()
+
+
 def make_client_call(
     entity: str, project: str, server_call: CallSchema, server: TraceServerInterface
 ) -> WeaveObject:
@@ -752,7 +756,7 @@ def __repr__(self) -> str:
 
 class WeaveClient:
     server: TraceServerInterface
-    future_executor: FutureExecutor
+    # future_executor: FutureExecutor
 
     """
     A client for interacting with the Weave trace server.
@@ -775,7 +779,8 @@ def __init__(
         self.project = project
         self.server = server
         self._anonymous_ops: dict[str, Op] = {}
-        self.future_executor = FutureExecutor(max_workers=client_parallelism())
+        self._future_executor_local = threading.local()
+        self._future_executor = FutureExecutor(max_workers=client_parallelism())
         self.ensure_project_exists = ensure_project_exists
 
         if ensure_project_exists:
@@ -787,6 +792,16 @@ def __init__(
         if isinstance(self.server, RemoteHTTPTraceServer):
             self._server_is_flushable = self.server.should_batch
 
+    @property
+    def future_executor(self) -> FutureExecutor:
+        # return self._future_executor
+
+        if not hasattr(self._future_executor_local, "future_executor"):
+            self._future_executor_local.future_executor = FutureExecutor(
+                max_workers=client_parallelism()
+            )
+        return self._future_executor_local.future_executor
+
     ################ High Level Convenience Methods ################
 
     @trace_sentry.global_trace_sentry.watch()
@@ -1873,7 +1888,23 @@ def _flush(self) -> None:
             self.server.call_processor.wait_until_all_processed()  # type: ignore
 
     def _send_file_create(self, req: FileCreateReq) -> Future[FileCreateRes]:
-        return self.future_executor.defer(self.server.file_create, req)
+        import base64
+        from threading import current_thread
+        from datetime import datetime
+
+        thread = current_thread()
+        name = base64.b64encode(req.content).decode("utf-8")
+        print(
+            "send_file_create",
+            datetime.now().strftime("%M:%S.%f")[:-3],
+            name[-30:],
+            thread.name,
+            "seen={}".format(name in images),
+            "len={}".format(len(images)),
+        )
+        images.add(name)
+        out = self.future_executor.defer(self.server.file_create, req)
+        return out
 
 
 def safe_current_wb_run_id() -> str | None:

weave/trace_server/requests.py

@@ -1,11 +1,12 @@
 """Helpers for printing HTTP requests and responses."""
 
+from contextlib import contextmanager
 import datetime
 import json
 import os
 import threading
 from time import time
-from typing import Any, Optional, Union
+from typing import Any, Iterator, Optional, Union
 
 from requests import HTTPError as HTTPError
 from requests import PreparedRequest, Response, Session
@@ -14,6 +15,8 @@
 from rich.syntax import Syntax
 from rich.text import Text
 
+from weave.trace.settings import client_parallelism
+
 console = Console()
 
 # See https://rich.readthedocs.io/en/stable/appendix/colors.html
@@ -147,16 +150,18 @@ def send(self, request: PreparedRequest, **kwargs: Any) -> Response:  # type: ig
         return response
 
 
-session = Session()
-if os.environ.get("WEAVE_DEBUG_HTTP") == "1":
-    adapter = LoggingHTTPAdapter()
-    session.mount("http://", adapter)
-    session.mount("https://", adapter)
+thread_local = threading.local()
+
+
+def get_session() -> Session:
+    if not hasattr(thread_local, "session"):
+        thread_local.session = Session()
+    return thread_local.session
 
 
 def get(url: str, params: Optional[dict[str, str]] = None, **kwargs: Any) -> Response:
     """Send a GET request with optional logging."""
-    return session.get(url, params=params, **kwargs)
+    return get_session().get(url, params=params, **kwargs)
 
 
 def post(
@@ -166,4 +171,4 @@ def post(
     **kwargs: Any,
 ) -> Response:
     """Send a POST request with optional logging."""
-    return session.post(url, data=data, json=json, **kwargs)
+    return get_session().post(url, data=data, json=json, **kwargs)