chore: resolving flaky test.

PiperOrigin-RevId: 932661713

Author: jscud

tests/system/aiplatform/test_experiment_model.py

@@ -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()
@@ -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()
@@ -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(
@@ -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 == {
@@ -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(
@@ -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()
@@ -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(
@@ -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()
@@ -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(
@@ -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(
@@ -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)