diff --git a/fastdeploy/metrics/metrics.py b/fastdeploy/metrics/metrics.py index 0798d89af..ca8b6b391 100644 --- a/fastdeploy/metrics/metrics.py +++ b/fastdeploy/metrics/metrics.py @@ -19,6 +19,7 @@ metrics """ import os import shutil +import uuid from typing import Set from prometheus_client import ( @@ -35,24 +36,20 @@ from fastdeploy.metrics import build_1_2_5_buckets from fastdeploy.metrics.work_metrics import work_process_metrics -def cleanup_prometheus_files(is_main): +def cleanup_prometheus_files(is_main: bool, instance_id: str = None): """ Cleans and recreates the Prometheus multiprocess directory. - - Depending on whether it's the main process or a worker, this function removes the corresponding - Prometheus multiprocess directory (/tmp/prom_main or /tmp/prom_worker) and recreates it as an empty directory. - - Args: - is_main (bool): Indicates whether the current process is the main process. - - Returns: - str: The path to the newly created Prometheus multiprocess directory. """ - PROM_DIR = "/tmp/prom_main" if is_main else "/tmp/prom_worker" - if os.path.exists(PROM_DIR): - shutil.rmtree(PROM_DIR) - os.makedirs(PROM_DIR, exist_ok=True) - return PROM_DIR + base_dir = "/tmp/prom_main" if is_main else "/tmp/prom_worker" + if instance_id is None: + instance_id = str(uuid.uuid4()) + prom_dir = f"{base_dir}_{instance_id}" + + if os.path.exists(prom_dir): + shutil.rmtree(prom_dir, ignore_errors=True) + os.makedirs(prom_dir, exist_ok=True) + + return prom_dir class SimpleCollector(Collector): diff --git a/tests/metrics/test_metrics.py b/tests/metrics/test_metrics.py index f49e4b4f4..23c84c8f7 100644 --- a/tests/metrics/test_metrics.py +++ b/tests/metrics/test_metrics.py @@ -14,12 +14,13 @@ # limitations under the License. """ +import os import unittest from unittest.mock import patch from prometheus_client import Gauge -from fastdeploy.metrics.metrics import get_filtered_metrics +from fastdeploy.metrics.metrics import cleanup_prometheus_files, get_filtered_metrics class TestGetFilteredMetrics(unittest.TestCase): @@ -68,5 +69,30 @@ class TestGetFilteredMetrics(unittest.TestCase): self.assertIn("custom_metric_total", result) +class TestCleanupPrometheusFiles(unittest.TestCase): + def test_cleanup_prometheus_files(self): + prom_dir = cleanup_prometheus_files(is_main=True) + self.assertTrue(os.path.exists(prom_dir)) + self.assertTrue(os.path.isdir(prom_dir)) + + prom_dir2 = cleanup_prometheus_files(is_main=True) + self.assertNotEqual(prom_dir, prom_dir2) + self.assertTrue(os.path.exists(prom_dir2)) + self.assertTrue(os.path.isdir(prom_dir2)) + + prom_dir = cleanup_prometheus_files(is_main=True, instance_id="001") + self.assertTrue(os.path.exists(prom_dir)) + self.assertTrue(os.path.isdir(prom_dir)) + test_file = os.path.join(prom_dir, "test.txt") + with open(test_file, "w") as f: + f.write("hello") + self.assertTrue(os.path.exists(test_file)) + prom_dir2 = cleanup_prometheus_files(is_main=True, instance_id="001") + self.assertEqual(prom_dir, prom_dir2) + self.assertTrue(os.path.exists(prom_dir2)) + self.assertTrue(os.path.isdir(prom_dir2)) + self.assertFalse(os.path.exists(test_file)) + + if __name__ == "__main__": unittest.main()