feat: add adaptive GPU scheduling
This commit is contained in:
@@ -11,6 +11,7 @@ if str(PACKAGE_DIR) in sys.path:
|
||||
sys.path.insert(0, str(PACKAGE_DIR))
|
||||
|
||||
from hf_discovery import HuggingFaceDiscovery # noqa: E402
|
||||
from http_json import HttpJsonError # noqa: E402
|
||||
|
||||
|
||||
class RecordingHttpClient:
|
||||
@@ -22,6 +23,30 @@ class RecordingHttpClient:
|
||||
return {"success": True, "data": {"models": []}}
|
||||
|
||||
|
||||
class RateLimitedPageClient:
|
||||
def __init__(self) -> None:
|
||||
self.page_calls: list[int] = []
|
||||
self.page_two_attempts = 0
|
||||
|
||||
def request_json(self, _method: str, _path: str, *, query: dict) -> dict:
|
||||
page = int(query["page_number"])
|
||||
self.page_calls.append(page)
|
||||
if page == 2:
|
||||
self.page_two_attempts += 1
|
||||
if self.page_two_attempts == 1:
|
||||
raise HttpJsonError("rate limited", status_code=429)
|
||||
start = (page - 1) * 50
|
||||
items = [
|
||||
{
|
||||
"id": f"owner/model-{index}",
|
||||
"downloads": 100,
|
||||
"last_modified": "2026-01-01T00:00:00Z",
|
||||
}
|
||||
for index in range(start, start + 50)
|
||||
]
|
||||
return {"success": True, "data": {"models": items}}
|
||||
|
||||
|
||||
class ModelScopeDiscoveryTests(unittest.TestCase):
|
||||
def test_openapi_page_size_never_exceeds_platform_limit(self) -> None:
|
||||
http_client = RecordingHttpClient()
|
||||
@@ -36,6 +61,29 @@ class ModelScopeDiscoveryTests(unittest.TestCase):
|
||||
self.assertEqual([], models)
|
||||
self.assertEqual(50, http_client.queries[0]["page_size"])
|
||||
|
||||
def test_rate_limited_page_keeps_previous_pages_and_resumes_from_cache(self) -> None:
|
||||
http_client = RateLimitedPageClient()
|
||||
discovery = HuggingFaceDiscovery(
|
||||
http_client=http_client, # type: ignore[arg-type]
|
||||
page_interval_seconds=0,
|
||||
page_cache_ttl_seconds=60,
|
||||
)
|
||||
|
||||
partial = discovery.list_recent_models(
|
||||
pipeline_tags=["text-generation"],
|
||||
limit=100,
|
||||
min_downloads=0,
|
||||
)
|
||||
resumed = discovery.list_recent_models(
|
||||
pipeline_tags=["text-generation"],
|
||||
limit=100,
|
||||
min_downloads=0,
|
||||
)
|
||||
|
||||
self.assertEqual(50, len(partial))
|
||||
self.assertEqual(100, len(resumed))
|
||||
self.assertEqual([1, 2, 2], http_client.page_calls)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user