model-registry/test/python/test_mr_conn.py

30 lines
682 B
Python

from model_registry import ModelRegistry
def main(server: str, port: int):
mr = ModelRegistry(server, port, author="test", is_secure=False)
model = mr.register_model(
"my-model",
"https://mybucket.uri/",
version="2.0.0",
model_format_name="onnx",
model_format_version="1",
storage_key="my-data-connection",
storage_path="path/to/model",
metadata={
"day": 1,
"split": "train",
},
)
m = mr.get_registered_model("my-model")
assert m
assert model.id == m.id, f"{model} != {m}"
if __name__ == "__main__":
import sys
main(sys.argv[1], int(sys.argv[2]))