def test_deploy_with_traffic_percent(self, deploy_model_mock, sync):
with mock.patch.object(endpoint_service_client.EndpointServiceClient,
"get_endpoint") as get_endpoint_mock:
get_endpoint_mock.return_value = gca_endpoint.Endpoint(
display_name=_TEST_DISPLAY_NAME,
name=_TEST_ENDPOINT_NAME,
traffic_split={"model1": 100},
)
test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME)
test_model = models.Model(_TEST_ID)
test_endpoint.deploy(model=test_model,
traffic_percentage=70,
sync=sync)
if not sync:
test_endpoint.wait()
automatic_resources = gca_machine_resources.AutomaticResources(
min_replica_count=1,
max_replica_count=1,
)
deployed_model = gca_endpoint.DeployedModel(
automatic_resources=automatic_resources,
model=test_model.resource_name,
display_name=None,
)
deploy_model_mock.assert_called_once_with(
endpoint=test_endpoint.resource_name,
deployed_model=deployed_model,
traffic_split={
"model1": 30,
"0": 70
},
metadata=(),
)
def test_undeploy_with_traffic_split(self, undeploy_model_mock, sync):
with mock.patch.object(endpoint_service_client.EndpointServiceClient,
"get_endpoint") as get_endpoint_mock:
get_endpoint_mock.return_value = gca_endpoint.Endpoint(
display_name=_TEST_DISPLAY_NAME,
name=_TEST_ENDPOINT_NAME,
traffic_split={
"model1": 40,
"model2": 60
},
)
test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME)
test_endpoint.undeploy(
deployed_model_id="model1",
traffic_split={
"model1": 0,
"model2": 100
},
sync=sync,
)
if not sync:
test_endpoint.wait()
undeploy_model_mock.assert_called_once_with(
endpoint=test_endpoint.resource_name,
deployed_model_id="model1",
traffic_split={"model2": 100},
metadata=(),
)
def create_endpoint_mock():
with mock.patch.object(endpoint_service_client.EndpointServiceClient,
"create_endpoint") as create_endpoint_mock:
create_endpoint_lro_mock = mock.Mock(ga_operation.Operation)
create_endpoint_lro_mock.result.return_value = gca_endpoint.Endpoint(
name=_TEST_ENDPOINT_NAME, display_name=_TEST_DISPLAY_NAME)
create_endpoint_mock.return_value = create_endpoint_lro_mock
yield create_endpoint_mock
def get_endpoint_with_models_mock():
with mock.patch.object(endpoint_service_client.EndpointServiceClient,
"get_endpoint") as get_endpoint_mock:
get_endpoint_mock.return_value = gca_endpoint.Endpoint(
display_name=_TEST_DISPLAY_NAME,
name=_TEST_ENDPOINT_NAME,
deployed_models=_TEST_DEPLOYED_MODELS,
)
yield get_endpoint_mock
def get_endpoint_alt_location_mock():
with mock.patch.object(endpoint_service_client.EndpointServiceClient,
"get_endpoint") as get_endpoint_mock:
get_endpoint_mock.return_value = gca_endpoint.Endpoint(
display_name=_TEST_DISPLAY_NAME,
name=_TEST_ENDPOINT_NAME_ALT_LOCATION,
encryption_spec=_TEST_ENCRYPTION_SPEC,
)
yield get_endpoint_mock
def get_endpoint_mock():
with mock.patch.object(endpoint_service_client.EndpointServiceClient,
"get_endpoint") as get_endpoint_mock:
test_endpoint_resource_name = endpoint_service_client.EndpointServiceClient.endpoint_path(
_TEST_PROJECT, _TEST_LOCATION, _TEST_ID)
get_endpoint_mock.return_value = gca_endpoint.Endpoint(
display_name=_TEST_MODEL_NAME,
name=test_endpoint_resource_name,
)
yield get_endpoint_mock
def test_create_with_description(self, create_endpoint_mock, sync):
aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
my_endpoint = models.Endpoint.create(
display_name=_TEST_DISPLAY_NAME, description=_TEST_DESCRIPTION, sync=sync
)
if not sync:
my_endpoint.wait()
expected_endpoint = gca_endpoint.Endpoint(
display_name=_TEST_DISPLAY_NAME, description=_TEST_DESCRIPTION,
)
create_endpoint_mock.assert_called_once_with(
parent=_TEST_PARENT, endpoint=expected_endpoint, metadata=(),
)
def test_create_with_labels(self, create_endpoint_mock, sync):
my_endpoint = models.Endpoint.create(display_name=_TEST_DISPLAY_NAME,
labels=_TEST_LABELS,
sync=sync)
if not sync:
my_endpoint.wait()
expected_endpoint = gca_endpoint.Endpoint(
display_name=_TEST_DISPLAY_NAME,
labels=_TEST_LABELS,
)
create_endpoint_mock.assert_called_once_with(
parent=_TEST_PARENT,
endpoint=expected_endpoint,
metadata=(),
)
def test_create(self, create_endpoint_mock, sync):
aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
my_endpoint = models.Endpoint.create(
display_name=_TEST_DISPLAY_NAME,
encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME,
sync=sync,
)
if not sync:
my_endpoint.wait()
expected_endpoint = gca_endpoint.Endpoint(
display_name=_TEST_DISPLAY_NAME, encryption_spec=_TEST_ENCRYPTION_SPEC
)
create_endpoint_mock.assert_called_once_with(
parent=_TEST_PARENT, endpoint=expected_endpoint, metadata=(),
)
expected_endpoint.name = _TEST_ENDPOINT_NAME
assert my_endpoint._gca_resource == expected_endpoint
def test_create(self, create_endpoint_mock, sync):
my_endpoint = models.Endpoint.create(
display_name=_TEST_DISPLAY_NAME,
encryption_spec_key_name=_TEST_ENCRYPTION_KEY_NAME,
sync=sync,
)
if not sync:
my_endpoint.wait()
expected_endpoint = gca_endpoint.Endpoint(
display_name=_TEST_DISPLAY_NAME,
encryption_spec=_TEST_ENCRYPTION_SPEC)
create_endpoint_mock.assert_called_once_with(
parent=_TEST_PARENT,
endpoint=expected_endpoint,
metadata=(),
)
expected_endpoint.name = _TEST_ENDPOINT_NAME
assert my_endpoint.gca_resource == expected_endpoint
assert my_endpoint.network is None
def test_undeploy(self, undeploy_model_mock, sync):
aiplatform.init(project=_TEST_PROJECT, location=_TEST_LOCATION)
with mock.patch.object(
endpoint_service_client.EndpointServiceClient, "get_endpoint"
) as get_endpoint_mock:
get_endpoint_mock.return_value = gca_endpoint.Endpoint(
display_name=_TEST_DISPLAY_NAME,
name=_TEST_ENDPOINT_NAME,
traffic_split={"model1": 100},
)
test_endpoint = models.Endpoint(_TEST_ENDPOINT_NAME)
assert dict(test_endpoint._gca_resource.traffic_split) == {"model1": 100}
test_endpoint.undeploy("model1", sync=sync)
if not sync:
test_endpoint.wait()
undeploy_model_mock.assert_called_once_with(
endpoint=test_endpoint.resource_name,
deployed_model_id="model1",
traffic_split={},
# traffic_split={"model1": 0},
metadata=(),
)
"medv":
aiplatform.explain.ExplanationMetadata.OutputMetadata(
{"output_tensor_name": "dense_2"})
},
)
_TEST_EXPLANATION_PARAMETERS = aiplatform.explain.ExplanationParameters(
{"sampled_shapley_attribution": {
"path_count": 10
}})
# CMEK encryption
_TEST_ENCRYPTION_KEY_NAME = "key_1234"
_TEST_ENCRYPTION_SPEC = gca_encryption_spec.EncryptionSpec(
kms_key_name=_TEST_ENCRYPTION_KEY_NAME)
_TEST_ENDPOINT_GAPIC = gca_endpoint.Endpoint(display_name=_TEST_DISPLAY_NAME,
name=_TEST_ENDPOINT_NAME)
_TEST_ENDPOINT_LIST = [
gca_endpoint.Endpoint(
name=_TEST_ENDPOINT_NAME,
display_name="aac",
create_time=datetime.now() - timedelta(minutes=15),
),
gca_endpoint.Endpoint(
name=_TEST_ENDPOINT_NAME,
display_name="aab",
create_time=datetime.now() - timedelta(minutes=5),
),
gca_endpoint.Endpoint(
name=_TEST_ENDPOINT_NAME,
display_name="aaa",
},
)
_TEST_EXPLANATION_PARAMETERS = aiplatform.explain.ExplanationParameters(
{"sampled_shapley_attribution": {"path_count": 10}}
)
# CMEK encryption
_TEST_ENCRYPTION_KEY_NAME = "key_1234"
_TEST_ENCRYPTION_SPEC = gca_encryption_spec.EncryptionSpec(
kms_key_name=_TEST_ENCRYPTION_KEY_NAME
)
_TEST_ENDPOINT_LIST = [
gca_endpoint.Endpoint(
display_name="aac", create_time=datetime.now() - timedelta(minutes=15)
),
gca_endpoint.Endpoint(
display_name="aab", create_time=datetime.now() - timedelta(minutes=5)
),
gca_endpoint.Endpoint(
display_name="aaa", create_time=datetime.now() - timedelta(minutes=10)
),
]
_TEST_LIST_FILTER = 'display_name="abc"'
_TEST_LIST_ORDER_BY_CREATE_TIME = "create_time desc"
_TEST_LIST_ORDER_BY_DISPLAY_NAME = "display_name"
@pytest.fixture
