feat: add code samples for tuning with intermediate checkpoints (#13366)

yishan-pu · Yishan Pu · web-flow · commit 286361a6bb3e · 2025-05-14T14:30:45.000+10:00
* feat: add code samples for tuning with intermediate checkpoints

* feat: add code samples for tuning with intermediate checkpoints

---------

Co-authored-by: Yishan Pu &lt;yishanpu@google.com&gt;
diff --git a/genai/tuning/requirements.txt b/genai/tuning/requirements.txt
@@ -1 +1 @@
-google-genai==1.7.0
+google-genai==1.15.0
diff --git a/genai/tuning/test_tuning_examples.py b/genai/tuning/test_tuning_examples.py
@@ -12,14 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from unittest.mock import MagicMock, patch
+from unittest.mock import call, MagicMock, patch
 
 from google.genai import types
 
 import tuning_job_create
 import tuning_job_get
 import tuning_job_list
 import tuning_textgen_with_txt
+import tuning_with_checkpoints_create
+import tuning_with_checkpoints_get_model
+import tuning_with_checkpoints_list_checkpoints
+import tuning_with_checkpoints_set_default_checkpoint
+import tuning_with_checkpoints_textgen_with_txt
 
 
 @patch("google.genai.Client")
@@ -113,3 +118,171 @@ def test_tuning_textgen_with_txt(mock_genai_client: MagicMock) -> None:
     mock_genai_client.assert_called_once_with(http_options=types.HttpOptions(api_version="v1"))
     mock_genai_client.return_value.tunings.get.assert_called_once()
     mock_genai_client.return_value.models.generate_content.assert_called_once()
+
+
+@patch("google.genai.Client")
+def test_tuning_job_create_with_checkpoints(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",
+            endpoint="test-endpoint-2",
+            checkpoints=[
+                types.TunedModelCheckpoint(checkpoint_id="1", epoch=1, step=10, endpoint="test-endpoint-1"),
+                types.TunedModelCheckpoint(checkpoint_id="2", epoch=2, step=20, endpoint="test-endpoint-2"),
+            ]
+        )
+    )
+    mock_genai_client.return_value.tunings.tune.return_value = mock_tuning_job
+
+    response = tuning_with_checkpoints_create.create_with_checkpoints()
+
+    mock_genai_client.assert_called_once_with(http_options=types.HttpOptions(api_version="v1"))
+    mock_genai_client.return_value.tunings.tune.assert_called_once()
+    assert response == "test-tuning-job"
+
+
+@patch("google.genai.Client")
+def test_tuning_with_checkpoints_get_model(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",
+            endpoint="test-endpoint-2",
+            checkpoints=[
+                types.TunedModelCheckpoint(checkpoint_id="1", epoch=1, step=10, endpoint="test-endpoint-1"),
+                types.TunedModelCheckpoint(checkpoint_id="2", epoch=2, step=20, endpoint="test-endpoint-2"),
+            ]
+        )
+    )
+    mock_model = types.Model(
+        name="test-model",
+        default_checkpoint_id="2",
+        checkpoints=[
+            types.Checkpoint(checkpoint_id="1", epoch=1, step=10),
+            types.Checkpoint(checkpoint_id="2", epoch=2, step=20),
+        ]
+    )
+    mock_genai_client.return_value.tunings.get.return_value = mock_tuning_job
+    mock_genai_client.return_value.models.get.return_value = mock_model
+
+    response = tuning_with_checkpoints_get_model.get_tuned_model_with_checkpoints("test-tuning-job")
+
+    mock_genai_client.assert_called_once_with(http_options=types.HttpOptions(api_version="v1"))
+    mock_genai_client.return_value.tunings.get.assert_called_once_with(name="test-tuning-job")
+    mock_genai_client.return_value.models.get.assert_called_once_with(model="test-model")
+    assert response == "test-model"
+
+
+@patch("google.genai.Client")
+def test_tuning_with_checkpoints_list_checkpoints(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",
+            endpoint="test-endpoint-2",
+            checkpoints=[
+                types.TunedModelCheckpoint(checkpoint_id="1", epoch=1, step=10, endpoint="test-endpoint-1"),
+                types.TunedModelCheckpoint(checkpoint_id="2", epoch=2, step=20, endpoint="test-endpoint-2"),
+            ]
+        )
+    )
+    mock_genai_client.return_value.tunings.get.return_value = mock_tuning_job
+
+    response = tuning_with_checkpoints_list_checkpoints.list_checkpoints("test-tuning-job")
+
+    mock_genai_client.assert_called_once_with(http_options=types.HttpOptions(api_version="v1"))
+    mock_genai_client.return_value.tunings.get.assert_called_once_with(name="test-tuning-job")
+    assert response == "test-tuning-job"
+
+
+@patch("google.genai.Client")
+def test_tuning_with_checkpoints_set_default_checkpoint(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",
+            endpoint="test-endpoint-2",
+            checkpoints=[
+                types.TunedModelCheckpoint(checkpoint_id="1", epoch=1, step=10, endpoint="test-endpoint-1"),
+                types.TunedModelCheckpoint(checkpoint_id="2", epoch=2, step=20, endpoint="test-endpoint-2"),
+            ]
+        )
+    )
+    mock_model = types.Model(
+        name="test-model",
+        default_checkpoint_id="2",
+        checkpoints=[
+            types.Checkpoint(checkpoint_id="1", epoch=1, step=10),
+            types.Checkpoint(checkpoint_id="2", epoch=2, step=20),
+        ]
+    )
+    mock_updated_model = types.Model(
+        name="test-model",
+        default_checkpoint_id="1",
+        checkpoints=[
+            types.Checkpoint(checkpoint_id="1", epoch=1, step=10),
+            types.Checkpoint(checkpoint_id="2", epoch=2, step=20),
+        ]
+    )
+    mock_genai_client.return_value.tunings.get.return_value = mock_tuning_job
+    mock_genai_client.return_value.models.get.return_value = mock_model
+    mock_genai_client.return_value.models.update.return_value = mock_updated_model
+
+    response = tuning_with_checkpoints_set_default_checkpoint.set_default_checkpoint("test-tuning-job", "1")
+
+    mock_genai_client.assert_called_once_with(http_options=types.HttpOptions(api_version="v1"))
+    mock_genai_client.return_value.tunings.get.assert_called_once_with(name="test-tuning-job")
+    mock_genai_client.return_value.models.get.assert_called_once_with(model="test-model")
+    mock_genai_client.return_value.models.update.assert_called_once()
+    assert response == "1"
+
+
+@patch("google.genai.Client")
+def test_tuning_with_checkpoints_textgen_with_txt(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",
+            endpoint="test-endpoint-2",
+            checkpoints=[
+                types.TunedModelCheckpoint(checkpoint_id="1", epoch=1, step=10, endpoint="test-endpoint-1"),
+                types.TunedModelCheckpoint(checkpoint_id="2", epoch=2, step=20, endpoint="test-endpoint-2"),
+            ]
+        )
+    )
+    mock_response = types.GenerateContentResponse._from_response(  # pylint: disable=protected-access
+        response={
+            "candidates": [
+                {
+                    "content": {
+                        "parts": [{"text": "This is a mocked answer."}]
+                    }
+                }
+            ]
+        },
+        kwargs={},
+    )
+
+    mock_genai_client.return_value.tunings.get.return_value = mock_tuning_job
+    mock_genai_client.return_value.models.generate_content.return_value = mock_response
+
+    tuning_with_checkpoints_textgen_with_txt.test_checkpoint("test-tuning-job")
+
+    mock_genai_client.assert_called_once_with(http_options=types.HttpOptions(api_version="v1"))
+    mock_genai_client.return_value.tunings.get.assert_called_once()
+    assert mock_genai_client.return_value.models.generate_content.call_args_list == [
+        call(model="test-endpoint-2", contents="Why is the sky blue?"),
+        call(model="test-endpoint-1", contents="Why is the sky blue?"),
+        call(model="test-endpoint-2", contents="Why is the sky blue?"),
+    ]
diff --git a/genai/tuning/tuning_job_create.py b/genai/tuning/tuning_job_create.py
@@ -48,6 +48,13 @@ def create_tuning_job() -> str:
     # projects/123456789012/locations/us-central1/endpoints/123456789012345
     # projects/123456789012/locations/us-central1/metadataStores/default/contexts/tuning-experiment-2025010112345678
 
+    if tuning_job.tuned_model.checkpoints:
+        for i, checkpoint in enumerate(tuning_job.tuned_model.checkpoints):
+            print(f"Checkpoint {i + 1}: ", checkpoint)
+        # Example response:
+        # Checkpoint 1:  checkpoint_id='1' epoch=1 step=10 endpoint='projects/123456789012/locations/us-central1/endpoints/123456789000000'
+        # Checkpoint 2:  checkpoint_id='2' epoch=2 step=20 endpoint='projects/123456789012/locations/us-central1/endpoints/123456789012345'
+
     # [END googlegenaisdk_tuning_job_create]
     return tuning_job.name
 
diff --git a/genai/tuning/tuning_with_checkpoints_create.py b/genai/tuning/tuning_with_checkpoints_create.py
@@ -0,0 +1,65 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+def create_with_checkpoints() -> str:
+    # [START googlegenaisdk_tuning_with_checkpoints_create]
+    import time
+
+    from google import genai
+    from google.genai.types import HttpOptions, CreateTuningJobConfig
+
+    client = genai.Client(http_options=HttpOptions(api_version="v1"))
+
+    tuning_job = client.tunings.tune(
+        base_model="gemini-2.0-flash-lite-001",
+        training_dataset="gs://cloud-samples-data/ai-platform/generative_ai/gemini-2_0/text/sft_train_data.jsonl",
+        config=CreateTuningJobConfig(
+            tuned_model_display_name="Example tuning job",
+            # Set to True to disable tuning intermediate checkpoints. Default is False.
+            export_last_checkpoint_only=False,
+        ),
+    )
+
+    running_states = set([
+        "JOB_STATE_PENDING",
+        "JOB_STATE_RUNNING",
+    ])
+
+    while tuning_job.state in running_states:
+        print(tuning_job.state)
+        tuning_job = client.tunings.get(name=tuning_job.name)
+        time.sleep(60)
+
+    print(tuning_job.tuned_model.model)
+    print(tuning_job.tuned_model.endpoint)
+    print(tuning_job.experiment)
+    # Example response:
+    # projects/123456789012/locations/us-central1/models/1234567890@1
+    # projects/123456789012/locations/us-central1/endpoints/123456789012345
+    # projects/123456789012/locations/us-central1/metadataStores/default/contexts/tuning-experiment-2025010112345678
+
+    if tuning_job.tuned_model.checkpoints:
+        for i, checkpoint in enumerate(tuning_job.tuned_model.checkpoints):
+            print(f"Checkpoint {i + 1}: ", checkpoint)
+        # Example response:
+        # Checkpoint 1:  checkpoint_id='1' epoch=1 step=10 endpoint='projects/123456789012/locations/us-central1/endpoints/123456789000000'
+        # Checkpoint 2:  checkpoint_id='2' epoch=2 step=20 endpoint='projects/123456789012/locations/us-central1/endpoints/123456789012345'
+
+    # [END googlegenaisdk_tuning_with_checkpoints_create]
+    return tuning_job.name
+
+
+if __name__ == "__main__":
+    create_with_checkpoints()
diff --git a/genai/tuning/tuning_with_checkpoints_get_model.py b/genai/tuning/tuning_with_checkpoints_get_model.py
@@ -0,0 +1,48 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+def get_tuned_model_with_checkpoints(name: str) -> str:
+    # [START googlegenaisdk_tuning_with_checkpoints_get_model]
+    from google import genai
+    from google.genai.types import HttpOptions
+
+    client = genai.Client(http_options=HttpOptions(api_version="v1"))
+
+    # Get the tuning job and the tuned model.
+    # Eg. name = "projects/123456789012/locations/us-central1/tuningJobs/123456789012345"
+    tuning_job = client.tunings.get(name=name)
+    tuned_model = client.models.get(model=tuning_job.tuned_model.model)
+    print(tuned_model)
+    # Example response:
+    # Model(name='projects/123456789012/locations/us-central1/models/1234567890@1', ...)
+
+    print(f"Default checkpoint: {tuned_model.default_checkpoint_id}")
+    # Example response:
+    # Default checkpoint: 2
+
+    if tuned_model.checkpoints:
+        for _, checkpoint in enumerate(tuned_model.checkpoints):
+            print(f"Checkpoint {checkpoint.checkpoint_id}: ", checkpoint)
+    # Example response:
+    # Checkpoint 1:  checkpoint_id='1' epoch=1 step=10
+    # Checkpoint 2:  checkpoint_id='2' epoch=2 step=20
+
+    # [END googlegenaisdk_tuning_with_checkpoints_get_model]
+    return tuned_model.name
+
+
+if __name__ == "__main__":
+    tuning_job_name = input("Tuning job name: ")
+    get_tuned_model_with_checkpoints(tuning_job_name)
diff --git a/genai/tuning/tuning_with_checkpoints_list_checkpoints.py b/genai/tuning/tuning_with_checkpoints_list_checkpoints.py
@@ -0,0 +1,40 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+def list_checkpoints(name: str) -> str:
+    # [START googlegenaisdk_tuning_with_checkpoints_list_checkpoints]
+    from google import genai
+    from google.genai.types import HttpOptions
+
+    client = genai.Client(http_options=HttpOptions(api_version="v1"))
+
+    # Get the tuning job and the tuned model.
+    # Eg. name = "projects/123456789012/locations/us-central1/tuningJobs/123456789012345"
+    tuning_job = client.tunings.get(name=name)
+
+    if tuning_job.tuned_model.checkpoints:
+        for i, checkpoint in enumerate(tuning_job.tuned_model.checkpoints):
+            print(f"Checkpoint {i + 1}: ", checkpoint)
+    # Example response:
+    # Checkpoint 1:  checkpoint_id='1' epoch=1 step=10 endpoint='projects/123456789012/locations/us-central1/endpoints/123456789000000'
+    # Checkpoint 2:  checkpoint_id='2' epoch=2 step=20 endpoint='projects/123456789012/locations/us-central1/endpoints/123456789012345'
+
+    # [END googlegenaisdk_tuning_with_checkpoints_list_checkpoints]
+    return tuning_job.name
+
+
+if __name__ == "__main__":
+    tuning_job_name = input("Tuning job name: ")
+    list_checkpoints(tuning_job_name)
diff --git a/genai/tuning/tuning_with_checkpoints_set_default_checkpoint.py b/genai/tuning/tuning_with_checkpoints_set_default_checkpoint.py
diff --git a/genai/tuning/tuning_with_checkpoints_textgen_with_txt.py b/genai/tuning/tuning_with_checkpoints_textgen_with_txt.py

Original file line number	Diff line number	Diff line change
`@@ -1 +1 @@`
`1`		`-google-genai==1.7.0`
	`1`	`+google-genai==1.15.0`