diff --git a/py/routes/update_routes.py b/py/routes/update_routes.py index 79012582..14dc6c95 100644 --- a/py/routes/update_routes.py +++ b/py/routes/update_routes.py @@ -225,7 +225,7 @@ class UpdateRoutes: logger.debug("Could not close downloaded-version history database", exc_info=True) # Skip settings.json, civitai, model cache and runtime cache folders - UpdateRoutes._clean_plugin_folder(plugin_root, skip_files=['settings.json', 'civitai', 'model_cache', 'cache', 'wildcards', 'backups']) + UpdateRoutes._clean_plugin_folder(plugin_root, skip_files=['settings.json', 'civitai', 'model_cache', 'cache', 'wildcards', 'backups', 'stats']) # Extract ZIP to temp dir with tempfile.TemporaryDirectory() as tmp_dir: @@ -235,7 +235,7 @@ class UpdateRoutes: extracted_root = next(os.scandir(tmp_dir)).path # Copy files, skipping user data that should be preserved - skip_items = {'settings.json', 'civitai', 'wildcards', 'backups'} + skip_items = {'settings.json', 'civitai', 'wildcards', 'backups', 'stats'} for item in os.listdir(extracted_root): if item in skip_items: continue @@ -252,7 +252,7 @@ class UpdateRoutes: # for ComfyUI Manager to work properly tracking_info_file = os.path.join(plugin_root, '.tracking') tracking_files = [] - skip_tracked = {'civitai', 'wildcards', 'backups'} + skip_tracked = {'civitai', 'wildcards', 'backups', 'stats'} for root, dirs, files in os.walk(extracted_root): # Skip user data directories and their contents rel_root = os.path.relpath(root, extracted_root) diff --git a/py/services/backup_service.py b/py/services/backup_service.py index 6864ff8f..c852f7ca 100644 --- a/py/services/backup_service.py +++ b/py/services/backup_service.py @@ -141,6 +141,16 @@ class BackupService: ) ) + stats_path = os.path.join(get_settings_dir(create=True), "stats", "lora_manager_stats.json") + if os.path.exists(stats_path): + targets.append( + ( + "usage_stats", + "stats/lora_manager_stats.json", + stats_path, + ) + ) + return targets @staticmethod @@ -348,6 +358,8 @@ class BackupService: if kind == "model_update": filename = os.path.basename(archive_member) return str(Path(get_cache_file_path(CacheType.MODEL_UPDATE, create_dir=True)).parent / filename) + if kind == "usage_stats": + return os.path.join(get_settings_dir(create=True), "stats", "lora_manager_stats.json") return None async def create_auto_snapshot_if_due(self) -> Optional[dict[str, Any]]: diff --git a/py/utils/usage_stats.py b/py/utils/usage_stats.py index a603668c..066b8314 100644 --- a/py/utils/usage_stats.py +++ b/py/utils/usage_stats.py @@ -10,6 +10,7 @@ from typing import Dict, Set from ..config import config from ..services.service_registry import ServiceRegistry +from ..utils.settings_paths import get_settings_dir # Check if running in standalone mode standalone_mode = os.environ.get("LORA_MANAGER_STANDALONE", "0") == "1" or os.environ.get("HF_HUB_DISABLE_TELEMETRY", "0") == "0" @@ -83,6 +84,7 @@ class UsageStats: # Load existing stats if available self._stats_file_path = self._get_stats_file_path() + self._migrate_from_old_location() self._load_stats() # Save interval in seconds @@ -95,14 +97,38 @@ class UsageStats: logger.debug("Usage statistics tracker initialized") def _get_stats_file_path(self) -> str: - """Get the path to the stats JSON file""" + """Get the path to the stats JSON file in the settings directory.""" + settings_dir = get_settings_dir(create=True) + return os.path.join(settings_dir, "stats", self.STATS_FILENAME) + + @staticmethod + def _get_old_stats_file_path() -> str: + """Get the legacy stats file path in the first lora root directory.""" if not config.loras_roots or len(config.loras_roots) == 0: - # If no lora roots are available, we can't save stats - # This will be handled by the caller - raise RuntimeError("No LoRA root directories configured. Cannot initialize usage statistics.") - - # Use the first lora root - return os.path.join(config.loras_roots[0], self.STATS_FILENAME) + return "" + return os.path.join(config.loras_roots[0], UsageStats.STATS_FILENAME) + + def _migrate_from_old_location(self) -> None: + """Migrate stats file from old location (first lora root) to new location (settings_dir/stats/).""" + new_path = self._stats_file_path + if os.path.exists(new_path): + return + + old_path = self._get_old_stats_file_path() + if not old_path or not os.path.exists(old_path): + return + + try: + os.makedirs(os.path.dirname(new_path), exist_ok=True) + shutil.copy2(old_path, new_path) + logger.info("Migrated usage stats from %s to %s", old_path, new_path) + try: + os.remove(old_path) + logger.info("Cleaned up old stats file: %s", old_path) + except Exception as e: + logger.warning("Failed to remove old stats file %s: %s", old_path, e) + except Exception as e: + logger.error("Failed to migrate usage stats from %s to %s: %s", old_path, new_path, e) def _backup_old_stats(self): """Backup the old stats file before conversion""" diff --git a/tests/utils/test_usage_stats.py b/tests/utils/test_usage_stats.py index 5794a265..9d61b8c9 100644 --- a/tests/utils/test_usage_stats.py +++ b/tests/utils/test_usage_stats.py @@ -27,9 +27,12 @@ async def _finalize_usage_stats(tasks): def _prepare_usage_stats(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, *, sleep_override=None): UsageStats._instance = None - stats_root = tmp_path / "loras" - stats_root.mkdir(parents=True, exist_ok=True) - monkeypatch.setattr(usage_stats_module.config, "loras_roots", [str(stats_root)]) + settings_dir = tmp_path / "settings" + settings_dir.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr(usage_stats_module, "get_settings_dir", lambda create=True: str(settings_dir)) + loras_root = tmp_path / "loras" + loras_root.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr(usage_stats_module.config, "loras_roots", [str(loras_root)]) created_tasks = [] real_create_task = usage_stats_module.asyncio.create_task @@ -45,7 +48,7 @@ def _prepare_usage_stats(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, *, sle monkeypatch.setattr(usage_stats_module.asyncio, "sleep", sleep_override) stats = UsageStats() - return stats, created_tasks, stats_root + return stats, created_tasks, settings_dir, loras_root async def test_usage_stats_converts_legacy_format(tmp_path, monkeypatch): @@ -57,12 +60,15 @@ async def test_usage_stats_converts_legacy_format(tmp_path, monkeypatch): } UsageStats._instance = None - stats_root = tmp_path / "loras" - stats_root.mkdir(parents=True, exist_ok=True) - monkeypatch.setattr(usage_stats_module.config, "loras_roots", [str(stats_root)]) + settings_dir = tmp_path / "settings" + settings_dir.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr(usage_stats_module, "get_settings_dir", lambda create=True: str(settings_dir)) + loras_root = tmp_path / "loras" + loras_root.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr(usage_stats_module.config, "loras_roots", [str(loras_root)]) - stats_path = stats_root / UsageStats.STATS_FILENAME - stats_path.write_text(json.dumps(legacy_stats), encoding="utf-8") + old_stats_path = loras_root / UsageStats.STATS_FILENAME + old_stats_path.write_text(json.dumps(legacy_stats), encoding="utf-8") created_tasks = [] real_create_task = usage_stats_module.asyncio.create_task @@ -83,20 +89,23 @@ async def test_usage_stats_converts_legacy_format(tmp_path, monkeypatch): assert converted["checkpoints"]["hash1"] == {"total": 3, "history": {today: 3}} assert converted["loras"]["hash2"] == {"total": 5, "history": {today: 5}} - backup_path = stats_path.with_suffix(stats_path.suffix + UsageStats.BACKUP_SUFFIX) + new_stats_path = settings_dir / "stats" / UsageStats.STATS_FILENAME + assert new_stats_path.exists() + + backup_path = new_stats_path.with_suffix(new_stats_path.suffix + UsageStats.BACKUP_SUFFIX) assert backup_path.exists() await _finalize_usage_stats(created_tasks) async def test_usage_stats_save_stats_persists_file(tmp_path, monkeypatch): - stats, tasks, stats_root = _prepare_usage_stats(tmp_path, monkeypatch) + stats, tasks, settings_dir, _ = _prepare_usage_stats(tmp_path, monkeypatch) stats.stats["total_executions"] = 4 saved = await stats.save_stats(force=True) assert saved is True - stats_path = stats_root / UsageStats.STATS_FILENAME + stats_path = settings_dir / "stats" / UsageStats.STATS_FILENAME persisted = json.loads(stats_path.read_text(encoding="utf-8")) assert persisted["total_executions"] == 4 assert persisted["last_save_time"] == stats.stats["last_save_time"] @@ -110,7 +119,7 @@ async def test_usage_stats_background_processor_handles_pending_prompts(tmp_path async def fast_sleep(_seconds): await real_sleep(0.01) - stats, tasks, _ = _prepare_usage_stats(tmp_path, monkeypatch, sleep_override=fast_sleep) + stats, tasks, _, _ = _prepare_usage_stats(tmp_path, monkeypatch, sleep_override=fast_sleep) metadata_calls = [] # Use string literals directly to avoid dependency on conditional imports @@ -155,7 +164,7 @@ async def test_usage_stats_background_processor_handles_pending_prompts(tmp_path async def test_usage_stats_calculates_pending_checkpoint_hash_on_demand(tmp_path, monkeypatch): - stats, tasks, _ = _prepare_usage_stats(tmp_path, monkeypatch) + stats, tasks, _, _ = _prepare_usage_stats(tmp_path, monkeypatch) metadata_payload = { "models": { @@ -195,7 +204,7 @@ async def test_usage_stats_calculates_pending_checkpoint_hash_on_demand(tmp_path async def test_usage_stats_skips_failed_checkpoint_hash_retry(tmp_path, monkeypatch): - stats, tasks, _ = _prepare_usage_stats(tmp_path, monkeypatch) + stats, tasks, _, _ = _prepare_usage_stats(tmp_path, monkeypatch) metadata_payload = { "models": { @@ -234,7 +243,7 @@ async def test_usage_stats_skips_failed_checkpoint_hash_retry(tmp_path, monkeypa async def test_usage_stats_resolves_manually_copied_checkpoint_from_disk(tmp_path, monkeypatch): - stats, tasks, _ = _prepare_usage_stats(tmp_path, monkeypatch) + stats, tasks, _, _ = _prepare_usage_stats(tmp_path, monkeypatch) checkpoints_root = tmp_path / "checkpoints" checkpoints_root.mkdir() @@ -273,7 +282,7 @@ async def test_usage_stats_resolves_manually_copied_checkpoint_from_disk(tmp_pat async def test_usage_stats_skips_name_fallback_for_missing_lora_hash(tmp_path, monkeypatch): - stats, tasks, _ = _prepare_usage_stats(tmp_path, monkeypatch) + stats, tasks, _, _ = _prepare_usage_stats(tmp_path, monkeypatch) metadata_payload = { "models": {}, @@ -294,3 +303,79 @@ async def test_usage_stats_skips_name_fallback_for_missing_lora_hash(tmp_path, m assert not any(key.startswith("name:") for key in stats.stats["loras"]) await _finalize_usage_stats(tasks) + + +async def test_usage_stats_migrates_from_old_location(tmp_path, monkeypatch): + settings_dir = tmp_path / "settings" + settings_dir.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr(usage_stats_module, "get_settings_dir", lambda create=True: str(settings_dir)) + loras_root = tmp_path / "loras" + loras_root.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr(usage_stats_module.config, "loras_roots", [str(loras_root)]) + + old_data = { + "checkpoints": {}, + "loras": {"lora-hash": {"total": 3, "history": {"2025-01-01": 3}}}, + "embeddings": {}, + "total_executions": 3, + "last_save_time": 100.0, + } + old_path = loras_root / UsageStats.STATS_FILENAME + old_path.write_text(json.dumps(old_data), encoding="utf-8") + + created_tasks = [] + real_create_task = usage_stats_module.asyncio.create_task + + def _track_task(coro): + task = real_create_task(coro) + created_tasks.append(task) + return task + + monkeypatch.setattr(usage_stats_module.asyncio, "create_task", _track_task) + + stats = UsageStats() + + new_path = settings_dir / "stats" / UsageStats.STATS_FILENAME + assert new_path.exists(), "Stats file should be migrated to new location" + assert not old_path.exists(), "Old stats file should be removed after migration" + assert stats.stats["total_executions"] == 3 + assert stats.stats["loras"]["lora-hash"]["total"] == 3 + + await _finalize_usage_stats(created_tasks) + + +async def test_usage_stats_uses_new_location_directly(tmp_path, monkeypatch): + settings_dir = tmp_path / "settings" + settings_dir.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr(usage_stats_module, "get_settings_dir", lambda create=True: str(settings_dir)) + loras_root = tmp_path / "loras" + loras_root.mkdir(parents=True, exist_ok=True) + monkeypatch.setattr(usage_stats_module.config, "loras_roots", [str(loras_root)]) + + new_data = { + "checkpoints": {}, + "loras": {}, + "embeddings": {}, + "total_executions": 7, + "last_save_time": 200.0, + } + new_path = settings_dir / "stats" / UsageStats.STATS_FILENAME + new_path.parent.mkdir(parents=True, exist_ok=True) + new_path.write_text(json.dumps(new_data), encoding="utf-8") + + created_tasks = [] + real_create_task = usage_stats_module.asyncio.create_task + + def _track_task(coro): + task = real_create_task(coro) + created_tasks.append(task) + return task + + monkeypatch.setattr(usage_stats_module.asyncio, "create_task", _track_task) + + stats = UsageStats() + + assert stats.stats["total_executions"] == 7 + assert not loras_root.joinpath(UsageStats.STATS_FILENAME).exists() + + await _finalize_usage_stats(created_tasks)