atomgit 1.0.5__tar.gz → 1.0.7__tar.gz

{atomgit-1.0.5 → atomgit-1.0.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: atomgit
-Version: 1.0.5
+Version: 1.0.7
 Summary: AtomGit模型文件上传下载CLI工具
 Home-page: https://atomgit.com/gitcode-ai/atomgit_cli
 Author: AtomGit CLI Team

{atomgit-1.0.5 → atomgit-1.0.7}/__init__.py

@@ -8,7 +8,7 @@ AtomGit CLI - 基于Transformers和Hugging Face Hub的模型文件上传下载
 支持模型和数据集的上传、下载等操作。
 """
 
-__version__ = '1.0.5'
+__version__ = '1.0.7'
 __author__ = 'AtomGit CLI Team'
 __description__ = 'AtomGit模型文件上传下载CLI工具'
 

{atomgit-1.0.5 → atomgit-1.0.7}/api.py

@@ -16,6 +16,12 @@ os.environ["HF_HOME"] = cache_dir
 
 
 from huggingface_hub import hf_hub_download, upload_folder, create_repo, snapshot_download, constants as hf_constants
+try:
+    from .rate_limiter import install_hf_rate_limiter
+except ImportError:
+    from rate_limiter import install_hf_rate_limiter
+
+install_hf_rate_limiter()
 
 try:
     from .config import config
@@ -221,7 +227,8 @@ class HuggingFaceAPI:
                     repo_id=normalized_repo_id,
                     local_dir=str(local_path),
                     force_download=force_download,  # 根据用户选择决定是否强制下载
-                    token=credentials['token'] if credentials and 'token' in credentials else None
+                    token=credentials['token'] if credentials and 'token' in credentials else None,
+                    max_workers=8
                 )
                 print(f"✅ 仓库下载成功")
                 return True
@@ -247,11 +254,14 @@ class HuggingFaceAPI:
             local_path.mkdir(parents=True, exist_ok=True)
             
             # 首先尝试不使用token下载（适用于公开仓库）
+            credentials = config.get_credentials()
             try:
                 hf_hub_download(
                     repo_id=normalized_repo_id,
                     filename=filename,
-                    local_dir=str(local_path)
+                    local_dir=str(local_path),
+                    token=credentials['token'] if credentials else None,
+                    force_download=force_download
                 )
                 print(f"✅ 文件下载成功")
                 return True
@@ -264,7 +274,6 @@ class HuggingFaceAPI:
                     print("检测到认证问题，尝试使用token下载...")
                     
                     # 如果公开下载失败，尝试使用token下载
-                    credentials = config.get_credentials()
                     if credentials:
                         try:
                             hf_hub_download(
@@ -307,4 +316,4 @@ class HuggingFaceAPI:
 
 
 # 全局API实例
-api = HuggingFaceAPI() 
+api = HuggingFaceAPI()

{atomgit-1.0.5 → atomgit-1.0.7}/atomgit.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: atomgit
-Version: 1.0.5
+Version: 1.0.7
 Summary: AtomGit模型文件上传下载CLI工具
 Home-page: https://atomgit.com/gitcode-ai/atomgit_cli
 Author: AtomGit CLI Team

{atomgit-1.0.5 → atomgit-1.0.7}/atomgit.egg-info/SOURCES.txt

@@ -6,6 +6,7 @@ api.py
 atomgit_hub.py
 cli.py
 config.py
+rate_limiter.py
 requirements.txt
 setup.py
 test.py
@@ -16,6 +17,7 @@ utils.py
 ./atomgit_hub.py
 ./cli.py
 ./config.py
+./rate_limiter.py
 ./test.py
 ./utils.py
 atomgit.egg-info/PKG-INFO
@@ -23,4 +25,5 @@ atomgit.egg-info/SOURCES.txt
 atomgit.egg-info/dependency_links.txt
 atomgit.egg-info/entry_points.txt
 atomgit.egg-info/requires.txt
-atomgit.egg-info/top_level.txt
+atomgit.egg-info/top_level.txt
+tests/test_rate_limiter.py

{atomgit-1.0.5 → atomgit-1.0.7}/atomgit_hub.py

@@ -23,6 +23,15 @@ os.environ["HF_HOME"] = cache_dir
 
 from huggingface_hub import snapshot_download as hf_snapshot_download
 from huggingface_hub import hf_hub_download, upload_folder as hf_upload_folder, create_repo
+try:
+    from .rate_limiter import install_hf_rate_limiter
+except ImportError:
+    try:
+        from rate_limiter import install_hf_rate_limiter
+    except ImportError:
+        from atomgit.rate_limiter import install_hf_rate_limiter
+
+install_hf_rate_limiter()
 
 
 try:
@@ -331,8 +340,10 @@ def upload_folder(
         # 如果要上传到根目录，直接使用源文件夹
         upload_path = str(folder_path)
     else:
-        # 如果要上传到特定路径，需要重新组织目录结构
-        with tempfile.TemporaryDirectory() as temp_dir:
+        # 如果要上传到特定路径，需要重新组织目录结构。
+        # 注意：不能提前退出 with TemporaryDirectory()，否则上传时临时目录已被删除。
+        temp_dir = tempfile.mkdtemp(prefix="atomgit_upload_")
+        try:
             temp_path = Path(temp_dir)
             
             # 创建目标路径
@@ -342,7 +353,11 @@ def upload_folder(
             # 复制整个目录树
             shutil.copytree(folder_path, target_path, dirs_exist_ok=True)
             
-            upload_path = str(temp_path)
+        except Exception:
+            shutil.rmtree(temp_dir, ignore_errors=True)
+            raise
+
+        upload_path = str(temp_path)
     # 使用 Monkey Patch 方式临时修改 huggingface_hub 的默认超时配置
     from huggingface_hub import constants as hf_constants
     
@@ -372,6 +387,10 @@ def upload_folder(
             raise Exception(f"仓库不存在：{repo_id}")
         else:
             raise Exception(f"上传失败：{error_msg}")
+    finally:
+        # 清理临时目录（仅重新组织目录结构时创建）
+        if path_in_repo not in ("./", ".", ""):
+            shutil.rmtree(temp_dir, ignore_errors=True)
 
 
 def create_repository(

{atomgit-1.0.5 → atomgit-1.0.7}/cli.py

@@ -39,7 +39,7 @@ except ImportError:
 
 
 @click.group()
-@click.version_option(version='1.0.5')
+@click.version_option(version='1.0.7')
 def cli():
     """AtomGit CLI - 基于Transformers和Hugging Face Hub的AtomGit平台模型文件上传下载工具"""
     pass

atomgit-1.0.7/rate_limiter.py

@@ -0,0 +1,107 @@
+"""Shared client-side rate limiting for AtomGit Hub requests."""
+
+from __future__ import annotations
+
+import json
+import os
+import time
+import threading
+from pathlib import Path
+from typing import Optional
+from urllib.parse import urlparse
+
+
+class SharedRateLimiter:
+    """A cross-process leaky bucket using a small lock-protected state file."""
+
+    _thread_lock = threading.Lock()
+
+    def __init__(self, rate_per_minute: float = 45.0, burst: int = 2,
+                 state_file: Optional[Path] = None) -> None:
+        self.rate = max(1.0, rate_per_minute) / 60.0
+        self.burst = max(1, burst)
+        self.state_file = state_file or (
+            Path(os.path.expanduser("~/.cache/atomgit")) / "rate_limit.json"
+        )
+        self.state_file.parent.mkdir(parents=True, exist_ok=True)
+
+    def acquire(self) -> None:
+        """Reserve one request slot, sharing reservations across processes."""
+        while True:
+            wait_for = 0.0
+            # flock is process-scoped on some Unix systems; the thread lock is
+            # therefore also required when a single CLI uses a worker pool.
+            with self._thread_lock, self.state_file.open("a+", encoding="utf-8") as handle:
+                try:
+                    import fcntl
+                    fcntl.flock(handle.fileno(), fcntl.LOCK_EX)
+                except ImportError:
+                    pass
+                handle.seek(0)
+                try:
+                    state = json.load(handle)
+                except (json.JSONDecodeError, EOFError):
+                    state = {"next_slot": 0.0}
+
+                now = time.monotonic()
+                next_slot = float(state.get("next_slot", 0.0))
+                # A reservation may be at most ``burst`` intervals ahead.
+                reservation = max(next_slot, now - self.burst / self.rate)
+                reservation += 1.0 / self.rate
+                state["next_slot"] = reservation
+                handle.seek(0)
+                handle.truncate()
+                json.dump(state, handle)
+                handle.flush()
+                wait_for = max(0.0, reservation - now)
+                try:
+                    fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
+                except (NameError, UnboundLocalError):
+                    pass
+            if wait_for <= 0:
+                return
+            time.sleep(min(wait_for, 1.0))
+
+
+def install_hf_rate_limiter() -> bool:
+    """Install the limiter into huggingface_hub v1.x when available.
+
+    The function is intentionally best-effort so the CLI remains compatible
+    with older huggingface_hub releases.
+    """
+    try:
+        import httpx
+        from huggingface_hub import set_client_factory
+    except (ImportError, AttributeError):
+        return False
+
+    # AtomGit's file metadata endpoint is limited to 50 requests/minute.
+    # Keep a safety margin and leave ordinary file transfer requests untouched.
+    limiter = SharedRateLimiter(rate_per_minute=45.0, burst=2)
+
+    def is_file_api_request(request) -> bool:
+        path = urlparse(str(request.url)).path.rstrip("/")
+        return "/api/v1/file/" in f"{path}/"
+
+    class LimitedClient(httpx.Client):
+        def send(self, request, *args, **kwargs):
+            attempts = 0
+            while True:
+                limited = is_file_api_request(request)
+                if limited:
+                    limiter.acquire()
+                response = super().send(request, *args, **kwargs)
+                if not limited or response.status_code != 429 or attempts >= 5:
+                    return response
+                retry_after = response.headers.get("Retry-After")
+                try:
+                    delay = max(1.0, float(retry_after)) if retry_after else 2 ** attempts
+                except ValueError:
+                    delay = 2 ** attempts
+                response.close()
+                time.sleep(min(delay, 60.0))
+                attempts += 1
+
+    set_client_factory(lambda: LimitedClient())
+    return True
+    _thread_lock = threading.Lock()

{atomgit-1.0.5 → atomgit-1.0.7}/setup.py

@@ -31,7 +31,7 @@ def read_requirements():
 
 setup(
     name='atomgit',
-    version='1.0.5',
+    version='1.0.7',
     author='AtomGit CLI Team',
     author_email='sa@atomgit.com',
     description='AtomGit模型文件上传下载CLI工具',

{atomgit-1.0.5 → atomgit-1.0.7}/test.py

@@ -130,7 +130,7 @@ def run_all_tests():
     # 2. 测试upload
     # test_upload_folder(repo_id=model, folder_path=model_folder_path)
     # test_upload_folder(repo_id=dataset, folder_path=dataset_folder_path)
-    test_upload_folder(repo_id="yanlp/dataset-t1", folder_path="/Users/yanlp/csdn/IdeaProjects/gitcode/gitcode-hf-registry/dataset-t11")
+    test_upload_folder(repo_id="yanlp/dataset-t2", folder_path="/Users/yanlp/csdn/IdeaProjects/gitcode/gitcode-hf-registry/dataset-t11")
 
     # 3. 测试下载
     # test_snapshot_download(repo_id=model, local_dir="./test_downloads/model/snapshot_full")

atomgit-1.0.7/tests/test_rate_limiter.py

@@ -0,0 +1,72 @@
+import json
+import subprocess
+import sys
+import tempfile
+import threading
+import time
+import unittest
+from pathlib import Path
+
+from rate_limiter import SharedRateLimiter, install_hf_rate_limiter
+
+
+class RateLimiterTests(unittest.TestCase):
+    def test_concurrent_access_keeps_valid_shared_state(self):
+        with tempfile.TemporaryDirectory() as directory:
+            state_file = Path(directory) / "rate.json"
+            limiter = SharedRateLimiter(60000, 4, state_file)
+            errors = []
+
+            def worker():
+                try:
+                    for _ in range(10):
+                        limiter.acquire()
+                except Exception as error:  # pragma: no cover - assertion below
+                    errors.append(error)
+
+            threads = [threading.Thread(target=worker) for _ in range(4)]
+            for thread in threads:
+                thread.start()
+            for thread in threads:
+                thread.join()
+
+            self.assertFalse(errors)
+            self.assertIsInstance(json.loads(state_file.read_text()), dict)
+
+    def test_multiple_processes_share_state_file(self):
+        with tempfile.TemporaryDirectory() as directory:
+            state_file = Path(directory) / "rate.json"
+            code = (
+                "from pathlib import Path; "
+                "from rate_limiter import SharedRateLimiter; "
+                f"SharedRateLimiter(60000, 4, Path({str(state_file)!r})).acquire()"
+            )
+            processes = [subprocess.run([sys.executable, "-c", code]) for _ in range(3)]
+            self.assertTrue(all(process.returncode == 0 for process in processes))
+            self.assertTrue(state_file.exists())
+
+    def test_huggingface_client_factory_is_limited(self):
+        import huggingface_hub
+
+        self.assertTrue(install_hf_rate_limiter())
+        self.assertEqual(type(huggingface_hub.get_session()).__name__, "LimitedClient")
+        huggingface_hub.close_session()
+
+    def test_file_api_limit_is_conservative(self):
+        limiter = SharedRateLimiter()
+        self.assertAlmostEqual(limiter.rate * 60, 45.0)
+        self.assertEqual(limiter.burst, 2)
+
+
+class InstalledCliTests(unittest.TestCase):
+    def test_installed_cli_help(self):
+        cli = Path(__file__).parents[1] / ".venv" / "bin" / "atomgit"
+        if not cli.exists():
+            self.skipTest(".venv/bin/atomgit is not installed")
+        result = subprocess.run([str(cli), "--help"], capture_output=True, text=True)
+        self.assertEqual(result.returncode, 0, result.stderr)
+        self.assertIn("download", result.stdout)
+
+
+if __name__ == "__main__":
+    unittest.main()