diff --git a/python/cog/__init__.py b/python/cog/__init__.py index dd3bcedf4f..d6ad24d9af 100644 --- a/python/cog/__init__.py +++ b/python/cog/__init__.py @@ -4,7 +4,7 @@ from .base_predictor import BasePredictor from .mimetypes_ext import install_mime_extensions -from .server.scope import current_scope +from .server.scope import current_scope, emit_metric from .types import ( ConcatenateIterator, ExperimentalFeatureWarning, @@ -24,10 +24,11 @@ __all__ = [ "__version__", + "current_scope", + "emit_metric", "BaseModel", "BasePredictor", "ConcatenateIterator", - "current_scope", "ExperimentalFeatureWarning", "File", "Input", diff --git a/python/cog/server/scope.py b/python/cog/server/scope.py index 131a434684..7487c63b7e 100644 --- a/python/cog/server/scope.py +++ b/python/cog/server/scope.py @@ -37,3 +37,21 @@ def scope(sc: Scope) -> Generator[None, None, None]: yield finally: _current_scope.reset(s) + + +def emit_metric(name: str, value: Union[float, int]) -> None: + """ + DEPRECATED: This function will be removed in a future version of cog. + + Records a metric event from the model. Intended to be called from + within the `predict` function. + + This allows older models using an experimental `cog.emit_metric` function + to run using newer releases without requiring code changes. + """ + warnings.warn( + "emit_metric is deprecated and will be removed in a future version. Use `current_scope().record_metric()` instead", + category=DeprecationWarning, + stacklevel=1, + ) + current_scope().record_metric(name, value) diff --git a/python/tests/server/fixtures/emit_metric.py b/python/tests/server/fixtures/emit_metric.py new file mode 100644 index 0000000000..c4ffe868f0 --- /dev/null +++ b/python/tests/server/fixtures/emit_metric.py @@ -0,0 +1,13 @@ +from cog import emit_metric + + +class Predictor: + def setup(self): + print("did setup") + + def predict(self, *, name: str) -> str: + print(f"hello, {name}") + + emit_metric("foo", 123) + + return f"hello, {name}" diff --git a/python/tests/server/fixtures/emit_metric_async.py b/python/tests/server/fixtures/emit_metric_async.py new file mode 100644 index 0000000000..805b03ef77 --- /dev/null +++ b/python/tests/server/fixtures/emit_metric_async.py @@ -0,0 +1,13 @@ +from cog import emit_metric + + +class Predictor: + def setup(self): + print("did setup") + + async def predict(self, *, name: str) -> str: + print(f"hello, {name}") + + emit_metric("foo", 123) + + return f"hello, {name}" diff --git a/python/tests/server/test_worker.py b/python/tests/server/test_worker.py index 456cae84e2..cb6a469430 100644 --- a/python/tests/server/test_worker.py +++ b/python/tests/server/test_worker.py @@ -82,6 +82,20 @@ "foo": 123, }, ), + ( + WorkerConfig("emit_metric"), + {"name": ST_NAMES}, + { + "foo": 123, + }, + ), + ( + WorkerConfig("emit_metric_async", is_async=True), + {"name": ST_NAMES}, + { + "foo": 123, + }, + ), ] OUTPUT_FIXTURES = [