def test_client_download_model(watchman_service):
"""
Test client's ability to download the model
"""
client = Client(project=tu.GORDO_PROJECT, target=tu.GORDO_SINGLE_TARGET)
models = client.download_model()
assert isinstance(models, dict)
assert isinstance(models[tu.GORDO_SINGLE_TARGET], BaseEstimator)
# Can't download model for non-existent target
with pytest.raises(ValueError):
client = Client(project=tu.GORDO_PROJECT, target="non-existent-target")
client.download_model()
def download_model(ctx: click.Context, output_dir: str):
"""
Download the actual model from the target and write to an output directory
"""
client = Client(*ctx.obj["args"], **ctx.obj["kwargs"])
models = client.download_model()
# Iterate over mapping of models and save into their own sub dirs of the output_dir
for target, model in models.items():
model_out_dir = os.path.join(output_dir, target)
os.mkdir(model_out_dir)
click.secho(
f"Writing model '{target}' to directory: '{model_out_dir}'...",
nl=False)
serializer.dump(model, model_out_dir)
click.secho(f"done")
click.secho(f"Wrote all models to directory: {output_dir}", fg="green")