Skip to content

chore: resolving flaky test. #6882

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

chore: resolving flaky test. #6882

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 12 additions & 21 deletions tests/system/aiplatform/test_experiment_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,9 +60,8 @@ def test_sklearn_model(self, shared_state):

# Test get ExperimentModel with aritfact id
model_artifact = aiplatform.get_experiment_model("sk-model")
assert model_artifact.uri.endswith("sklearn-model")

shared_state["resources"] = [model_artifact]
assert model_artifact.uri.endswith("sklearn-model")

# Test get model info from ExperimentModel
model_info = model_artifact.get_model_info()
Expand Down Expand Up @@ -117,9 +116,8 @@ def test_xgboost_booster_with_custom_uri(self, shared_state):

# Test get ExperimentModel with aritfact id
model_artifact = aiplatform.get_experiment_model("xgb-booster")
assert model_artifact.uri == uri

shared_state["resources"].append(model_artifact)
assert model_artifact.uri == uri

# Test get model info from ExperimentModel
model_info = model_artifact.get_model_info()
Expand All @@ -144,10 +142,9 @@ def test_xgboost_booster_with_custom_uri(self, shared_state):
registered_model = model_artifact.register_model(
serving_container_image_uri=container_uri
)
assert registered_model.display_name.startswith("xgboost model")

self.registered_models_cpu.append(registered_model)
shared_state["resources"].append(registered_model)
assert registered_model.display_name.startswith("xgboost model")

def test_xgboost_xgbmodel_with_custom_names(self, shared_state):
aiplatform.init(
Expand All @@ -170,11 +167,10 @@ def test_xgboost_xgbmodel_with_custom_names(self, shared_state):

# Test get ExperimentModel with aritfact id
model_artifact = aiplatform.get_experiment_model("xgboost-xgbmodel")
shared_state["resources"].append(model_artifact)
assert model_artifact.uri.endswith("xgboost-model")
assert model_artifact.display_name == "custom-experiment-model-name"

shared_state["resources"].append(model_artifact)

# Test get model info from ExperimentModel
model_info = model_artifact.get_model_info()
assert model_info == {
Expand All @@ -201,10 +197,9 @@ def test_xgboost_xgbmodel_with_custom_names(self, shared_state):
serving_container_image_uri=container_uri,
display_name="custom-registered-model-name",
)
assert registered_model.display_name == "custom-registered-model-name"

self.registered_models_cpu.append(registered_model)
shared_state["resources"].append(registered_model)
assert registered_model.display_name == "custom-registered-model-name"

def test_tensorflow_keras_model_with_input_example(self, shared_state):
aiplatform.init(
Expand All @@ -230,9 +225,8 @@ def test_tensorflow_keras_model_with_input_example(self, shared_state):

# Test get ExperimentModel with aritfact id
model_artifact = aiplatform.get_experiment_model("keras-model")
assert model_artifact.uri.endswith("tensorflow-model")

shared_state["resources"].append(model_artifact)
assert model_artifact.uri.endswith("tensorflow-model")

# Test get model info from ExperimentModel
model_info = model_artifact.get_model_info()
Expand Down Expand Up @@ -261,10 +255,9 @@ def test_tensorflow_keras_model_with_input_example(self, shared_state):
registered_model = model_artifact.register_model(
serving_container_image_uri=container_uri
)
assert registered_model.display_name.startswith("tensorflow model")

self.registered_models_cpu.append(registered_model)
shared_state["resources"].append(registered_model)
assert registered_model.display_name.startswith("tensorflow model")

def test_tensorflow_module_with_gpu_container(self, shared_state):
aiplatform.init(
Expand Down Expand Up @@ -294,9 +287,8 @@ def add(self, x):

# Test get ExperimentModel with aritfact id
model_artifact = aiplatform.get_experiment_model("tf-module")
assert model_artifact.uri.endswith("tensorflow-model")

shared_state["resources"].append(model_artifact)
assert model_artifact.uri.endswith("tensorflow-model")

# Test get model info from ExperimentModel
model_info = model_artifact.get_model_info()
Expand All @@ -322,10 +314,9 @@ def add(self, x):
serving_container_image_uri=container_uri,
use_gpu=True,
)
assert registered_model.display_name.startswith("tensorflow model")

self.registered_models_gpu.append(registered_model)
shared_state["resources"].append(registered_model)
assert registered_model.display_name.startswith("tensorflow model")

def test_deploy_model_with_cpu_container(self, shared_state):
aiplatform.init(
Expand All @@ -342,8 +333,8 @@ def test_deploy_model_with_cpu_container(self, shared_state):
endpoint = registered_model.deploy()

pred = endpoint.predict([[1, 2]])
assert isinstance(pred, models.Prediction)
shared_state["resources"].append(endpoint)
assert isinstance(pred, models.Prediction)

def test_deploy_model_with_gpu_container(self, shared_state):
aiplatform.init(
Expand All @@ -359,9 +350,9 @@ def test_deploy_model_with_gpu_container(self, shared_state):
# Deploy the registered model
# Since we are using gpu, we need to specify accelerator_type and count
endpoint = registered_model.deploy(
accelerator_type="NVIDIA_TESLA_T4", accelerator_count=1, sync=False
accelerator_type="NVIDIA_TESLA_T4", accelerator_count=1
)
shared_state["resources"].append(endpoint)

pred = endpoint.predict([[1, 2]])
assert isinstance(pred, models.Prediction)
shared_state["resources"].append(endpoint)
Loading