add posthog (#410)

* add posthog * add posthog 2 * add posthog fixes * add posthog fixes * update: posthog * add posthog: dynamic import * bump: version
michaelfeil · Oct 10, 2024 · 6ce0eb5 · 6ce0eb5
1 parent 98c5895
commit 6ce0eb5
Show file tree

Hide file tree

Showing 13 changed files with 377 additions and 12 deletions.
diff --git a/docs/assets/openapi.json b/docs/assets/openapi.json
diff --git a/docs/docs/cli_v2.md b/docs/docs/cli_v2.md
@@ -75,8 +75,8 @@ $ infinity_emb v2 --help
 │                                                                                       `INFINITY_LENGTHS_VIA_TOKENIZ… │
 │                                                                                       [default:                      │
 │                                                                                       lengths-via-tokenize]          │
-│ --dtype                                                [float32|float16|int8|fp8|aut  dtype for the model weights.   │
-│                                                        o]                             [env var: `INFINITY_DTYPE`]    │
+│ --dtype                                                [float32|float16|bfloat16|int  dtype for the model weights.   │
+│                                                        8|fp8|auto]                    [env var: `INFINITY_DTYPE`]    │
 │                                                                                       [default: auto]                │
 │ --embedding-dtype                                      [float32|int8|uint8|binary|ub  dtype post-forward pass. If != │
 │                                                        inary]                         `float32`, using Post-Forward  │

diff --git a/docs/docs/telemetry.md b/docs/docs/telemetry.md
@@ -0,0 +1,16 @@
+# Telemetry
+
+All CLI arguments are currently logged, as well as system info such as os.
+The feedback will be used to optimzed common models / OS support.
+
+## Disable Telemetry
+You can disable tracking like the following:
+
+```
+# set 
+export DO_NOT_TRACK="1"
+# infinity specific setting
+export INFINITY_ANONYMOUS_USAGE_STATS="0"
+```
+
+or by CLI argument.
diff --git a/libs/client_infinity/infinity_client/pyproject.toml b/libs/client_infinity/infinity_client/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "infinity_client"
-version = "0.0.62"
+version = "0.0.63"
 description = "A client library for accessing ♾️ Infinity - Embedding Inference Server"
 authors = []
 readme = "README.md"

diff --git a/libs/infinity_emb/infinity_emb/__init__.py b/libs/infinity_emb/infinity_emb/__init__.py
@@ -28,7 +28,7 @@
 from infinity_emb.log_handler import logger  # noqa: E402
 from infinity_emb.sync_engine import SyncEngineArray  # noqa: E402
 
-__version__ = importlib.metadata.version("infinity_emb")
+__version__: str = importlib.metadata.version("infinity_emb")
 
 __all__ = [
     "__version__",

diff --git a/libs/infinity_emb/infinity_emb/_optional_imports.py b/libs/infinity_emb/infinity_emb/_optional_imports.py
@@ -67,6 +67,7 @@ def _raise_error(self) -> None:
 CHECK_TRANSFORMERS = OptionalImports("transformers", "torch")
 CHECK_TORCH = OptionalImports("torch.nn", "torch")
 # CHECK_REQUESTS = OptionalImports("requests", "server")
+CHECK_POSTHOG = OptionalImports("posthog", "server")
 CHECK_AIOHTTP = OptionalImports("aiohttp", "server")
 CHECK_PIL = OptionalImports("PIL", "vision")
 CHECK_SOUNDFILE = OptionalImports("soundfile", "audio")

diff --git a/libs/infinity_emb/infinity_emb/env.py b/libs/infinity_emb/infinity_emb/env.py
@@ -68,7 +68,7 @@ def _optional_infinity_var_multiple(
 
     @staticmethod
     def _to_bool(value: str) -> bool:
-        return value.lower() in {"true", "1"}
+        return value.lower() in {"true", "1", "yes", "y"}
 
     @staticmethod
     def _to_bool_multiple(value: list[str]) -> list[bool]:
@@ -153,6 +153,17 @@ def calibration_dataset_url(self):
             default="https://raw.githubusercontent.com/michaelfeil/infinity/2da1f32d610b8edbe4ce58d0c44fc27c963abca6/docs/assets/multilingual_calibration.utf8",
         )
 
+    @cached_property
+    def anonymous_usage_stats(self):
+        tracking_allowed = self._to_bool(
+            self._optional_infinity_var(
+                "anonymous_usage_stats",
+                default="true",
+            )
+        )
+        tracking_allowed_2 = not self._to_bool(os.getenv("DO_NOT_TRACK", "0"))
+        return tracking_allowed and tracking_allowed_2
+
     @cached_property
     def cache_dir(self) -> Path:
         """gets the cache directory for infinity_emb."""

diff --git a/libs/infinity_emb/infinity_emb/infinity_server.py b/libs/infinity_emb/infinity_emb/infinity_server.py
@@ -6,6 +6,7 @@
 import signal
 import sys
 import time
+import uuid
 from contextlib import asynccontextmanager
 from typing import Any, Optional, Union
 
@@ -36,9 +37,11 @@
     ImageCorruption,
     InferenceEngine,
     Modality,
+    ModelCapabilites,
     ModelNotDeployedError,
     PoolingMethod,
 )
+from infinity_emb.telemetry import PostHog, StartupTelemetry
 
 
 def create_server(
@@ -61,10 +64,32 @@ def create_server(
     from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
     from prometheus_fastapi_instrumentator import Instrumentator
 
+    def send_telemetry_start(
+        engine_args_list: list[EngineArgs],
+        capabilities_list: list[set[ModelCapabilites]],
+    ):
+        session_id = uuid.uuid4().hex
+        for arg, capabilities in zip(engine_args_list, capabilities_list):
+            PostHog.capture(
+                StartupTelemetry(
+                    engine_args=arg,
+                    num_engines=len(engine_args_list),
+                    capabilities=capabilities,
+                    session_id=session_id,
+                )
+            )
+
     @asynccontextmanager
     async def lifespan(app: FastAPI):
         instrumentator.expose(app)  # type: ignore
         app.engine_array = AsyncEngineArray.from_args(engine_args_list)  # type: ignore
+        asyncio.create_task(
+            asyncio.to_thread(
+                send_telemetry_start,
+                engine_args_list,
+                [e.capabilities for e in app.engine_array],  # type: ignore
+            )
+        )
         # start in a threadpool
         await app.engine_array.astart()  # type: ignore
 
@@ -87,6 +112,7 @@ async def kill_later(seconds: int):
                 " -> exit ."
             )
             asyncio.create_task(kill_later(3))
+
         yield
         await app.engine_array.astop()  # type: ignore
         # shutdown!

diff --git a/libs/infinity_emb/infinity_emb/primitives.py b/libs/infinity_emb/infinity_emb/primitives.py
@@ -67,7 +67,7 @@ class ClassifyReturnType(TypedDict):
 UnionReturnType = Union[EmbeddingReturnType, ReRankReturnType, ClassifyReturnType]
 
 
-class EnumType(enum.Enum):
+class EnumType(str, enum.Enum):
     @classmethod
     @lru_cache
     def names_enum(cls) -> enum.Enum:
@@ -125,6 +125,7 @@ def resolve(self) -> Optional[str]:
 class Dtype(EnumType):
     float32: str = "float32"
     float16: str = "float16"
+    bfloat16: str = "bfloat16"
     int8: str = "int8"
     fp8: str = "fp8"
     auto: str = "auto"