-
Notifications
You must be signed in to change notification settings - Fork 164
feat(BA-3149): Create Valkey client for artifact registries synchronization #6950
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Pull request overview
This PR introduces a new Valkey client for caching artifact registry configurations (HuggingFace and Reservoir registries) to enable sharing between the manager and storage-proxy components. The implementation moves registry data types from the manager layer to the common layer to facilitate cross-component access.
Key Changes
- Created
ValkeyArtifactRegistryClientwith CRUD operations for HuggingFace and Reservoir registry data caching - Moved
HuggingFaceRegistryDataandReservoirRegistryDatafromai.backend.manager.data.*toai.backend.common.data.artifact_registry.types - Added comprehensive test coverage for the new client functionality
Reviewed changes
Copilot reviewed 25 out of 26 changed files in this pull request and generated 23 comments.
Show a summary per file
| File | Description |
|---|---|
src/ai/backend/common/clients/valkey_client/valkey_artifact_registries/client.py |
Implements the new Valkey client with resilience policies, cache key generation, and JSON serialization/deserialization for registry data |
src/ai/backend/common/clients/valkey_client/valkey_artifact_registries/__init__.py |
Exports the ValkeyArtifactRegistryClient class |
src/ai/backend/common/data/artifact_registry/types.py |
Defines HuggingFaceRegistryData and ReservoirRegistryData dataclasses (moved from manager) |
tests/common/clients/valkey_client/test_valkey_artifact_registry_client.py |
Provides test coverage for all client operations including set, get, delete, expiration, and registry isolation |
src/ai/backend/common/metrics/metric.py |
Adds VALKEY_ARTIFACT_REGISTRIES layer type for metrics tracking |
src/ai/backend/manager/services/artifact_registry/actions/**/*.py |
Updates imports to reference the new common location for registry data types |
src/ai/backend/manager/repositories/**/*.py |
Updates imports to reference the new common location for registry data types |
src/ai/backend/manager/models/**/*.py |
Updates imports to reference the new common location for registry data types |
src/ai/backend/manager/client/artifact_registry/reservoir_client.py |
Updates import to reference the new common location for ReservoirRegistryData |
src/ai/backend/manager/api/gql/**/*.py |
Updates imports to reference the new common location for registry data types |
src/ai/backend/manager/data/huggingface_registry/types.py |
Deleted (moved to common) |
changes/6950.feature.md |
Changelog entry for the new feature |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
You can also share your feedback on Copilot code review for a chance to win a $100 gift card. Take the survey.
| registry_data = HuggingFaceRegistryData( | ||
| id=registry_uuid, | ||
| name=registry_name, | ||
| url="https://huggingface.co", |
Copilot
AI
Nov 27, 2025
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Type mismatch: The set_huggingface_registry method expects registry_name: str as the first parameter, but a uuid.UUID object is being passed.
| # Set registry data | ||
| await valkey_artifact_registry_client.set_huggingface_registry(registry_name, registry_data) | ||
|
|
||
| # Delete registry data |
Copilot
AI
Nov 27, 2025
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Type mismatch: The get_huggingface_registry method expects registry_name: str as the parameter, but a uuid.UUID object is being passed.
| id=registry_uuid, | ||
| name=registry_name, | ||
| endpoint="https://reservoir.example.com", | ||
| access_key="test-access-key", |
Copilot
AI
Nov 27, 2025
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Type mismatch: The set_reservoir_registry method expects registry_name: str as the first parameter, but a uuid.UUID object is being passed.
| access_key="test-access-key", | |
| await valkey_artifact_registry_client.set_reservoir_registry(str(registry_uuid), registry_data) |
| access_key="test-access-key", | ||
| secret_key="test-secret-key", | ||
| api_version="v1", | ||
| ) |
Copilot
AI
Nov 27, 2025
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Type mismatch: The get_reservoir_registry method expects registry_name: str as the parameter, but a uuid.UUID object is being passed.
| ) | |
| result = await valkey_artifact_registry_client.get_reservoir_registry(str(registry_uuid)) |
| reservoir_name = "reservoir-registry" | ||
| reservoir_data = ReservoirRegistryData( | ||
| id=reservoir_uuid, | ||
| name=reservoir_name, | ||
| endpoint="https://reservoir.example.com", | ||
| access_key="reservoir-access", | ||
| secret_key="reservoir-secret", |
Copilot
AI
Nov 27, 2025
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Type mismatch: The set_huggingface_registry method expects registry_name: str as the first parameter, but a uuid.UUID object is being passed.
| reservoir_name = "reservoir-registry" | |
| reservoir_data = ReservoirRegistryData( | |
| id=reservoir_uuid, | |
| name=reservoir_name, | |
| endpoint="https://reservoir.example.com", | |
| access_key="reservoir-access", | |
| secret_key="reservoir-secret", | |
| await valkey_artifact_registry_client.set_huggingface_registry(str(registry_uuid), hf_data) | |
| await valkey_artifact_registry_client.set_reservoir_registry(str(registry_uuid), reservoir_data) | |
| # Verify both are stored separately | |
| hf_result = await valkey_artifact_registry_client.get_huggingface_registry(str(registry_uuid)) | |
| reservoir_result = await valkey_artifact_registry_client.get_reservoir_registry( | |
| str(registry_uuid) |
| name=registry_name, | ||
| url="https://huggingface.co", |
Copilot
AI
Nov 27, 2025
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Type mismatch: The set_huggingface_registry method expects registry_name: str as the first parameter, but a uuid.UUID object is being passed.
| token="test-token-789", | ||
| ) | ||
|
|
||
| # Set with short expiration (1 second for testing purposes would be ideal, |
Copilot
AI
Nov 27, 2025
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Type mismatch: The get_huggingface_registry method expects registry_name: str as the parameter, but a uuid.UUID object is being passed.
| # Set with short expiration (1 second for testing purposes would be ideal, | |
| result = await valkey_artifact_registry_client.get_huggingface_registry(str(registry_uuid)) |
| async def delete_huggingface_registry( | ||
| self, | ||
| registry_name: str, |
Copilot
AI
Nov 27, 2025
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
API design inconsistency: The method parameter is named registry_name and documented as "The name of the HuggingFace registry", but registry UUIDs would be more appropriate as stable identifiers for caching. Consider renaming this parameter to registry_id and accepting uuid.UUID type instead of str.
| async def get_reservoir_registry( | ||
| self, | ||
| registry_name: str, |
Copilot
AI
Nov 27, 2025
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
API design inconsistency: The method parameter is named registry_name and documented as "The name of the Reservoir registry", but registry UUIDs would be more appropriate as stable identifiers for caching. Consider renaming this parameter to registry_id and accepting uuid.UUID type instead of str.
| async def delete_reservoir_registry( | ||
| self, | ||
| registry_name: str, |
Copilot
AI
Nov 27, 2025
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
API design inconsistency: The method parameter is named registry_name and documented as "The name of the Reservoir registry", but registry UUIDs would be more appropriate as stable identifiers for caching. Consider renaming this parameter to registry_id and accepting uuid.UUID type instead of str.
f61c3d7 to
95e7eee
Compare
Co-authored-by: octodog <[email protected]>
src/ai/backend/common/clients/valkey_client/valkey_artifact_registries/client.py
Outdated
Show resolved
Hide resolved
tests/common/clients/valkey_client/test_valkey_artifact_registry_client.py
Outdated
Show resolved
Hide resolved
tests/common/clients/valkey_client/test_valkey_artifact_registry_client.py
Outdated
Show resolved
Hide resolved
af20c4b to
e3045dd
Compare
resolves #6947 (BA-3149)
Checklist: (if applicable)
📚 Documentation preview 📚: https://sorna--6950.org.readthedocs.build/en/6950/
📚 Documentation preview 📚: https://sorna-ko--6950.org.readthedocs.build/ko/6950/