Skip to content

Commit 19ac84e

Browse files
authored
feat(genai/tuning): Add sample for canceling tuning jobs (#14483)
* feat: add support for canceling tuning jobs and include corresponding unit tests * fix: suppress errors during GCS cleanup in tuning test teardown * refactor: refine exception handling in tests and remove unreachable main execution block in tuning job cancel script * refactor: move imports outside of tuning job cancellation function and add start marker * refactor: reorder imports in test file and fix comment indentation in tuning job cancellation script * refactor: cleanup GCS test utility and resolve imports in tuning examples * style: cleanup whitespace in tuning job files * fix: reorder imports in test_tuning_examples.py to resolve import sequence issues
1 parent aefdff9 commit 19ac84e

2 files changed

Lines changed: 39 additions & 2 deletions

File tree

genai/tuning/test_tuning_examples.py

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,14 +13,14 @@
1313
# limitations under the License.
1414

1515
from datetime import datetime as dt
16-
1716
from unittest.mock import call, MagicMock, patch
1817

1918
from google.cloud import storage
2019
from google.genai import types
21-
import pytest
2220

2321
import preference_tuning_job_create
22+
import pytest
23+
import tuning_job_cancel
2424
import tuning_job_create
2525
import tuning_job_get
2626
import tuning_job_list
@@ -89,6 +89,14 @@ def test_tuning_job_get(mock_genai_client: MagicMock) -> None:
8989
assert response == "test-tuning-job"
9090

9191

92+
@patch("google.genai.Client")
93+
def test_tuning_job_cancel(mock_genai_client: MagicMock) -> None:
94+
tuning_job_cancel.cancel_tuning_job("test-tuning-job")
95+
96+
mock_genai_client.assert_called_once_with(http_options=types.HttpOptions(api_version="v1"))
97+
mock_genai_client.return_value.tunings.cancel.assert_called_once_with(name="test-tuning-job")
98+
99+
92100
@patch("google.genai.Client")
93101
def test_tuning_job_list(mock_genai_client: MagicMock) -> None:
94102
# Mock the API response

genai/tuning/tuning_job_cancel.py

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
# [START googlegenaisdk_tuning_job_cancel]
16+
17+
from google import genai
18+
from google.genai.types import HttpOptions
19+
20+
21+
def cancel_tuning_job(tuning_job_name: str) -> None:
22+
client = genai.Client(http_options=HttpOptions(api_version="v1"))
23+
24+
# Cancel the tuning job.
25+
# Eg. tuning_job_name = "projects/123456789012/locations/us-central1/tuningJobs/123456789012345"
26+
client.tunings.cancel(name=tuning_job_name)
27+
28+
29+
# [END googlegenaisdk_tuning_job_cancel]

0 commit comments

Comments
 (0)