wehub-resource-sync
56 lines
1.9 KiB
Python
56 lines
1.9 KiB
Python
import logging
|
|
import os
|
|
from functools import lru_cache
|
|
|
|
from mlflow.tracking.context.git_context import GitRunContext
|
|
from mlflow.tracking.context.registry import resolve_tags
|
|
from mlflow.utils.databricks_utils import is_in_databricks_notebook
|
|
from mlflow.utils.git_utils import get_git_branch, get_git_commit, get_git_repo_url
|
|
from mlflow.utils.mlflow_tags import (
|
|
MLFLOW_GIT_BRANCH,
|
|
MLFLOW_GIT_COMMIT,
|
|
MLFLOW_GIT_REPO_URL,
|
|
TRACE_RESOLVE_TAGS_ALLOWLIST,
|
|
)
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def resolve_env_metadata():
|
|
"""
|
|
Resolve common environment metadata to be saved in the trace info. These should not
|
|
# change over time, so we resolve them only once. These will be stored in trace
|
|
# metadata rather than tags, because they are immutable.
|
|
"""
|
|
# GitRunContext does not property work in notebook because _get_main_file()
|
|
# points to the kernel launcher file, not the actual notebook file.
|
|
metadata = resolve_tags(ignore=[GitRunContext])
|
|
if not is_in_databricks_notebook():
|
|
# Get Git metadata for the script or notebook. If the notebook is in a
|
|
# Databricks managed Git repo, DatabricksRepoRunContext the metadata
|
|
# so we don't need to run this logic.
|
|
metadata.update(_resolve_git_metadata())
|
|
|
|
return {key: value for key, value in metadata.items() if key in TRACE_RESOLVE_TAGS_ALLOWLIST}
|
|
|
|
|
|
def _resolve_git_metadata():
|
|
try:
|
|
import git # noqa: F401
|
|
except ImportError:
|
|
_logger.debug("Git python package is not installed. Skipping git metadata resolution.")
|
|
return {}
|
|
|
|
try:
|
|
repo = os.getcwd()
|
|
return {
|
|
MLFLOW_GIT_COMMIT: get_git_commit(repo) or "",
|
|
MLFLOW_GIT_REPO_URL: get_git_repo_url(repo) or "",
|
|
MLFLOW_GIT_BRANCH: get_git_branch(repo) or "",
|
|
}
|
|
except Exception:
|
|
_logger.debug("Failed to resolve git metadata", exc_info=True)
|
|
|
|
return {}
|