Avoid warnings

replicate · Jan 16, 2025 · 2f12ead · 2f12ead
1 parent 1007849
commit 2f12ead
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git 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
@@ -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]: