feat(ingest): show custom model info (#10259)

Author: hsheth2

metadata-ingestion/src/datahub/entrypoints.py

@@ -37,6 +37,7 @@
 from datahub.cli.timeline_cli import timeline
 from datahub.configuration.common import should_show_stack_trace
 from datahub.telemetry import telemetry
+from datahub.utilities._custom_package_loader import model_version_name
 from datahub.utilities.logging_manager import configure_logging
 from datahub.utilities.server_config_util import get_gms_config
 
@@ -99,6 +100,7 @@ def version() -> None:
     """Print version number and exit."""
 
     click.echo(f"DataHub CLI version: {datahub_package.nice_version_name()}")
+    click.echo(f"Models: {model_version_name()}")
     click.echo(f"Python version: {sys.version}")
 
 

metadata-ingestion/src/datahub/ingestion/run/pipeline.py

@@ -39,6 +39,7 @@
 from datahub.ingestion.transformer.transform_registry import transform_registry
 from datahub.metadata.schema_classes import MetadataChangeProposalClass
 from datahub.telemetry import stats, telemetry
+from datahub.utilities._custom_package_loader import model_version_name
 from datahub.utilities.global_warning_util import (
     clear_global_warnings,
     get_global_warnings,
@@ -127,6 +128,7 @@ def _add_init_error_context(step: str) -> Iterator[None]:
 class CliReport(Report):
     cli_version: str = datahub.nice_version_name()
     cli_entry_location: str = datahub.__file__
+    models_version: str = model_version_name()
     py_version: str = sys.version
     py_exec_path: str = sys.executable
     os_details: str = platform.platform()

metadata-ingestion/src/datahub/utilities/_custom_package_loader.py

@@ -21,7 +21,7 @@ def _get_all_registered_custom_packages() -> List[EntryPoint]:
     return list(entry_points(group=_CUSTOM_PACKAGE_GROUP_KEY))
 
 
-def _get_custom_package_for_name(name: str) -> Optional[str]:
+def _get_custom_entrypoint_for_name(name: str) -> Optional[EntryPoint]:
     entrypoints = [
         ep for ep in _get_all_registered_custom_packages() if ep.name == name
     ]
@@ -37,7 +37,37 @@ def _get_custom_package_for_name(name: str) -> Optional[str]:
             f"Multiple custom packages registered for {name}: cannot pick between {all_package_options}"
         )
 
-    return entrypoints[0].value
+    return entrypoints[0]
+
+
+def _get_custom_package_for_name(name: str) -> Optional[str]:
+    entrypoint = _get_custom_entrypoint_for_name(name)
+    if entrypoint:
+        return entrypoint.value
+    return None
+
+
+def model_version_name() -> str:
+    custom_entrypoint = _get_custom_entrypoint_for_name(_MODELS_KEY)
+    if not custom_entrypoint:
+        return "bundled"
+
+    if not custom_entrypoint.dist:
+        return custom_entrypoint.value
+
+    custom_model_package = custom_entrypoint.dist.name
+    version = custom_entrypoint.dist.version
+
+    try:
+        package = __import__(custom_model_package)
+
+        nice_version_name = getattr(package, "nice_version_name", None)
+        if nice_version_name:
+            version = nice_version_name()
+    except ImportError:
+        pass
+
+    return f"{custom_model_package} ({version})"
 
 
 def get_custom_models_package() -> Optional[str]: