def version(project, model, mlflow_url, api_client):
mlflow.set_tracking_uri(mlflow_url)
r = mlflow.start_run()
mlflow.end_run()
v = cl.Version(id=1, model_id=model.id, mlflow_run_id=r.info.run_id, mlflow_url=mlflow_url,
artifact_uri="gs://artifacts",
created_at="2019-08-29T08:13:12.377Z",
updated_at="2019-08-29T08:13:12.377Z")
return ModelVersion(v, model, api_client)
class TestModel:
v1 = cl.Version(id=1, model_id=1)
v2 = cl.Version(id=2, model_id=1)
v3 = cl.Version(id=3, model_id=1, labels={"model": "T-800"})
@responses.activate
def test_list_version(self, model):
responses.add("GET",
"/v1/models/1/versions?limit=50&cursor=&search=",
match_querystring=True,
body=json.dumps([self.v1.to_dict()]),
status=200,
adding_headers={"Next-Cursor": "abcdef"},
content_type='application/json')
responses.add("GET",
"/v1/models/1/versions?limit=50&cursor=abcdef&search=",
match_querystring=True,
body=json.dumps([self.v2.to_dict()]),
status=200,
content_type='application/json')
versions = model.list_version()
assert len(versions) == 2
assert versions[0].id == 1
assert versions[1].id == 2
@responses.activate
def test_list_version_with_labels(self, model):
responses.add(
"GET",
"/v1/models/1/versions?limit=50&cursor=&search=labels%3Amodel+in+%28T-800%29",
body=json.dumps([self.v3.to_dict()]),
status=200,
match_querystring=True,
content_type='application/json')
versions = model.list_version({"model": ["T-800"]})
assert len(versions) == 1
assert versions[0].id == 3
assert versions[0].labels["model"] == "T-800"
@responses.activate
def test_list_endpoint(self, model):
responses.add("GET",
'/v1/models/1/endpoints',
body=json.dumps([
mdl_endpoint_1.to_dict(),
mdl_endpoint_2.to_dict()
]),
status=200,
content_type='application/json')
endpoints = model.list_endpoint()
assert len(endpoints) == 2
assert endpoints[0].id == str(mdl_endpoint_1.id)
assert endpoints[1].id == str(mdl_endpoint_2.id)
v = model.get_version(1)
assert v.id == 1
assert model.get_version(3) is None
@responses.activate
def test_serve_traffic(self, model):
ve = VersionEndpoint(ep1)
with pytest.raises(ValueError):
model.serve_traffic([ve], environment_name=env_1.name)
with pytest.raises(ValueError):
model.serve_traffic({ve: -1}, environment_name=env_1.name)
with pytest.raises(ValueError):
model.serve_traffic({ve: 101}, environment_name=env_1.name)
with pytest.raises(ValueError):
model.serve_traffic({VersionEndpoint(ep2): 100},
environment_name=env_1.name)
# test create
responses.add("GET",
'/v1/models/1/endpoints',
body=json.dumps([]),
status=200,
content_type='application/json')
responses.add("POST",
'/v1/models/1/endpoints',
body=json.dumps(mdl_endpoint_1.to_dict()),
status=200,
content_type='application/json')
endpoint = model.serve_traffic({ve: 100},
environment_name=env_1.name)
assert endpoint.id == str(mdl_endpoint_1.id)
assert endpoint.environment_name == env_1.name == mdl_endpoint_1.environment_name
responses.reset()
# test update
responses.add("GET",
'/v1/models/1/endpoints',
body=json.dumps([mdl_endpoint_1.to_dict()]),
status=200,
content_type='application/json')
responses.add("GET",
'/v1/models/1/endpoints/1',
body=json.dumps(mdl_endpoint_1.to_dict()),
status=200,
content_type='application/json')
responses.add("PUT",
'/v1/models/1/endpoints/1',
body=json.dumps(mdl_endpoint_1.to_dict()),
status=200,
content_type='application/json')
endpoint = model.serve_traffic({ve: 100},
environment_name=env_1.name)
assert endpoint.id == str(mdl_endpoint_1.id)
assert endpoint.environment_name == env_1.name == mdl_endpoint_1.environment_name
@responses.activate
def test_stop_serving_traffic(self, model):
ve = VersionEndpoint(ep1)
with pytest.raises(ValueError):
model.serve_traffic([ve], environment_name=env_1.name)
with pytest.raises(ValueError):
model.serve_traffic({ve: -1}, environment_name=env_1.name)
with pytest.raises(ValueError):
model.serve_traffic({ve: 101}, environment_name=env_1.name)
with pytest.raises(ValueError):
model.serve_traffic({VersionEndpoint(ep2): 100},
environment_name=env_1.name)
# test create
responses.add("GET",
'/v1/models/1/endpoints',
body=json.dumps([]),
status=200,
content_type='application/json')
responses.add("POST",
'/v1/models/1/endpoints',
body=json.dumps(mdl_endpoint_1.to_dict()),
status=200,
content_type='application/json')
endpoint = model.serve_traffic({ve: 100},
environment_name=env_1.name)
assert endpoint.id == str(mdl_endpoint_1.id)
assert endpoint.environment_name == env_1.name == mdl_endpoint_1.environment_name
responses.reset()
# test DELETE
responses.reset()
responses.add("GET",
'/v1/models/1/endpoints',
body=json.dumps([mdl_endpoint_1.to_dict()]),
status=200,
content_type='application/json')
responses.add("GET",
'/v1/models/1/endpoints/1',
body=json.dumps(mdl_endpoint_1.to_dict()),
status=200,
content_type='application/json')
responses.add("DELETE",
'/v1/models/1/endpoints/1',
status=200,
content_type='application/json')
model.stop_serving_traffic(endpoint.environment_name)
assert len(responses.calls) == 2
@responses.activate
def test_serve_traffic_default_env(self, model):
ve = VersionEndpoint(ep1)
# no default environment
responses.add("GET",
'/v1/environments',
body=json.dumps([env_2.to_dict()]),
status=200,
content_type='application/json')
with pytest.raises(ValueError):
model.serve_traffic({ve: 100})
responses.reset()
# test create
responses.add("GET",
'/v1/environments',
body=json.dumps([env_1.to_dict(),
env_2.to_dict()]),
status=200,
content_type='application/json')
responses.add("GET",
'/v1/models/1/endpoints',
body=json.dumps([]),
status=200,
content_type='application/json')
responses.add("POST",
'/v1/models/1/endpoints',
body=json.dumps(mdl_endpoint_1.to_dict()),
status=200,
content_type='application/json')
endpoint = model.serve_traffic({ve: 100})
assert endpoint.id == str(mdl_endpoint_1.id)
assert endpoint.environment_name == env_1.name == mdl_endpoint_1.environment_name
responses.reset()
# test update
responses.add("GET",
'/v1/environments',
body=json.dumps([env_1.to_dict(),
env_2.to_dict()]),
status=200,
content_type='application/json')
responses.add("GET",
'/v1/models/1/endpoints',
body=json.dumps([mdl_endpoint_1.to_dict()]),
status=200,
content_type='application/json')
responses.add("GET",
'/v1/models/1/endpoints/1',
body=json.dumps(mdl_endpoint_1.to_dict()),
status=200,
content_type='application/json')
responses.add("PUT",
'/v1/models/1/endpoints/1',
body=json.dumps(mdl_endpoint_1.to_dict()),
status=200,
content_type='application/json')
endpoint = model.serve_traffic({ve: 100})
assert endpoint.id == str(mdl_endpoint_1.id)
assert endpoint.environment_name == env_1.name == mdl_endpoint_1.environment_name
class TestModel:
v1 = cl.Version(id=1, model_id=1)
v2 = cl.Version(id=2, model_id=1)
@responses.activate
def test_get_version(self, model):
responses.add("GET",
"/v1/models/1/versions",
body=json.dumps([self.v1.to_dict(),
self.v2.to_dict()]),
status=200,
content_type='application/json')
v = model.get_version(1)
assert v.id == 1
assert model.get_version(3) is None
@responses.activate
def test_endpoint(self, model):
responses.add("GET",
'/v1/models/1/endpoints',
body=json.dumps([mdl_endpoint_2.to_dict()]),
status=200,
content_type='application/json')
ep = model.endpoint
assert ep is None
responses.reset()
responses.add("GET",
'/v1/models/1/endpoints',
body=json.dumps([
mdl_endpoint_1.to_dict(),
mdl_endpoint_2.to_dict()
]),
status=200,
content_type='application/json')
ep = model.endpoint
assert ep.id == str(mdl_endpoint_1.id)
assert ep.status.value == mdl_endpoint_1.status
assert ep.environment_name == mdl_endpoint_1.environment_name
assert ep.url.startswith(f"http://{mdl_endpoint_1.url}")
@responses.activate
def test_list_endpoint(self, model):
responses.add("GET",
'/v1/models/1/endpoints',
body=json.dumps([
mdl_endpoint_1.to_dict(),
mdl_endpoint_2.to_dict()
]),
status=200,
content_type='application/json')
endpoints = model.list_endpoint()
assert len(endpoints) == 2
assert endpoints[0].id == str(mdl_endpoint_1.id)
assert endpoints[1].id == str(mdl_endpoint_2.id)
@responses.activate
def test_serve_traffic(self, model):
ve = VersionEndpoint(ep1)
with pytest.raises(ValueError):
model.serve_traffic([ve], environment_name=env_1.name)
with pytest.raises(ValueError):
model.serve_traffic({ve: -1}, environment_name=env_1.name)
with pytest.raises(ValueError):
model.serve_traffic({ve: 101}, environment_name=env_1.name)
with pytest.raises(ValueError):
model.serve_traffic({VersionEndpoint(ep2): 100},
environment_name=env_1.name)
# test create
responses.add("GET",
'/v1/models/1/endpoints',
body=json.dumps([]),
status=200,
content_type='application/json')
responses.add("POST",
'/v1/models/1/endpoints',
body=json.dumps(mdl_endpoint_1.to_dict()),
status=200,
content_type='application/json')
endpoint = model.serve_traffic({ve: 100},
environment_name=env_1.name)
assert endpoint.id == str(mdl_endpoint_1.id)
assert endpoint.environment_name == env_1.name == mdl_endpoint_1.environment_name
responses.reset()
# test update
responses.add("GET",
'/v1/models/1/endpoints',
body=json.dumps([mdl_endpoint_1.to_dict()]),
status=200,
content_type='application/json')
responses.add("GET",
'/v1/models/1/endpoints/1',
body=json.dumps(mdl_endpoint_1.to_dict()),
status=200,
content_type='application/json')
responses.add("PUT",
'/v1/models/1/endpoints/1',
body=json.dumps(mdl_endpoint_1.to_dict()),
status=200,
content_type='application/json')
endpoint = model.serve_traffic({ve: 100},
environment_name=env_1.name)
assert endpoint.id == str(mdl_endpoint_1.id)
assert endpoint.environment_name == env_1.name == mdl_endpoint_1.environment_name
@responses.activate
def test_stop_serving_traffic(self, model):
ve = VersionEndpoint(ep1)
with pytest.raises(ValueError):
model.serve_traffic([ve], environment_name=env_1.name)
with pytest.raises(ValueError):
model.serve_traffic({ve: -1}, environment_name=env_1.name)
with pytest.raises(ValueError):
model.serve_traffic({ve: 101}, environment_name=env_1.name)
with pytest.raises(ValueError):
model.serve_traffic({VersionEndpoint(ep2): 100},
environment_name=env_1.name)
# test create
responses.add("GET",
'/v1/models/1/endpoints',
body=json.dumps([]),
status=200,
content_type='application/json')
responses.add("POST",
'/v1/models/1/endpoints',
body=json.dumps(mdl_endpoint_1.to_dict()),
status=200,
content_type='application/json')
endpoint = model.serve_traffic({ve: 100},
environment_name=env_1.name)
assert endpoint.id == str(mdl_endpoint_1.id)
assert endpoint.environment_name == env_1.name == mdl_endpoint_1.environment_name
responses.reset()
# test DELETE
responses.reset()
responses.add("GET",
'/v1/models/1/endpoints',
body=json.dumps([mdl_endpoint_1.to_dict()]),
status=200,
content_type='application/json')
responses.add("GET",
'/v1/models/1/endpoints/1',
body=json.dumps(mdl_endpoint_1.to_dict()),
status=200,
content_type='application/json')
responses.add("DELETE",
'/v1/models/1/endpoints/1',
status=200,
content_type='application/json')
model.stop_serving_traffic(endpoint.environment_name)
assert len(responses.calls) == 2
@responses.activate
def test_serve_traffic_default_env(self, model):
ve = VersionEndpoint(ep1)
# no default environment
responses.add("GET",
'/v1/environments',
body=json.dumps([env_2.to_dict()]),
status=200,
content_type='application/json')
with pytest.raises(ValueError):
model.serve_traffic({ve: 100})
responses.reset()
# test create
responses.add("GET",
'/v1/environments',
body=json.dumps([env_1.to_dict(),
env_2.to_dict()]),
status=200,
content_type='application/json')
responses.add("GET",
'/v1/models/1/endpoints',
body=json.dumps([]),
status=200,
content_type='application/json')
responses.add("POST",
'/v1/models/1/endpoints',
body=json.dumps(mdl_endpoint_1.to_dict()),
status=200,
content_type='application/json')
endpoint = model.serve_traffic({ve: 100})
assert endpoint.id == str(mdl_endpoint_1.id)
assert endpoint.environment_name == env_1.name == mdl_endpoint_1.environment_name
responses.reset()
# test update
responses.add("GET",
'/v1/environments',
body=json.dumps([env_1.to_dict(),
env_2.to_dict()]),
status=200,
content_type='application/json')
responses.add("GET",
'/v1/models/1/endpoints',
body=json.dumps([mdl_endpoint_1.to_dict()]),
status=200,
content_type='application/json')
responses.add("GET",
'/v1/models/1/endpoints/1',
body=json.dumps(mdl_endpoint_1.to_dict()),
status=200,
content_type='application/json')
responses.add("PUT",
'/v1/models/1/endpoints/1',
body=json.dumps(mdl_endpoint_1.to_dict()),
status=200,
content_type='application/json')
endpoint = model.serve_traffic({ve: 100})
assert endpoint.id == str(mdl_endpoint_1.id)
assert endpoint.environment_name == env_1.name == mdl_endpoint_1.environment_name