Merge pull request #778 from roboflow/feature/store-hostname-in-usage

Add hostname with optional DEDICATED_DEPLOYMENT_ID to usage payload

inference/core/env.py

@@ -446,3 +446,5 @@
     os.getenv("ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE", "True")
 )
 WORKFLOW_BLOCKS_WRITE_DIRECTORY = os.getenv("WORKFLOW_BLOCKS_WRITE_DIRECTORY")
+
+DEDICATED_DEPLOYMENT_ID = os.getenv("DEDICATED_DEPLOYMENT_ID")

inference/core/version.py

@@ -1,4 +1,4 @@
-__version__ = "0.25.0"
+__version__ = "0.25.1"
 
 
 if __name__ == "__main__":

inference/usage_tracking/collector.py

@@ -15,7 +15,7 @@
 
 from typing_extensions import ParamSpec
 
-from inference.core.env import API_KEY, LAMBDA, REDIS_HOST
+from inference.core.env import API_KEY, DEDICATED_DEPLOYMENT_ID, LAMBDA, REDIS_HOST
 from inference.core.logger import logger
 from inference.core.version import __version__ as inference_version
 from inference.core.workflows.execution_engine.v1.compiler.entities import (
@@ -136,6 +136,7 @@ def empty_usage_dict(exec_session_id: str) -> APIKeyUsage:
                     "timestamp_start": None,
                     "timestamp_stop": None,
                     "exec_session_id": exec_session_id,
+                    "hostname": "",
                     "ip_address_hash": "",
                     "processed_frames": 0,
                     "fps": 0,
@@ -232,13 +233,26 @@ def record_resource_details(
     @staticmethod
     def system_info(
         ip_address: Optional[str] = None,
+        hostname: Optional[str] = None,
+        dedicated_deployment_id: Optional[str] = None,
     ) -> SystemDetails:
+        if not dedicated_deployment_id:
+            dedicated_deployment_id = DEDICATED_DEPLOYMENT_ID
+        if not hostname:
+            try:
+                hostname = socket.gethostname()
+            except Exception as exc:
+                logger.warning("Could not obtain hostname, %s", exc)
+                hostname = ""
+        if dedicated_deployment_id:
+            hostname = f"{dedicated_deployment_id}:{hostname}"
         if ip_address:
             ip_address_hash_hex = sha256_hash(ip_address)
         else:
             try:
                 ip_address: str = socket.gethostbyname(socket.gethostname())
-            except:
+            except Exception as exc:
+                logger.warning("Could not obtain IP address, %s", exc)
                 s = None
                 try:
                     s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
@@ -253,6 +267,7 @@ def system_info(
             ip_address_hash_hex = sha256_hash(ip_address)
 
         return {
+            "hostname": hostname,
             "ip_address_hash": ip_address_hash_hex,
             "is_gpu_available": False,  # TODO
         }
@@ -311,6 +326,7 @@ def _update_usage_payload(
         with self._system_info_lock:
             ip_address_hash = self._system_info["ip_address_hash"]
             is_gpu_available = self._system_info["is_gpu_available"]
+            hostname = self._system_info["hostname"]
         with UsageCollector._lock:
             source_usage = self._usage[api_key_hash][f"{category}:{resource_id}"]
             if not source_usage["timestamp_start"]:
@@ -325,6 +341,7 @@ def _update_usage_payload(
             source_usage["resource_id"] = resource_id
             source_usage["resource_details"] = json.dumps(resource_details)
             source_usage["api_key_hash"] = api_key_hash
+            source_usage["hostname"] = hostname
             source_usage["ip_address_hash"] = ip_address_hash
             source_usage["is_gpu_available"] = is_gpu_available
             logger.debug("Updated usage: %s", source_usage)

tests/inference/unit_tests/usage_tracking/test_collector.py

@@ -33,6 +33,7 @@ def test_create_empty_usage_dict():
                     "timestamp_start": None,
                     "timestamp_stop": None,
                     "exec_session_id": "exec_session_id",
+                    "hostname": "",
                     "ip_address_hash": "",
                     "processed_frames": 0,
                     "fps": 0,
@@ -806,12 +807,27 @@ def test_zip_usage_payloads_with_different_exec_session_ids():
     ]
 
 
-def test_system_info():
+def test_system_info_with_dedicated_deployment_id():
     # given
-    system_info = UsageCollector.system_info(ip_address="w.x.y.z")
+    system_info = UsageCollector.system_info(ip_address="w.x.y.z", hostname="hostname01", dedicated_deployment_id="deployment01")
 
     # then
     expected_system_info = {
+        "hostname": f"deployment01:hostname01",
+        "ip_address_hash": hashlib.sha256("w.x.y.z".encode()).hexdigest()[:5],
+        "is_gpu_available": False,
+    }
+    for k, v in expected_system_info.items():
+        assert system_info[k] == v
+
+
+def test_system_info_with_no_dedicated_deployment_id():
+    # given
+    system_info = UsageCollector.system_info(ip_address="w.x.y.z", hostname="hostname01")
+
+    # then
+    expected_system_info = {
+        "hostname": f"hostname01",
         "ip_address_hash": hashlib.sha256("w.x.y.z".encode()).hexdigest()[:5],
         "is_gpu_available": False,
     }