def test_with_latest_model_versions(self): name = random_str() mvd_1 = ModelVersion( name, "1", 1000, 2000, "version 1", "user 1", "Production", "source 1", "run ID 1", "PENDING_REGISTRATION", "Model version is in production!", ) mvd_2 = ModelVersion( name, "4", 1300, 2002, "version 4", "user 2", "Staging", "source 4", "run ID 12", "READY", "Model copied over!", ) as_dict = { "name": name, "creation_timestamp": 1, "last_updated_timestamp": 4000, "description": random_str(), "latest_versions": [mvd_1, mvd_2], "tags": [], } rmd_1 = RegisteredModel.from_dictionary(as_dict) as_dict["tags"] = {} self.assertEqual(dict(rmd_1), as_dict) proto = rmd_1.to_proto() self.assertEqual(proto.creation_timestamp, 1) self.assertEqual(proto.last_updated_timestamp, 4000) self.assertEqual(set([mvd.version for mvd in proto.latest_versions]), set(["1", "4"])) self.assertEqual(set([mvd.name for mvd in proto.latest_versions]), set([name])) self.assertEqual( set([mvd.current_stage for mvd in proto.latest_versions]), set(["Production", "Staging"]), ) self.assertEqual( set([mvd.last_updated_timestamp for mvd in proto.latest_versions]), set([2000, 2002]) ) self.assertEqual( set([mvd.creation_timestamp for mvd in proto.latest_versions]), set([1300, 1000]) )
def test_creation_and_hydration(self): name = random_str() t1, t2 = 100, 150 source = "path/to/source" run_id = uuid.uuid4().hex run_link = "http://localhost:5000/path/to/run" tags = [ModelVersionTag("key", "value"), ModelVersionTag("randomKey", "not a random value")] mvd = ModelVersion(name, "5", t1, t2, "version five", "user 1", "Production", source, run_id, "READY", "Model version #5 is ready to use.", tags, run_link) self._check(mvd, name, "5", t1, t2, "version five", "user 1", "Production", source, run_id, "READY", "Model version #5 is ready to use.", {tag.key: tag.value for tag in (tags or [])}) expected_dict = { "name": name, "version": "5", "creation_timestamp": t1, "last_updated_timestamp": t2, "description": "version five", "user_id": "user 1", "current_stage": "Production", "source": source, "run_id": run_id, "run_link": run_link, "status": "READY", "status_message": "Model version #5 is ready to use.", "tags": {tag.key: tag.value for tag in (tags or [])}} model_version_as_dict = dict(mvd) self.assertEqual(model_version_as_dict, expected_dict) proto = mvd.to_proto() self.assertEqual(proto.name, name) self.assertEqual(proto.version, "5") self.assertEqual(proto.status, ModelVersionStatus.from_string("READY")) self.assertEqual(proto.status_message, "Model version #5 is ready to use.") self.assertEqual(set([tag.key for tag in proto.tags]), set(["key", "randomKey"])) self.assertEqual(set([tag.value for tag in proto.tags]), set(["value", "not a random value"])) mvd_2 = ModelVersion.from_proto(proto) self._check(mvd_2, name, "5", t1, t2, "version five", "user 1", "Production", source, run_id, "READY", "Model version #5 is ready to use.", {tag.key: tag.value for tag in (tags or [])}) expected_dict.update({"registered_model": RegisteredModel(name)}) expected_dict["tags"] = tags mvd_3 = ModelVersion.from_dictionary(expected_dict) self._check(mvd_3, name, "5", t1, t2, "version five", "user 1", "Production", source, run_id, "READY", "Model version #5 is ready to use.", {tag.key: tag.value for tag in (tags or [])})
def test_string_repr(self): model_version = ModelVersion( name="myname", version="43", creation_timestamp=12, last_updated_timestamp=100, description="This is a test model.", user_id="user one", current_stage="Archived", source="path/to/a/notebook", run_id="some run", run_link="http://localhost:5000/path/to/run", status="PENDING_REGISTRATION", status_message="Copying!", tags=[], ) assert ( str(model_version) == "<ModelVersion: creation_timestamp=12, " "current_stage='Archived', description='This is a test " "model.', last_updated_timestamp=100, " "name='myname', " "run_id='some run', run_link='http://localhost:5000/path/" "to/run', source='path/to/a/notebook', " "status='PENDING_REGISTRATION', status_message='Copying!', " "tags={}, user_id='user one', version='43'>" )
def test_creation_and_hydration(self): name = random_str() t1, t2 = 100, 150 source = "path/to/source" run_id = uuid.uuid4().hex mvd = ModelVersion(name, "5", t1, t2, "version five", "user 1", "Production", source, run_id, "READY", "Model version #5 is ready to use.") self._check(mvd, name, "5", t1, t2, "version five", "user 1", "Production", source, run_id, "READY", "Model version #5 is ready to use.") expected_dict = { "name": name, "version": "5", "creation_timestamp": t1, "last_updated_timestamp": t2, "description": "version five", "user_id": "user 1", "current_stage": "Production", "source": source, "run_id": run_id, "status": "READY", "status_message": "Model version #5 is ready to use." } model_version_as_dict = dict(mvd) self.assertEqual(model_version_as_dict, expected_dict) proto = mvd.to_proto() self.assertEqual(proto.name, name) self.assertEqual(proto.version, "5") self.assertEqual(proto.status, ModelVersionStatus.from_string("READY")) self.assertEqual(proto.status_message, "Model version #5 is ready to use.") mvd_2 = ModelVersion.from_proto(proto) self._check(mvd_2, name, "5", t1, t2, "version five", "user 1", "Production", source, run_id, "READY", "Model version #5 is ready to use.") expected_dict.update({"registered_model": RegisteredModel(name)}) mvd_3 = ModelVersion.from_dictionary(expected_dict) self._check(mvd_3, name, "5", t1, t2, "version five", "user 1", "Production", source, run_id, "READY", "Model version #5 is ready to use.")
def from_proto(cls, proto): # input: mlflow.protos.model_registry_pb2.RegisteredModel # returns RegisteredModel entity registered_model = cls( proto.name, proto.creation_timestamp, proto.last_updated_timestamp, proto.description, [ModelVersion.from_proto(mvd) for mvd in proto.latest_versions]) for tag in proto.tags: registered_model._add_tag(RegisteredModelTag.from_proto(tag)) return registered_model
def test_creation_and_hydration(self): name = random_str() rm = RegisteredModel(name) model_version = ModelVersion(rm, 100) self._check(model_version, name, 100) expected_dict = {"version": 100} model_version_as_dict = dict(model_version) self.assertEqual(model_version_as_dict["version"], 100) self.assertIsInstance(model_version_as_dict["registered_model"], RegisteredModel) self.assertEqual(model_version_as_dict["registered_model"].name, name) model_version_as_dict.pop("registered_model") self.assertEqual(model_version_as_dict, expected_dict) proto = model_version.to_proto() self.assertEqual(proto.registered_model.name, name) self.assertEqual(proto.version, 100) model_version_2 = ModelVersion.from_proto(proto) self._check(model_version_2, name, 100) expected_dict.update({"registered_model": RegisteredModel(name)}) model_version_3 = ModelVersion.from_dictionary(expected_dict) self._check(model_version_3, name, 100)
def test_string_repr(self): model_version = ModelVersion(name="myname", version="43", creation_timestamp=12, last_updated_timestamp=100, description="This is a test model.", user_id="user one", current_stage="Archived", source="path/to/a/notebook", run_id="some run", status="PENDING_REGISTRATION", status_message="Copying!") assert str(model_version) == "<ModelVersion: creation_timestamp=12, " \ "current_stage='Archived', description='This is a test " \ "model.', last_updated_timestamp=100, " \ "name='myname', " \ "run_id='some run', source='path/to/a/notebook', " \ "status='PENDING_REGISTRATION', status_message='Copying!', " \ "user_id='user one', version='43'>"
def _properties(cls): # aggregate with base class properties since cls.__dict__ does not do it automatically return sorted(cls._get_properties_helper() + ModelVersion._properties())
def test_string_repr(self): model_version = ModelVersion(RegisteredModel(name="myname"), version=4) assert str(model_version) == "<ModelVersion: registered_model=<RegisteredModel: " \ "name='myname'>, version=4>"