chore: enhance cache and ObjectRef#296
Conversation
Signed-off-by: Klaus Ma <klausm@nvidia.com>
Summary of ChangesHello @k82cn, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly refactors the object caching mechanism by simplifying the Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the caching mechanism and the ObjectRef type. The changes simplify the ObjectRef to only represent remote objects and move the serialization (pickling) logic into the cache helper functions. This is a good improvement for encapsulation and clarity. My review includes suggestions to enhance error handling in the cache functions to make the SDK more robust by using specific FlameError exceptions. I also pointed out and suggested removing a potential issue with log suppression that goes against the repository's guidelines.
| if context._cache_endpoint is None: | ||
| raise Exception("Cache endpoint is not configured") | ||
|
|
||
| # Serialize the object using pickle | ||
| data = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL) | ||
|
|
||
| with suppress_dependency_logs(): | ||
| response = httpx.post(f"{context._cache_endpoint}/objects/{session_id}", data=data) | ||
| response.raise_for_status() | ||
|
|
||
| metadata = ObjectMetadata.model_validate(response.json()) | ||
| return ObjectRef(source=DataSource.REMOTE, url=metadata.endpoint, data=data, version=metadata.version) | ||
|
|
||
|
|
||
| def get_object(de: ObjectRef) -> "ObjectRef": | ||
| """Get an object from the cache.""" | ||
| if de.source != DataSource.REMOTE: | ||
| return de | ||
|
|
||
| return ObjectRef(url=metadata.endpoint, version=metadata.version) |
There was a problem hiding this comment.
This function can be improved in a couple of ways:
- Error Handling: Raise a specific
FlameErrorinstead of a genericException. Also, catch potential exceptions fromhttpxandpydanticand wrap them inFlameErrorfor consistent error handling. - Logging: The use of
suppress_dependency_logsgoes against the repository's guidelines, as it can interfere with the application's logging configuration. It's better to let the application developer control logging levels.
The suggestion below applies these improvements. You'll also need to import FlameError and FlameErrorCode from .types.
if context._cache_endpoint is None:
raise FlameError(FlameErrorCode.INVALID_CONFIG, "Cache endpoint is not configured")
# Serialize the object using pickle
data = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)
try:
response = httpx.post(f"{context._cache_endpoint}/objects/{session_id}", data=data)
response.raise_for_status()
metadata = ObjectMetadata.model_validate(response.json())
return ObjectRef(url=metadata.endpoint, version=metadata.version)
except httpx.HTTPError as e:
raise FlameError(FlameErrorCode.INTERNAL, f"Failed to put object in cache: {e}") from e
except Exception as e:
raise FlameError(FlameErrorCode.INTERNAL, f"Failed to process cache response: {e}") from eReferences
- In a library, avoid modifying the global logging configuration for dependencies. This can cause conflicts with the application's own logging setup. Instead, document how users can configure logging or provide an opt-in utility (like a context manager) for temporary log level adjustments.
| with suppress_dependency_logs(): | ||
| response = httpx.get(de.url) | ||
| response = httpx.get(ref.url) | ||
| response.raise_for_status() | ||
|
|
||
| obj = Object.model_validate(response.json()) | ||
|
|
||
| de.data = bytes(obj.data) | ||
| de.version = obj.version | ||
|
|
||
| return de | ||
|
|
||
|
|
||
| def update_object(de: ObjectRef) -> "ObjectRef": | ||
| """Update an object in the cache.""" | ||
| if de.source != DataSource.REMOTE: | ||
| return de | ||
|
|
||
| obj = Object(version=de.version, data=list(de.data)) | ||
| data = bytes(obj.data) | ||
|
|
||
| # Deserialize the object using pickle | ||
| return pickle.loads(data) |
There was a problem hiding this comment.
This function can be improved with more robust error handling and by adhering to logging guidelines:
- Error Handling: Catch potential exceptions from
httpx,pydantic, andpickleand wrap them inFlameError. - Logging: The
suppress_dependency_logsshould be removed to comply with repository guidelines.
The suggestion below applies these improvements.
try:
response = httpx.get(ref.url)
response.raise_for_status()
obj = Object.model_validate(response.json())
data = bytes(obj.data)
# Deserialize the object using pickle
return pickle.loads(data)
except httpx.HTTPError as e:
raise FlameError(FlameErrorCode.INTERNAL, f"Failed to get object from cache: {e}") from e
except pickle.UnpicklingError as e:
raise FlameError(FlameErrorCode.INTERNAL, f"Failed to deserialize object from cache: {e}") from e
except Exception as e:
raise FlameError(FlameErrorCode.INTERNAL, f"Failed to process cache response: {e}") from eReferences
- In a library, avoid modifying the global logging configuration for dependencies. This can cause conflicts with the application's own logging setup. Instead, document how users can configure logging or provide an opt-in utility (like a context manager) for temporary log level adjustments.
| with suppress_dependency_logs(): | ||
| response = httpx.put(de.url, data=data) | ||
| response = httpx.put(ref.url, data=data) | ||
| response.raise_for_status() | ||
|
|
||
| metadata = ObjectMetadata.model_validate(response.json()) | ||
|
|
||
| de.version = metadata.version | ||
|
|
||
| return de | ||
| return ObjectRef(url=ref.url, version=metadata.version) |
There was a problem hiding this comment.
This function can be improved with more robust error handling and by adhering to logging guidelines:
- Error Handling: Catch potential exceptions from
httpxandpydanticand wrap them inFlameError. - Logging: The
suppress_dependency_logsshould be removed to comply with repository guidelines.
The suggestion below applies these improvements.
try:
response = httpx.put(ref.url, data=data)
response.raise_for_status()
metadata = ObjectMetadata.model_validate(response.json())
return ObjectRef(url=ref.url, version=metadata.version)
except httpx.HTTPError as e:
raise FlameError(FlameErrorCode.INTERNAL, f"Failed to update object in cache: {e}") from e
except Exception as e:
raise FlameError(FlameErrorCode.INTERNAL, f"Failed to process cache response: {e}") from eReferences
- In a library, avoid modifying the global logging configuration for dependencies. This can cause conflicts with the application's own logging setup. Instead, document how users can configure logging or provide an opt-in utility (like a context manager) for temporary log level adjustments.
1 participant
No description provided.