|
| 1 | +# Copyright 2024 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 | +# WITHcontent 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 | +import os |
| 16 | +from typing import MutableSequence, Optional |
| 17 | +from unittest import mock |
| 18 | +from unittest.mock import MagicMock |
| 19 | + |
| 20 | +from google.cloud.aiplatform_v1.types import prediction_service |
| 21 | +import google.protobuf.struct_pb2 as struct_pb2 |
| 22 | +from google.protobuf.struct_pb2 import Value |
| 23 | + |
| 24 | +from gemma2_predict_gpu import gemma2_predict_gpu |
| 25 | +from gemma2_predict_tpu import gemma2_predict_tpu |
| 26 | + |
| 27 | +# Global variables |
| 28 | +PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") |
| 29 | +GPU_ENDPOINT_REGION = "us-east1" |
| 30 | +GPU_ENDPOINT_ID = "123456789" # Mock ID used to check if GPU was called |
| 31 | + |
| 32 | +TPU_ENDPOINT_REGION = "us-west1" |
| 33 | +TPU_ENDPOINT_ID = "987654321" # Mock ID used to check if TPU was called |
| 34 | + |
| 35 | +# MOCKED RESPONSE |
| 36 | +MODEL_RESPONSES = """ |
| 37 | +The sky appears blue due to a phenomenon called **Rayleigh scattering**. |
| 38 | +
|
| 39 | +**Here's how it works:** |
| 40 | +
|
| 41 | +1. **Sunlight:** Sunlight is composed of all the colors of the rainbow. |
| 42 | +
|
| 43 | +2. **Earth's Atmosphere:** When sunlight enters the Earth's atmosphere, it collides with tiny particles like nitrogen and oxygen molecules. |
| 44 | +
|
| 45 | +3. **Scattering:** These particles scatter the sunlight in all directions. However, blue light (which has a shorter wavelength) is scattered more effectively than other colors. |
| 46 | +
|
| 47 | +4. **Our Perception:** As a result, we see a blue sky because the scattered blue light reaches our eyes from all directions. |
| 48 | +
|
| 49 | +**Why not other colors?** |
| 50 | +
|
| 51 | +* **Violet light** has an even shorter wavelength than blue and is scattered even more. However, our eyes are less sensitive to violet light, so we perceive the sky as blue. |
| 52 | +* **Longer wavelengths** like red, orange, and yellow are scattered less and travel more directly through the atmosphere. This is why we see these colors during sunrise and sunset, when sunlight has to travel through more of the atmosphere. |
| 53 | +""" |
| 54 | + |
| 55 | + |
| 56 | +# Mocked function - we check if proper format was used depending on selected architecture |
| 57 | +def mock_predict( |
| 58 | + endpoint: Optional[str] = None, |
| 59 | + instances: Optional[MutableSequence[struct_pb2.Value]] = None, |
| 60 | +) -> prediction_service.PredictResponse: |
| 61 | + gpu_endpoint = f"projects/{PROJECT_ID}/locations/{GPU_ENDPOINT_REGION}/endpoints/{GPU_ENDPOINT_ID}" |
| 62 | + tpu_endpoint = f"projects/{PROJECT_ID}/locations/{TPU_ENDPOINT_REGION}/endpoints/{TPU_ENDPOINT_ID}" |
| 63 | + instance_fields = instances[0].struct_value.fields |
| 64 | + |
| 65 | + if endpoint == gpu_endpoint: |
| 66 | + assert "string_value" in instance_fields["inputs"] |
| 67 | + assert "struct_value" in instance_fields["parameters"] |
| 68 | + parameters = instance_fields["parameters"].struct_value.fields |
| 69 | + assert "number_value" in parameters["max_tokens"] |
| 70 | + assert "number_value" in parameters["temperature"] |
| 71 | + assert "number_value" in parameters["top_p"] |
| 72 | + assert "number_value" in parameters["top_k"] |
| 73 | + elif endpoint == tpu_endpoint: |
| 74 | + assert "string_value" in instance_fields["prompt"] |
| 75 | + assert "number_value" in instance_fields["max_tokens"] |
| 76 | + assert "number_value" in instance_fields["temperature"] |
| 77 | + assert "number_value" in instance_fields["top_p"] |
| 78 | + assert "number_value" in instance_fields["top_k"] |
| 79 | + else: |
| 80 | + assert False |
| 81 | + |
| 82 | + response = prediction_service.PredictResponse() |
| 83 | + response.predictions.append(Value(string_value=MODEL_RESPONSES)) |
| 84 | + return response |
| 85 | + |
| 86 | + |
| 87 | +@mock.patch("google.cloud.aiplatform.gapic.PredictionServiceClient") |
| 88 | +def test_gemma2_predict_gpu(mock_client: MagicMock) -> None: |
| 89 | + mock_client_instance = mock_client.return_value |
| 90 | + mock_client_instance.predict = mock_predict |
| 91 | + |
| 92 | + response = gemma2_predict_gpu(GPU_ENDPOINT_REGION, GPU_ENDPOINT_ID) |
| 93 | + assert "Rayleigh scattering" in response |
| 94 | + |
| 95 | + |
| 96 | +@mock.patch("google.cloud.aiplatform.gapic.PredictionServiceClient") |
| 97 | +def test_gemma2_predict_tpu(mock_client: MagicMock) -> None: |
| 98 | + mock_client_instance = mock_client.return_value |
| 99 | + mock_client_instance.predict = mock_predict |
| 100 | + |
| 101 | + response = gemma2_predict_tpu(TPU_ENDPOINT_REGION, TPU_ENDPOINT_ID) |
| 102 | + assert "Rayleigh scattering" in response |
0 commit comments