feat: add code samples for continuous tuning #13579
Conversation
|
Here is the summary of changes. You are about to add 1 region tag.
This comment is generated by snippet-bot.
|
product-auto-label
bot
added
the
samples
Summary of ChangesHello @yishan-pu, 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 delivers a new code sample for continuous tuning of generative AI models, showcasing how to leverage the Google GenAI SDK to create and monitor tuning jobs. It enables users to specify a base model and an optional checkpoint ID for fine-tuning, ensuring the sample is functional and reliable through an updated SDK dependency and a dedicated unit test. 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 introduces a new code sample for continuous tuning in Google GenAI. The changes include a new script continuous_tuning_create.py with the sample logic, an update to the google-genai dependency, and a new test case. My review focuses on improving code quality and test correctness. I've suggested moving imports and constants to the module level for better structure, renaming a parameter for clarity, and significantly improving the new test case to correctly cover the polling logic in the sample code.
| @patch("google.genai.Client") | ||
| def test_continuous_tuning_create(mock_genai_client: MagicMock) -> None: | ||
| # Mock the API response | ||
| mock_tuning_job = types.TuningJob( | ||
| name="test-tuning-job", | ||
| experiment="test-experiment", | ||
| tuned_model=types.TunedModel( | ||
| model="test-model-2", | ||
| endpoint="test-endpoint" | ||
| ) | ||
| ) | ||
| mock_genai_client.return_value.tunings.tune.return_value = mock_tuning_job | ||
|
|
||
| response = continuous_tuning_create.create_continuous_tuning_job(tuned_model_name="test-model", checkpoint_id="1") | ||
|
|
||
| mock_genai_client.assert_called_once_with(http_options=types.HttpOptions(api_version="v1beta1")) | ||
| mock_genai_client.return_value.tunings.tune.assert_called_once() | ||
| assert response == "test-tuning-job" |
There was a problem hiding this comment.
This test doesn't fully cover the functionality of create_continuous_tuning_job. The polling while loop, which includes calls to client.tunings.get() and time.sleep(), is not being tested because the initial mocked TuningJob doesn't have a state that would cause the loop to execute.
To properly test the polling logic, you should:
- Patch
time.sleepto prevent test delays. - Mock
client.tunings.get()to return a sequence ofTuningJobobjects that simulate the job's state transitioning from a running state to a completed state.
Here is a more complete test case that replaces the current one:
from unittest.mock import patch, MagicMock
@patch("continuous_tuning_create.time.sleep")
@patch("google.genai.Client")
def test_continuous_tuning_create(mock_genai_client: MagicMock, mock_sleep: MagicMock) -> None:
# Mock the API response for the initial tune call
mock_initial_job = types.TuningJob(
name="test-tuning-job",
state="JOB_STATE_PENDING",
)
mock_genai_client.return_value.tunings.tune.return_value = mock_initial_job
# Mock the responses for the polling `get` call
mock_running_job = types.TuningJob(
name="test-tuning-job",
state="JOB_STATE_RUNNING",
)
mock_succeeded_job = types.TuningJob(
name="test-tuning-job",
state="JOB_STATE_SUCCEEDED",
experiment="test-experiment",
tuned_model=types.TunedModel(
model="test-model-2",
endpoint="test-endpoint",
checkpoints=[]
)
)
mock_genai_client.return_value.tunings.get.side_effect = [
mock_running_job,
mock_succeeded_job
]
response = continuous_tuning_create.create_continuous_tuning_job(tuned_model_name="test-model", checkpoint_id="1")
mock_genai_client.assert_called_once_with(http_options=types.HttpOptions(api_version="v1beta1"))
mock_genai_client.return_value.tunings.tune.assert_called_once()
# Assert that polling happened
assert mock_genai_client.return_value.tunings.get.call_count == 2
mock_sleep.assert_called_with(60)
assert mock_sleep.call_count == 2
assert response == "test-tuning-job"| # limitations under the License. | ||
|
|
||
|
|
||
| def create_continuous_tuning_job(tuned_model_name: str, checkpoint_id: str) -> str: |
There was a problem hiding this comment.
The parameter name tuned_model_name is confusing because it's used as the base_model on line 37. Renaming it to base_model_name would make the function's purpose clearer and improve readability. You will also need to update its usage on line 37.
| def create_continuous_tuning_job(tuned_model_name: str, checkpoint_id: str) -> str: | |
| def create_continuous_tuning_job(base_model_name: str, checkpoint_id: str) -> str: |
| import time | ||
|
|
||
| from google import genai | ||
| from google.genai.types import HttpOptions, TuningDataset, CreateTuningJobConfig |
There was a problem hiding this comment.
According to PEP 8, imports should be at the top of the file, outside of any functions.1 While this pattern is sometimes used in code samples for self-containment, moving the imports to the module level improves readability and is a standard Python convention.
Style Guide References
Footnotes
-
PEP 8 recommends that imports are always put at the top of the file, just after any module comments and docstrings, and before module globals and constants. ↩
| running_states = set([ | ||
| "JOB_STATE_PENDING", | ||
| "JOB_STATE_RUNNING", | ||
| ]) |
There was a problem hiding this comment.
This set of running states is a constant. It's better to define it at the module level (outside the function) with an all-caps name, like RUNNING_STATES, as per PEP 8 conventions for constants.1 This improves readability and avoids recreating the set on each function call. You would then update line 51 to use this new constant.
Style Guide References
Footnotes
-
PEP 8 suggests that constants are usually defined on a module level and written in all capital letters with underscores separating words. ↩
Description
Continuous tuning accepts a model resource name as the base_model and an optional checkpoint ID to run continuous tuning from.
Fixes #
Note: Before submitting a pull request, please open an issue for discussion if you are not associated with Google.
Checklist
nox -s py-3.9(see Test Environment Setup)nox -s lint(see Test Environment Setup)