From 2f12ead687006d087dd229e354fc8fbe0b9cf574 Mon Sep 17 00:00:00 2001
From: Philip Potter <phil@replicate.com>
Date: Thu, 16 Jan 2025 12:25:13 +0000
Subject: [PATCH] Avoid warnings

---
 python/cog/server/scope.py  | 6 +++++-
 python/cog/server/worker.py | 4 ++--
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/python/cog/server/scope.py b/python/cog/server/scope.py
index 30229d21be..d326cf4b46 100644
--- a/python/cog/server/scope.py
+++ b/python/cog/server/scope.py
@@ -23,6 +23,10 @@ def current_scope() -> Scope:
         category=ExperimentalFeatureWarning,
         stacklevel=1,
     )
+    return _get_current_scope()
+
+
+def _get_current_scope() -> Scope:
     s = _current_scope.get()
     if s is None:
         raise RuntimeError("No scope available")
@@ -40,7 +44,7 @@ def scope(sc: Scope) -> Generator[None, None, None]:
 
 @contextmanager
 def evolve_scope(**kwargs: Any) -> Generator[None, None, None]:
-    new_scope = evolve(current_scope(), **kwargs)
+    new_scope = evolve(_get_current_scope(), **kwargs)
     s = _current_scope.set(new_scope)
     try:
         yield
diff --git a/python/cog/server/worker.py b/python/cog/server/worker.py
index f283f789c6..40f8ff0b5a 100644
--- a/python/cog/server/worker.py
+++ b/python/cog/server/worker.py
@@ -57,7 +57,7 @@
     InvalidStateException,
 )
 from .helpers import SimpleStreamRedirector, StreamRedirector
-from .scope import Scope, current_scope, evolve_scope, scope
+from .scope import Scope, _get_current_scope, evolve_scope, scope
 
 if PYDANTIC_V2:
     from .helpers import unwrap_pydantic_serialization_iterators
@@ -480,7 +480,7 @@ def record_metric(self, name: str, value: Union[float, int]) -> None:
     @property
     def _current_tag(self) -> Optional[str]:
         if self._has_async_predictor:
-            return current_scope()._tag
+            return _get_current_scope()._tag
         return self._sync_tag
 
     def _load_predictor(self) -> Optional[BasePredictor]: