def test_search_model_versions(mock_store): mock_store.search_model_versions.return_value = [ ModelVersion(name="Model 1", version="1", creation_timestamp=123), ModelVersion(name="Model 1", version="2", creation_timestamp=124) ] result = newModelRegistryClient().search_model_versions("name=Model 1") mock_store.search_model_versions.assert_called_once_with("name=Model 1") assert len(result) == 2
def test_search_model_versions(mock_get_request_message, mock_model_registry_store): mock_get_request_message.return_value = SearchModelVersions( filter="source_path = 'A/B/CD'") mvds = [ ModelVersion(name="model_1", version="5", creation_timestamp=100, last_updated_timestamp=1200, description="v 5", user_id="u1", current_stage="Production", source="A/B/CD", run_id=uuid.uuid4().hex, status="READY", status_message=None), ModelVersion(name="model_1", version="12", creation_timestamp=110, last_updated_timestamp=2000, description="v 12", user_id="u2", current_stage="Production", source="A/B/CD", run_id=uuid.uuid4().hex, status="READY", status_message=None), ModelVersion(name="ads_model", version="8", creation_timestamp=200, last_updated_timestamp=2000, description="v 8", user_id="u1", current_stage="Staging", source="A/B/CD", run_id=uuid.uuid4().hex, status="READY", status_message=None), ModelVersion(name="fraud_detection_model", version="345", creation_timestamp=1000, last_updated_timestamp=1001, description="newest version", user_id="u12", current_stage="None", source="A/B/CD", run_id=uuid.uuid4().hex, status="READY", status_message=None), ] mock_model_registry_store.search_model_versions.return_value = mvds resp = _search_model_versions() args, _ = mock_model_registry_store.search_model_versions.call_args assert args == ("source_path = 'A/B/CD'", ) assert json.loads(resp.get_data()) == {"model_versions": jsonify(mvds)}
def test_get_latest_versions(mock_get_request_message, mock_model_registry_store): name = "model1" mock_get_request_message.return_value = GetLatestVersions(name=name) mvds = [ ModelVersion(name=name, version="5", creation_timestamp=1, last_updated_timestamp=12, description="v 5", user_id="u1", current_stage="Production", source="A/B", run_id=uuid.uuid4().hex, status="READY", status_message=None), ModelVersion(name=name, version="1", creation_timestamp=1, last_updated_timestamp=1200, description="v 1", user_id="u1", current_stage="Archived", source="A/B2", run_id=uuid.uuid4().hex, status="READY", status_message=None), ModelVersion(name=name, version="12", creation_timestamp=100, last_updated_timestamp=None, description="v 12", user_id="u2", current_stage="Staging", source="A/B3", run_id=uuid.uuid4().hex, status="READY", status_message=None), ] mock_model_registry_store.get_latest_versions.return_value = mvds resp = _get_latest_versions() _, args = mock_model_registry_store.get_latest_versions.call_args assert args == {"name": name, "stages": []} assert json.loads(resp.get_data()) == {"model_versions": jsonify(mvds)} for stages in [[], ["None"], ["Staging"], ["Staging", "Production"]]: mock_get_request_message.return_value = GetLatestVersions( name=name, stages=stages) _get_latest_versions() _, args = mock_model_registry_store.get_latest_versions.call_args assert args == {"name": name, "stages": stages}
def to_mlflow_entity(self): return ModelVersion( self.name, self.version, self.creation_time, self.last_updated_time, self.description, self.user_id, self.current_stage, self.source, self.run_id, self.status, self.status_message, [tag.to_mlflow_entity() for tag in self.model_version_tags])
def test_create_model_version(mock_get_request_message, mock_model_registry_store): run_id = uuid.uuid4().hex tags = [ ModelVersionTag(key="key", value="value"), ModelVersionTag(key="anotherKey", value="some other value") ] mock_get_request_message.return_value = CreateModelVersion( name="model_1", source="A/B", run_id=run_id, tags=[tag.to_proto() for tag in tags]) mv = ModelVersion(name="model_1", version="12", creation_timestamp=123, tags=tags) mock_model_registry_store.create_model_version.return_value = mv resp = _create_model_version() _, args = mock_model_registry_store.create_model_version.call_args assert args["name"] == "model_1" assert args["source"] == "A/B" assert args["run_id"] == run_id assert {tag.key: tag.value for tag in args["tags"]} == {tag.key: tag.value for tag in tags} assert json.loads(resp.get_data()) == {"model_version": jsonify(mv)}
def get_model_version(self, name, version): """ Get the model version instance by name and version. :param name: Registered model name. :param version: Registered model version. :return: A single :py:class:`mlflow.entities.model_registry.ModelVersion` object. """ req_body = message_to_json( GetModelVersion(name=name, version=str(version))) response_proto = self._call_endpoint(GetModelVersion, req_body) return ModelVersion.from_proto(response_proto.model_version)
def test_update_model_version(mock_store): name = "Model 1" version = "12" description = "new description" expected_result = ModelVersion(name, version, creation_timestamp=123, description=description) mock_store.update_model_version.return_value = expected_result actal_result = newModelRegistryClient().update_model_version( name, version, "new description") mock_store.update_model_version.assert_called_once_with( name=name, version=version, description="new description") assert expected_result == actal_result
def test_create_model_version(mock_store): name = "Model 1" version = "1" tags_dict = {"key": "value", "another key": "some other value"} tags = [ModelVersionTag(key, value) for key, value in tags_dict.items()] mock_store.create_model_version.return_value = ModelVersion( name=name, version=version, creation_timestamp=123, tags=tags) result = newModelRegistryClient().create_model_version( name, "uri:/for/source", "run123", tags_dict) mock_store.create_model_version.assert_called_once_with( name, "uri:/for/source", "run123", tags) assert result.name == name assert result.version == version assert result.tags == tags_dict
def update_model_version(self, name, version, description): """ Update metadata associated with a model version in backend. :param name: Registered model name. :param version: Registered model version. :param description: New model description. :return: A single :py:class:`mlflow.entities.model_registry.ModelVersion` object. """ req_body = message_to_json( UpdateModelVersion(name=name, version=str(version), description=description)) response_proto = self._call_endpoint(UpdateModelVersion, req_body) return ModelVersion.from_proto(response_proto.model_version)
def _model_version(name, version, stage, source="some:/source", run_id="run13579", tags=None): return ModelVersion(name, version, "2345671890", "234567890", "some description", "UserID", stage, source, run_id, tags=tags)
def test_register_model_with_non_runs_uri(): create_model_patch = mock.patch.object( MlflowClient, "create_registered_model", return_value=RegisteredModel("Model 1")) create_version_patch = mock.patch.object(MlflowClient, "create_model_version", return_value=ModelVersion( "Model 1", "1", creation_timestamp=123)) with create_model_patch, create_version_patch: register_model("s3:/some/path/to/model", "Model 1") MlflowClient.create_registered_model.assert_called_once_with("Model 1") MlflowClient.create_model_version.assert_called_once_with( "Model 1", run_id=None, source="s3:/some/path/to/model")
def get_latest_versions(self, name, stages=None): """ Latest version models for each requested stage. If no ``stages`` argument is provided, returns the latest version for each stage. :param name: Registered model name. :param stages: List of desired stages. If input list is None, return latest versions for for 'Staging' and 'Production' stages. :return: List of :py:class:`mlflow.entities.model_registry.ModelVersion` objects. """ req_body = message_to_json(GetLatestVersions(name=name, stages=stages)) response_proto = self._call_endpoint(GetLatestVersions, req_body) return [ ModelVersion.from_proto(model_version) for model_version in response_proto.model_versions ]
def search_model_versions(self, filter_string): """ Search for model versions in backend that satisfy the filter criteria. :param filter_string: A filter string expression. Currently supports a single filter condition either name of model like ``name = 'model_name'`` or ``run_id = '...'``. :return: PagedList of :py:class:`mlflow.entities.model_registry.ModelVersion` objects. """ req_body = message_to_json(SearchModelVersions(filter=filter_string)) response_proto = self._call_endpoint(SearchModelVersions, req_body) model_versions = [ ModelVersion.from_proto(mvd) for mvd in response_proto.model_versions ] return PagedList(model_versions, response_proto.next_page_token)
def test_register_model_with_existing_registered_model(): create_model_patch = mock.patch.object(MlflowClient, "create_registered_model", side_effect=MlflowException( "Some Message", RESOURCE_ALREADY_EXISTS)) create_version_patch = mock.patch.object(MlflowClient, "create_model_version", return_value=ModelVersion( "Model 1", "1", creation_timestamp=123)) with create_model_patch, create_version_patch: register_model("s3:/some/path/to/model", "Model 1") MlflowClient.create_registered_model.assert_called_once_with("Model 1") MlflowClient.create_model_version.assert_called_once_with( "Model 1", run_id=None, source="s3:/some/path/to/model")
def test_transition_model_version_stage(mock_store): name = "Model 1" version = "12" stage = "Production" expected_result = ModelVersion(name, version, creation_timestamp=123, current_stage=stage) mock_store.transition_model_version_stage.return_value = expected_result actual_result = newModelRegistryClient().transition_model_version_stage( name, version, stage) mock_store.transition_model_version_stage.assert_called_once_with( name=name, version=version, stage=stage, archive_existing_versions=False) assert expected_result == actual_result
def test_register_model_with_runs_uri(): create_model_patch = mock.patch.object( MlflowClient, "create_registered_model", return_value=RegisteredModel("Model 1")) get_uri_patch = mock.patch( "mlflow.store.artifact.runs_artifact_repo.RunsArtifactRepository.get_underlying_uri", return_value="s3:/path/to/source") create_version_patch = mock.patch.object(MlflowClient, "create_model_version", return_value=ModelVersion( "Model 1", "1", creation_timestamp=123)) with get_uri_patch, create_model_patch, create_version_patch: register_model("runs:/run12345/path/to/model", "Model 1") MlflowClient.create_registered_model.assert_called_once_with("Model 1") MlflowClient.create_model_version.assert_called_once_with( "Model 1", "s3:/path/to/source", "run12345")
def create_model_version(self, name, source, run_id, tags=None): """ Create a new model version from given source and run ID. :param name: Registered model name. :param source: Source path where the MLflow model is stored. :param run_id: Run ID from MLflow tracking server that generated the model. :param tags: A list of :py:class:`mlflow.entities.model_registry.ModelVersionTag` instances associated with this model version. :return: A single object of :py:class:`mlflow.entities.model_registry.ModelVersion` created in the backend. """ proto_tags = [tag.to_proto() for tag in tags or []] req_body = message_to_json( CreateModelVersion(name=name, source=source, run_id=run_id, tags=proto_tags)) response_proto = self._call_endpoint(CreateModelVersion, req_body) return ModelVersion.from_proto(response_proto.model_version)
def test_transition_model_version_stage(mock_get_request_message, mock_model_registry_store): name = "model1" version = "32" stage = "Production" mock_get_request_message.return_value = TransitionModelVersionStage( name=name, version=version, stage=stage) mv = ModelVersion(name=name, version=version, creation_timestamp=123, current_stage=stage) mock_model_registry_store.transition_model_version_stage.return_value = mv _transition_stage() _, args = mock_model_registry_store.transition_model_version_stage.call_args assert args == { "name": name, "version": version, "stage": stage, "archive_existing_versions": False }
def test_get_model_version_details(mock_get_request_message, mock_model_registry_store): mock_get_request_message.return_value = GetModelVersion(name="model1", version="32") mvd = ModelVersion(name="model1", version="5", creation_timestamp=1, last_updated_timestamp=12, description="v 5", user_id="u1", current_stage="Production", source="A/B", run_id=uuid.uuid4().hex, status="READY", status_message=None) mock_model_registry_store.get_model_version.return_value = mvd resp = _get_model_version() _, args = mock_model_registry_store.get_model_version.call_args assert args == {"name": "model1", "version": "32"} assert json.loads(resp.get_data()) == {"model_version": jsonify(mvd)}
def test_models_artifact_repo_init_with_stage_uri(host_creds_mock): # pylint: disable=unused-argument model_uri = "models:/MyModel/Production" artifact_location = "dbfs://databricks/mlflow-registry/12345/models/keras-model" model_version_detailed = ModelVersion("MyModel", "10", "2345671890", "234567890", "some description", "UserID", "Production", "source", "run12345") get_latest_versions_patch = mock.patch.object( MlflowClient, "get_latest_versions", return_value=[model_version_detailed]) get_model_version_download_uri_patch = mock.patch.object( MlflowClient, "get_model_version_download_uri", return_value=artifact_location) with get_latest_versions_patch, get_model_version_download_uri_patch: models_repo = ModelsArtifactRepository(model_uri) assert models_repo.artifact_uri == model_uri assert isinstance(models_repo.repo, DbfsRestArtifactRepository) assert models_repo.repo.artifact_uri == artifact_location
def test_update_model_version(mock_get_request_message, mock_model_registry_store): name = "model1" version = "32" description = "Great model!" mock_get_request_message.return_value = UpdateModelVersion( name=name, version=version, description=description) mv = ModelVersion(name=name, version=version, creation_timestamp=123, description=description) mock_model_registry_store.update_model_version.return_value = mv _update_model_version() _, args = mock_model_registry_store.update_model_version.call_args assert args == { "name": name, "version": version, "description": description }
def transition_model_version_stage(self, name, version, stage, archive_existing_versions): """ Update model version stage. :param name: Registered model name. :param version: Registered model version. :param new_stage: New desired stage for this model version. :param archive_existing_versions: If this flag is set to ``True``, all existing model versions in the stage will be automically moved to the "archived" stage. Only valid when ``stage`` is ``"staging"`` or ``"production"`` otherwise an error will be raised. :return: A single :py:class:`mlflow.entities.model_registry.ModelVersion` object. """ req_body = message_to_json( TransitionModelVersionStage( name=name, version=str(version), stage=stage, archive_existing_versions=archive_existing_versions)) response_proto = self._call_endpoint(TransitionModelVersionStage, req_body) return ModelVersion.from_proto(response_proto.model_version)