atomgit 1.0.4__tar.gz → 1.0.6__tar.gz

atomgit-1.0.6/MANIFEST.in

@@ -0,0 +1 @@
+include requirements.txt

{atomgit-1.0.4 → atomgit-1.0.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: atomgit
-Version: 1.0.4
+Version: 1.0.6
 Summary: AtomGit模型文件上传下载CLI工具
 Home-page: https://atomgit.com/gitcode-ai/atomgit_cli
 Author: AtomGit CLI Team
@@ -29,8 +29,8 @@ Requires-Dist: requests>=2.31.0
 Requires-Dist: tqdm>=4.66.1
 Requires-Dist: colorama>=0.4.6
 Requires-Dist: tabulate>=0.9.0
-Requires-Dist: huggingface-hub>=0.20.0
-Requires-Dist: datasets>=2.16.0
+Requires-Dist: huggingface-hub==1.1.7
+Requires-Dist: datasets==4.4.1
 Dynamic: author
 Dynamic: author-email
 Dynamic: classifier

{atomgit-1.0.4 → atomgit-1.0.6}/__init__.py

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

{atomgit-1.0.4 → atomgit-1.0.6}/api.py

@@ -7,6 +7,8 @@ import urllib.error
 
 # 设置Hugging Face Hub的API端点为AtomGit
 os.environ["HF_ENDPOINT"] = "https://hub.atomgit.com"
+# 禁用Xet协议，避免 xet-write-token 请求
+os.environ["HF_HUB_DISABLE_XET"] = "1"
 # 设置缓存目录
 cache_dir = os.path.expanduser("~/.cache/atomgit")
 os.makedirs(cache_dir, exist_ok=True)
@@ -14,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
@@ -151,12 +159,13 @@ class HuggingFaceAPI:
                 # 使用 Monkey Patch 方式临时修改 huggingface_hub 的默认超时配置
                 commit_message = message or "Upload folder using atomgit client"
                 hf_constants.DEFAULT_REQUEST_TIMEOUT = upload_timeout
-                upload_folder(
+                upload_kwargs = dict(
                     repo_id=repo_id,
                     folder_path=str(temp_dir),
                     token=credentials['token'],
-                    commit_message=commit_message
+                    commit_message=commit_message,
                 )
+                upload_folder(**upload_kwargs)
                 
                 return True
             finally:
@@ -185,12 +194,13 @@ class HuggingFaceAPI:
             # 使用 Monkey Patch 方式临时修改 huggingface_hub 的默认超时配置
             commit_message = message or "Upload folder using atomgit client"
             hf_constants.DEFAULT_REQUEST_TIMEOUT = upload_timeout
-            upload_folder(
+            upload_kwargs = dict(
                 repo_id=repo_id,
                 folder_path=str(dir_path),
                 token=credentials['token'],
-                commit_message=commit_message
+                commit_message=commit_message,
             )
+            upload_folder(**upload_kwargs)
             
             return True
             
@@ -217,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
@@ -243,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
@@ -260,7 +274,6 @@ class HuggingFaceAPI:
                     print("检测到认证问题，尝试使用token下载...")
                     
                     # 如果公开下载失败，尝试使用token下载
-                    credentials = config.get_credentials()
                     if credentials:
                         try:
                             hf_hub_download(
@@ -303,4 +316,4 @@ class HuggingFaceAPI:
 
 
 # 全局API实例
-api = HuggingFaceAPI() 
+api = HuggingFaceAPI()

{atomgit-1.0.4 → atomgit-1.0.6}/atomgit.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: atomgit
-Version: 1.0.4
+Version: 1.0.6
 Summary: AtomGit模型文件上传下载CLI工具
 Home-page: https://atomgit.com/gitcode-ai/atomgit_cli
 Author: AtomGit CLI Team
@@ -29,8 +29,8 @@ Requires-Dist: requests>=2.31.0
 Requires-Dist: tqdm>=4.66.1
 Requires-Dist: colorama>=0.4.6
 Requires-Dist: tabulate>=0.9.0
-Requires-Dist: huggingface-hub>=0.20.0
-Requires-Dist: datasets>=2.16.0
+Requires-Dist: huggingface-hub==1.1.7
+Requires-Dist: datasets==4.4.1
 Dynamic: author
 Dynamic: author-email
 Dynamic: classifier

{atomgit-1.0.4 → atomgit-1.0.6}/atomgit.egg-info/SOURCES.txt

@@ -1,3 +1,4 @@
+MANIFEST.in
 README.md
 __init__.py
 __main__.py
@@ -5,6 +6,8 @@ api.py
 atomgit_hub.py
 cli.py
 config.py
+rate_limiter.py
+requirements.txt
 setup.py
 test.py
 utils.py
@@ -14,6 +17,7 @@ utils.py
 ./atomgit_hub.py
 ./cli.py
 ./config.py
+./rate_limiter.py
 ./test.py
 ./utils.py
 atomgit.egg-info/PKG-INFO
@@ -21,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.4 → atomgit-1.0.6}/atomgit.egg-info/requires.txt

@@ -3,5 +3,5 @@ requests>=2.31.0
 tqdm>=4.66.1
 colorama>=0.4.6
 tabulate>=0.9.0
-huggingface-hub>=0.20.0
-datasets>=2.16.0
+huggingface-hub==1.1.7
+datasets==4.4.1

{atomgit-1.0.4 → atomgit-1.0.6}/atomgit_hub.py

@@ -14,6 +14,8 @@ from pathlib import Path
 
 # 设置Hugging Face Hub的API端点为AtomGit
 os.environ["HF_ENDPOINT"] = "https://hub.atomgit.com"
+# 禁用Xet协议，避免 xet-write-token 请求
+os.environ["HF_HUB_DISABLE_XET"] = "1"
 # 设置缓存目录
 cache_dir = os.path.expanduser("~/.cache/atomgit")
 os.makedirs(cache_dir, exist_ok=True)
@@ -21,6 +23,12 @@ 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:
+    from rate_limiter import install_hf_rate_limiter
+
+install_hf_rate_limiter()
 
 
 try:

{atomgit-1.0.4 → atomgit-1.0.6}/cli.py

@@ -9,6 +9,8 @@ from getpass import getpass
 
 # 设置Hugging Face Hub的API端点为AtomGit
 os.environ["HF_ENDPOINT"] = "https://hub.atomgit.com"
+# 禁用Xet协议，避免 xet-write-token 请求
+os.environ["HF_HUB_DISABLE_XET"] = "1"
 # 设置缓存目录
 cache_dir = os.path.expanduser("~/.cache/atomgit")
 os.makedirs(cache_dir, exist_ok=True)
@@ -37,7 +39,7 @@ except ImportError:
 
 
 @click.group()
-@click.version_option(version='1.0.4')
+@click.version_option(version='1.0.6')
 def cli():
     """AtomGit CLI - 基于Transformers和Hugging Face Hub的AtomGit平台模型文件上传下载工具"""
     pass

atomgit-1.0.6/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.6/requirements.txt

@@ -0,0 +1,7 @@
+click>=8.1.7
+requests>=2.31.0
+tqdm>=4.66.1
+colorama>=0.4.6
+tabulate>=0.9.0
+huggingface-hub==1.1.7
+datasets==4.4.1

{atomgit-1.0.4 → atomgit-1.0.6}/setup.py

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

{atomgit-1.0.4 → atomgit-1.0.6}/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.6/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()