def test_correct(self):
self.replace_corrupt_and_heavy_models()
request = service_pb2.LoadModelsRequest(models=(
model_pb2.ModelSpec(name="corrupt",
base_path="test/resources/models/corrupt"),
model_pb2.ModelSpec(name="heavy",
base_path="test/resources/models/heavy"),
))
response = self.stub.LoadModels(request)
self.assertTrue(response.success)
# Check both models' state is LOADED
request = service_pb2.ModelStatusRequest(model=model_pb2.ModelSpec(
name="corrupt"))
response = self.stub.GetModelStatus(request)
self.assertTrue(
model_pb2.ModelStatus.ModelState.Name(response.status.state) ==
"LOADED")
request = service_pb2.ModelStatusRequest(model=model_pb2.ModelSpec(
name="heavy"))
response = self.stub.GetModelStatus(request)
self.assertTrue(
model_pb2.ModelStatus.ModelState.Name(response.status.state) ==
"LOADED")
self.revert_model_changes()
def test_reload_models_in_config_file(self):
# Replace heavy and corrupt model and reload them
self.replace_corrupt_and_heavy_models()
response = self.stub.ReloadConfigModels(
service_pb2.ReloadModelsRequest())
# Get loaded models
response = self.stub.GetLoadedModels(service_pb2.LoadedModelsRequest())
loaded_models = [model.name for model in response.models]
# Check that only bad path model wasn't loaded
configured_models = [model["name"] for model in get_config()["models"]]
for model in configured_models:
if model not in loaded_models:
self.assertEqual(model, "bad_path")
# Check haeavy and corrupt model state is LOADED
request = service_pb2.ModelStatusRequest(model=model_pb2.ModelSpec(
name="corrupt"))
response = self.stub.GetModelStatus(request)
self.assertTrue(
model_pb2.ModelStatus.ModelState.Name(response.status.state) ==
"LOADED")
request = service_pb2.ModelStatusRequest(model=model_pb2.ModelSpec(
name="heavy"))
response = self.stub.GetModelStatus(request)
self.assertTrue(
model_pb2.ModelStatus.ModelState.Name(response.status.state) ==
"LOADED")
self.revert_model_changes()
def test_available(self):
request = service_pb2.ModelStatusRequest(model=model_pb2.ModelSpec(
name="heavy"))
response = self.stub.GetModelStatus(request)
self.assertTrue(
model_pb2.ModelStatus.ModelState.Name(response.status.state) ==
"AVAILABLE")
def test_failed(self):
request = service_pb2.ModelStatusRequest(model=model_pb2.ModelSpec(
name="corrupt"))
response = self.stub.GetModelStatus(request)
self.assertTrue(
model_pb2.ModelStatus.ModelState.Name(response.status.state) ==
"FAILED")
def test_unknown(self):
request = service_pb2.ModelStatusRequest(model=model_pb2.ModelSpec(
name="foo"))
response = self.stub.GetModelStatus(request)
self.assertTrue(
model_pb2.ModelStatus.ModelState.Name(response.status.state) ==
"UNKNOWN")
def test_missing_name_and_version(self):
try:
request = model_pb2.ModelSpec(version=1)
self.stub.GetModelStatus(request)
except grpc.RpcError as error:
self.assertTrue(
error._state.code == grpc.StatusCode.INVALID_ARGUMENT)
def test_update_bad_path_model(self):
copytree(self.CORRECT_MODEL_PATH, self.BAD_PATH)
time.sleep(3)
request = service_pb2.ModelStatusRequest(model=model_pb2.ModelSpec(
name="bad_path"))
response = self.stub.GetModelStatus(request)
self.assertTrue(
model_pb2.ModelStatus.ModelState.Name(response.status.state) ==
"LOADED")
rmtree(self.BAD_PATH)
def test_update_corrupt_model(self):
copytree((self.CORRECT_MODEL_PATH / "1"),
(self.CORRUPT_MODEL_PATH / "2"))
time.sleep(3)
request = service_pb2.ModelStatusRequest(model=model_pb2.ModelSpec(
name="corrupt"))
response = self.stub.GetModelStatus(request)
self.assertTrue(
model_pb2.ModelStatus.ModelState.Name(response.status.state) ==
"LOADED")
rmtree(self.CORRUPT_MODEL_PATH / "2")
def test_update_loaded_model(self):
request = service_pb2.ModelStatusRequest(model=model_pb2.ModelSpec(
name="correct"))
previous_response = self.stub.GetModelStatus(request)
self.assertEqual(previous_response.status.version, 1)
copytree((self.CORRECT_MODEL_PATH / "1"),
(self.CORRECT_MODEL_PATH / "2"))
time.sleep(3)
later_response = self.stub.GetModelStatus(request)
self.assertEqual(later_response.status.version, 2)
rmtree(self.CORRECT_MODEL_PATH / "2")
def test_highest_version(self):
request = service_pb2.ModelStatusRequest(model=model_pb2.ModelSpec(
name="correct"))
previous_response = self.stub.GetModelStatus(request)
self.assertEqual(previous_response.status.version, 1)
copytree((self.CORRECT_MODEL_PATH / "1"),
(self.CORRECT_MODEL_PATH / "2"))
self.stub.ReloadConfigModels(service_pb2.ReloadModelsRequest())
later_response = self.stub.GetModelStatus(request)
self.assertEqual(later_response.status.version, 2)
rmtree(self.CORRECT_MODEL_PATH / "2")
def test_corrupt_model(self):
model_spec = model_pb2.ModelSpec(
name="corrupt", base_path="test/resources/models/corrupt")
request = service_pb2.LoadModelsRequest(models=(model_spec, ))
response = self.stub.LoadModels(request)
self.assertFalse(response.success)
# Check model's state is FAILED
request = service_pb2.ModelStatusRequest(model=model_spec)
response = self.stub.GetModelStatus(request)
self.assertTrue(
model_pb2.ModelStatus.ModelState.Name(response.status.state) ==
"FAILED")
def _load_model(self, name: str, base_path: Path):
path = self._get_latest_version_path(base_path)
if path is not None:
# Check if model was already loaded
if (name in self._models and self._models[name].state
== model_pb2.ModelStatus.LOADED):
old_size = self._models[name].size
else:
old_size = 0
# Load model
size = path.stat().st_size * self._memory_factor
if self._available_memory > (size - old_size):
try:
self._models[name] = Model(
model_pb2.ModelSpec(
name=name,
base_path=str(base_path),
version=int(path.parent.name),
),
fasttext.load_model(str(path)),
size,
model_pb2.ModelStatus.LOADED,
)
self._available_memory -= size - old_size
logger.info(f"Model {name} loaded from {path}")
return True
except Exception as ex:
logger.warning(
f"Error loading model {name} from {path}: {ex}")
self._models[name] = Model(
None, None, None, state=model_pb2.ModelStatus.FAILED)
return False
logger.warning(
f"Not enough available memory to load model {name} from {path}"
)
self._models[name] = Model(None,
None,
None,
state=model_pb2.ModelStatus.AVAILABLE)
return False
logger.warning(f"Not model available in {base_path}")
return False
def test_reloading_new_model(self):
# Configure a new model
config_path = os.environ["SERVICE_CONFIG_PATH"]
with open(config_path, "r+") as f:
config = yaml.load(f.read())
config["models"].append({"base_path": "new", "name": "new"})
yaml.dump(config, f)
# Make the new model available and reload models
self.duplicate_correct_model()
# Check that the new model is loaded
request = service_pb2.ModelStatusRequest(model=model_pb2.ModelSpec(
name="new"))
response = self.stub.GetModelStatus(request)
self.assertTrue(
model_pb2.ModelStatus.ModelState.Name(response.status.state) ==
"LOADED")
# Revert changes
self.revert_config_changes()
self.revert_model_changes()
def test_bad_model_path(self):
request = service_pb2.LoadModelsRequest(models=(model_pb2.ModelSpec(
name="foo", base_path="test/resources/models/not_existing"), ))
response = self.stub.LoadModels(request)
self.assertFalse(response.success)
def test_missing_model_name(self):
request = service_pb2.LoadModelsRequest(models=(model_pb2.ModelSpec(
base_path="test/resources/headers"), ))
response = self.stub.LoadModels(request)
self.assertFalse(response.success)