wehub-resource-sync
725 lines
29 KiB
Python
725 lines
29 KiB
Python
import contextlib
|
|
import importlib
|
|
import importlib.metadata
|
|
import inspect
|
|
import logging
|
|
import threading
|
|
import time
|
|
from typing import Any, Callable
|
|
|
|
import mlflow
|
|
from mlflow.entities import Metric
|
|
from mlflow.utils.validation import MAX_METRICS_PER_BATCH
|
|
|
|
# Define the module-level logger for autologging utilities before importing utilities defined in
|
|
# submodules (e.g., `safety`, `events`) that depend on the module-level logger. Add the `noqa: E402`
|
|
# comment after each subsequent import to ignore "import not at top of file" code style errors
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
# Import autologging utilities used by this module
|
|
from mlflow.ml_package_versions import _ML_PACKAGE_VERSIONS, FLAVOR_TO_MODULE_NAME
|
|
from mlflow.utils.autologging_utils.client import MlflowAutologgingQueueingClient # noqa: F401
|
|
from mlflow.utils.autologging_utils.events import AutologgingEventLogger
|
|
from mlflow.utils.autologging_utils.logging_and_warnings import (
|
|
MlflowEventsAndWarningsBehaviorGlobally,
|
|
NonMlflowWarningsBehaviorForCurrentThread,
|
|
)
|
|
|
|
# Wildcard import other autologging utilities (e.g. safety utilities, event logging utilities) used
|
|
# in autologging integration implementations, which reference them via the
|
|
# `mlflow.utils.autologging_utils` module
|
|
from mlflow.utils.autologging_utils.safety import ( # noqa: F401
|
|
ExceptionSafeAbstractClass,
|
|
ExceptionSafeClass,
|
|
exception_safe_function_for_class,
|
|
is_testing,
|
|
picklable_exception_safe_function,
|
|
revert_patches,
|
|
safe_patch,
|
|
update_wrapper_extended,
|
|
with_managed_run,
|
|
)
|
|
from mlflow.utils.autologging_utils.versioning import (
|
|
get_min_max_version_and_pip_release,
|
|
is_flavor_supported_for_associated_package_versions,
|
|
)
|
|
|
|
INPUT_EXAMPLE_SAMPLE_ROWS = 5
|
|
ENSURE_AUTOLOGGING_ENABLED_TEXT = (
|
|
"please ensure that autologging is enabled before constructing the dataset."
|
|
)
|
|
|
|
# Flag indicating whether autologging is globally disabled for all integrations.
|
|
_AUTOLOGGING_GLOBALLY_DISABLED = False
|
|
|
|
# Autologging config key indicating whether or not a particular autologging integration
|
|
# was configured (i.e. its various `log_models`, `disable`, etc. configuration options
|
|
# were set) via a call to `mlflow.autolog()`, rather than via a call to the integration-specific
|
|
# autologging method (e.g., `mlflow.tensorflow.autolog()`, ...)
|
|
AUTOLOGGING_CONF_KEY_IS_GLOBALLY_CONFIGURED = "globally_configured"
|
|
|
|
# Dict mapping integration name to its config.
|
|
AUTOLOGGING_INTEGRATIONS = {}
|
|
|
|
# Global lock for turning on / off autologging
|
|
# Note "RLock" is required instead of plain lock, for avoid dead-lock
|
|
_autolog_conf_global_lock = threading.RLock()
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
def autologging_conf_lock(fn):
|
|
"""
|
|
Apply a global lock on functions that enable / disable autologging.
|
|
"""
|
|
|
|
def wrapper(*args, **kwargs):
|
|
with _autolog_conf_global_lock:
|
|
return fn(*args, **kwargs)
|
|
|
|
return update_wrapper_extended(wrapper, fn)
|
|
|
|
|
|
def get_mlflow_run_params_for_fn_args(fn, args, kwargs, unlogged=None):
|
|
"""Given arguments explicitly passed to a function, generate a dictionary of MLflow Run
|
|
parameter key / value pairs.
|
|
|
|
Args:
|
|
fn: function whose parameters are to be logged.
|
|
args: arguments explicitly passed into fn. If `fn` is defined on a class,
|
|
`self` should not be part of `args`; the caller is responsible for
|
|
filtering out `self` before calling this function.
|
|
kwargs: kwargs explicitly passed into fn.
|
|
unlogged: parameters not to be logged.
|
|
|
|
Returns:
|
|
A dictionary of MLflow Run parameter key / value pairs.
|
|
"""
|
|
unlogged = unlogged or []
|
|
param_spec = inspect.signature(fn).parameters
|
|
# Filter out `self` from the signature under the assumption that it is not contained
|
|
# within the specified `args`, as stipulated by the documentation
|
|
relevant_params = [param for param in param_spec.values() if param.name != "self"]
|
|
|
|
# Fetch the parameter names for specified positional arguments from the function
|
|
# signature & create a mapping from positional argument name to specified value
|
|
params_to_log = {
|
|
param_info.name: param_val
|
|
for param_info, param_val in zip(list(relevant_params)[: len(args)], args)
|
|
}
|
|
# Add all user-specified keyword arguments to the set of parameters to log
|
|
params_to_log.update(kwargs)
|
|
# Add parameters that were not explicitly specified by the caller to the mapping,
|
|
# using their default values
|
|
params_to_log.update({
|
|
param.name: param.default
|
|
for param in list(relevant_params)[len(args) :]
|
|
if param.name not in kwargs
|
|
})
|
|
# Filter out any parameters that should not be logged, as specified by the `unlogged` parameter
|
|
return {key: value for key, value in params_to_log.items() if key not in unlogged}
|
|
|
|
|
|
def log_fn_args_as_params(fn, args, kwargs, unlogged=None):
|
|
"""Log arguments explicitly passed to a function as MLflow Run parameters to the current active
|
|
MLflow Run.
|
|
|
|
Args:
|
|
fn: function whose parameters are to be logged
|
|
args: arguments explicitly passed into fn. If `fn` is defined on a class,
|
|
`self` should not be part of `args`; the caller is responsible for
|
|
filtering out `self` before calling this function.
|
|
kwargs: kwargs explicitly passed into fn
|
|
unlogged: parameters not to be logged
|
|
|
|
Returns:
|
|
None
|
|
|
|
"""
|
|
params_to_log = get_mlflow_run_params_for_fn_args(fn, args, kwargs, unlogged)
|
|
mlflow.log_params(params_to_log)
|
|
|
|
|
|
class InputExampleInfo:
|
|
"""
|
|
Stores info about the input example collection before it is needed.
|
|
|
|
For example, in xgboost and lightgbm, an InputExampleInfo object is attached to the dataset,
|
|
where its value is read later by the train method.
|
|
|
|
Exactly one of input_example or error_msg should be populated.
|
|
"""
|
|
|
|
def __init__(self, input_example=None, error_msg=None):
|
|
self.input_example = input_example
|
|
self.error_msg = error_msg
|
|
|
|
|
|
def resolve_input_example_and_signature(
|
|
get_input_example, infer_model_signature, log_input_example, log_model_signature, logger
|
|
):
|
|
"""Handles the logic of calling functions to gather the input example and infer the model
|
|
signature.
|
|
|
|
Args:
|
|
get_input_example: Function which returns an input example, usually sliced from a
|
|
dataset. This function can raise an exception, its message will be
|
|
shown to the user in a warning in the logs.
|
|
infer_model_signature: Function which takes an input example and returns the signature
|
|
of the inputs and outputs of the model. This function can raise
|
|
an exception, its message will be shown to the user in a warning
|
|
in the logs.
|
|
log_input_example: Whether to log errors while collecting the input example, and if it
|
|
succeeds, whether to return the input example to the user. We collect
|
|
it even if this parameter is False because it is needed for inferring
|
|
the model signature.
|
|
log_model_signature: Whether to infer and return the model signature.
|
|
logger: The logger instance used to log warnings to the user during input example
|
|
collection and model signature inference.
|
|
|
|
Returns:
|
|
A tuple of input_example and signature. Either or both could be None based on the
|
|
values of log_input_example and log_model_signature.
|
|
|
|
"""
|
|
|
|
input_example = None
|
|
input_example_user_msg = None
|
|
input_example_failure_msg = None
|
|
if log_input_example or log_model_signature:
|
|
try:
|
|
input_example = get_input_example()
|
|
except Exception as e:
|
|
input_example_failure_msg = str(e)
|
|
input_example_user_msg = "Failed to gather input example: " + str(e)
|
|
|
|
model_signature = None
|
|
model_signature_user_msg = None
|
|
if log_model_signature:
|
|
try:
|
|
if input_example is None:
|
|
raise Exception(
|
|
"could not sample data to infer model signature: " + input_example_failure_msg
|
|
)
|
|
model_signature = infer_model_signature(input_example)
|
|
except Exception as e:
|
|
model_signature_user_msg = "Failed to infer model signature: " + str(e)
|
|
|
|
# disable input_example signature inference in model logging if `log_model_signature`
|
|
# is set to `False` or signature inference in autologging fails
|
|
if (
|
|
model_signature is None
|
|
and input_example is not None
|
|
and (not log_model_signature or model_signature_user_msg is not None)
|
|
):
|
|
model_signature = False
|
|
|
|
if log_input_example and input_example_user_msg is not None:
|
|
logger.warning(input_example_user_msg)
|
|
if log_model_signature and model_signature_user_msg is not None:
|
|
logger.warning(model_signature_user_msg)
|
|
|
|
return input_example if log_input_example else None, model_signature
|
|
|
|
|
|
class BatchMetricsLogger:
|
|
"""
|
|
The BatchMetricsLogger will log metrics in batch against an mlflow run.
|
|
If run_id is passed to to constructor then all recording and logging will
|
|
happen against that run_id.
|
|
If no run_id is passed into constructor, then the run ID will be fetched
|
|
from `mlflow.active_run()` each time `record_metrics()` or `flush()` is called; in this
|
|
case, callers must ensure that an active run is present before invoking
|
|
`record_metrics()` or `flush()`.
|
|
"""
|
|
|
|
def __init__(self, run_id=None, tracking_uri=None, model_id=None):
|
|
from mlflow.tracking.client import MlflowClient
|
|
|
|
self.run_id = run_id
|
|
self.model_id = model_id
|
|
self.client = MlflowClient(tracking_uri)
|
|
|
|
# data is an array of Metric objects
|
|
self.data = []
|
|
self.total_training_time = 0
|
|
self.total_log_batch_time = 0
|
|
self.previous_training_timestamp = None
|
|
|
|
def flush(self):
|
|
"""
|
|
The metrics accumulated by BatchMetricsLogger will be batch logged to an MLflow run.
|
|
"""
|
|
self._timed_log_batch()
|
|
self.data = []
|
|
|
|
def _timed_log_batch(self):
|
|
# Retrieving run_id from active mlflow run when run_id is empty.
|
|
current_run_id = mlflow.active_run().info.run_id if self.run_id is None else self.run_id
|
|
|
|
start = time.time()
|
|
metrics_slices = [
|
|
self.data[i : i + MAX_METRICS_PER_BATCH]
|
|
for i in range(0, len(self.data), MAX_METRICS_PER_BATCH)
|
|
]
|
|
for metrics_slice in metrics_slices:
|
|
self.client.log_batch(run_id=current_run_id, metrics=metrics_slice)
|
|
end = time.time()
|
|
self.total_log_batch_time += end - start
|
|
|
|
def _should_flush(self):
|
|
target_training_to_logging_time_ratio = 10
|
|
if (
|
|
self.total_training_time
|
|
>= self.total_log_batch_time * target_training_to_logging_time_ratio
|
|
):
|
|
return True
|
|
|
|
return False
|
|
|
|
def record_metrics(self, metrics, step=None):
|
|
"""
|
|
Submit a set of metrics to be logged. The metrics may not be immediately logged, as this
|
|
class will batch them in order to not increase execution time too much by logging
|
|
frequently.
|
|
|
|
Args:
|
|
metrics: Dictionary containing key, value pairs of metrics to be logged.
|
|
step: The training step that the metrics correspond to.
|
|
"""
|
|
current_timestamp = time.time()
|
|
if self.previous_training_timestamp is None:
|
|
self.previous_training_timestamp = current_timestamp
|
|
|
|
training_time = current_timestamp - self.previous_training_timestamp
|
|
|
|
self.total_training_time += training_time
|
|
|
|
# log_batch() requires step to be defined. Therefore will set step to 0 if not defined.
|
|
if step is None:
|
|
step = 0
|
|
|
|
for key, value in metrics.items():
|
|
self.data.append(
|
|
Metric(key, value, int(current_timestamp * 1000), step, model_id=self.model_id)
|
|
)
|
|
|
|
if self._should_flush():
|
|
self.flush()
|
|
|
|
self.previous_training_timestamp = current_timestamp
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def batch_metrics_logger(run_id: str | None = None, model_id: str | None = None):
|
|
"""
|
|
Context manager that yields a BatchMetricsLogger object, which metrics can be logged against.
|
|
The BatchMetricsLogger keeps metrics in a list until it decides they should be logged, at
|
|
which point the accumulated metrics will be batch logged. The BatchMetricsLogger ensures
|
|
that logging imposes no more than a 10% overhead on the training, where the training is
|
|
measured by adding up the time elapsed between consecutive calls to record_metrics.
|
|
|
|
If logging a batch fails, a warning will be emitted and subsequent metrics will continue to
|
|
be collected.
|
|
|
|
Once the context is closed, any metrics that have yet to be logged will be logged.
|
|
|
|
Args:
|
|
run_id: ID of the run that the metrics will be logged to.
|
|
model_id: ID of the model that the metrics will be associated with.
|
|
"""
|
|
|
|
batch_metrics_logger = BatchMetricsLogger(run_id, model_id=model_id)
|
|
yield batch_metrics_logger
|
|
batch_metrics_logger.flush()
|
|
|
|
|
|
def gen_autologging_package_version_requirements_doc(integration_name):
|
|
"""
|
|
Returns:
|
|
A document note string saying the compatibility for the specified autologging
|
|
integration's associated package versions.
|
|
"""
|
|
min_ver, max_ver, pip_release = get_min_max_version_and_pip_release(integration_name)
|
|
required_pkg_versions = f"``{min_ver}`` <= ``{pip_release}`` <= ``{max_ver}``"
|
|
|
|
return (
|
|
" .. Note:: Autologging is known to be compatible with the following package versions: "
|
|
+ required_pkg_versions
|
|
+ ". Autologging may not succeed when used with package versions outside of this range."
|
|
+ "\n\n"
|
|
)
|
|
|
|
|
|
def _check_and_log_warning_for_unsupported_package_versions(integration_name):
|
|
"""
|
|
If the package version is not supported for autologging, log a warning message.
|
|
|
|
Only check the minimum version, not the maximum version. This is because the "maximum" version
|
|
in the ml-package-versions.yml is only updated per release and it cannot keep up with the pace
|
|
of the package releases. The cross-version tests in MLflow CI runs tests against the latest
|
|
available version, not limited to the "maximum" version, so it is safe to assume it supports
|
|
up to the latest version.
|
|
"""
|
|
if (
|
|
integration_name in FLAVOR_TO_MODULE_NAME
|
|
and not get_autologging_config(integration_name, "disable", True)
|
|
and not get_autologging_config(integration_name, "disable_for_unsupported_versions", False)
|
|
and not is_flavor_supported_for_associated_package_versions(
|
|
integration_name, check_max_version=False
|
|
)
|
|
):
|
|
min_var, _, pip_release = get_min_max_version_and_pip_release(integration_name)
|
|
module = importlib.import_module(FLAVOR_TO_MODULE_NAME[integration_name])
|
|
installed_version = getattr(module, "__version__", None) or importlib.metadata.version(
|
|
pip_release
|
|
)
|
|
_logger.warning(
|
|
f"MLflow {integration_name} autologging is known to be compatible with "
|
|
f"{min_var} <= {pip_release}, but the installed version is {installed_version}. "
|
|
"If you encounter errors during autologging, try upgrading "
|
|
f"/ downgrading {pip_release} to a compatible version, or try upgrading MLflow.",
|
|
)
|
|
|
|
|
|
def autologging_integration(name):
|
|
"""
|
|
**All autologging integrations should be decorated with this wrapper.**
|
|
|
|
Wraps an autologging function in order to store its configuration arguments. This enables
|
|
patch functions to broadly obey certain configurations (e.g., disable=True) without
|
|
requiring specific logic to be present in each autologging integration.
|
|
"""
|
|
|
|
def validate_param_spec(param_spec):
|
|
if "disable" not in param_spec or param_spec["disable"].default is not False:
|
|
raise Exception(
|
|
f"Invalid `autolog()` function for integration '{name}'. `autolog()` functions"
|
|
" must specify a 'disable' argument with default value 'False'"
|
|
)
|
|
elif "silent" not in param_spec or param_spec["silent"].default is not False:
|
|
raise Exception(
|
|
f"Invalid `autolog()` function for integration '{name}'. `autolog()` functions"
|
|
" must specify a 'silent' argument with default value 'False'"
|
|
)
|
|
|
|
def wrapper(_autolog):
|
|
param_spec = inspect.signature(_autolog).parameters
|
|
validate_param_spec(param_spec)
|
|
|
|
AUTOLOGGING_INTEGRATIONS[name] = {}
|
|
default_params = {param.name: param.default for param in param_spec.values()}
|
|
|
|
@autologging_conf_lock
|
|
def autolog(*args, **kwargs):
|
|
config_to_store = dict(default_params)
|
|
config_to_store.update({
|
|
param.name: arg for arg, param in zip(args, param_spec.values())
|
|
})
|
|
config_to_store.update(kwargs)
|
|
AUTOLOGGING_INTEGRATIONS[name] = config_to_store
|
|
|
|
try:
|
|
# Pass `autolog()` arguments to `log_autolog_called` in keyword format to enable
|
|
# event loggers to more easily identify important configuration parameters
|
|
# (e.g., `disable`) without examining positional arguments. Passing positional
|
|
# arguments to `log_autolog_called` is deprecated in MLflow > 1.13.1
|
|
AutologgingEventLogger.get_logger().log_autolog_called(name, (), config_to_store)
|
|
except Exception:
|
|
pass
|
|
|
|
revert_patches(name)
|
|
|
|
# If disabling autologging using fluent api, then every active integration's autolog
|
|
# needs to be called with disable=True. So do not short circuit and let
|
|
# `mlflow.autolog()` invoke all active integrations with disable=True.
|
|
if name != "mlflow" and get_autologging_config(name, "disable", True):
|
|
return
|
|
|
|
is_silent_mode = get_autologging_config(name, "silent", False)
|
|
# Reroute non-MLflow warnings encountered during autologging enablement to an
|
|
# MLflow event logger, and enforce silent mode if applicable (i.e. if the corresponding
|
|
# autologging integration was called with `silent=True`)
|
|
with (
|
|
MlflowEventsAndWarningsBehaviorGlobally(
|
|
# MLflow warnings emitted during autologging setup / enablement are likely
|
|
# actionable and relevant to the user, so they should be emitted as normal
|
|
# when `silent=False`. For reference, see recommended warning and event logging
|
|
# behaviors from https://docs.python.org/3/howto/logging.html#when-to-use-logging
|
|
reroute_warnings=False,
|
|
disable_event_logs=is_silent_mode,
|
|
disable_warnings=is_silent_mode,
|
|
),
|
|
NonMlflowWarningsBehaviorForCurrentThread(
|
|
# non-MLflow warnings emitted during autologging setup / enablement are not
|
|
# actionable for the user, as they are a byproduct of the autologging
|
|
# implementation. Accordingly, they should be rerouted to `logger.warning()`.
|
|
# For reference, see recommended warning and event logging
|
|
# behaviors from https://docs.python.org/3/howto/logging.html#when-to-use-logging
|
|
reroute_warnings=True,
|
|
disable_warnings=is_silent_mode,
|
|
),
|
|
):
|
|
_check_and_log_warning_for_unsupported_package_versions(name)
|
|
|
|
return _autolog(*args, **kwargs)
|
|
|
|
wrapped_autolog = update_wrapper_extended(autolog, _autolog)
|
|
# Set the autologging integration name as a function attribute on the wrapped autologging
|
|
# function, allowing the integration name to be extracted from the function. This is used
|
|
# during the execution of import hooks for `mlflow.autolog()`.
|
|
wrapped_autolog.integration_name = name
|
|
|
|
if name in FLAVOR_TO_MODULE_NAME:
|
|
wrapped_autolog.__doc__ = gen_autologging_package_version_requirements_doc(name) + (
|
|
wrapped_autolog.__doc__ or ""
|
|
)
|
|
return wrapped_autolog
|
|
|
|
return wrapper
|
|
|
|
|
|
def get_autologging_config(flavor_name, config_key, default_value=None):
|
|
"""
|
|
Returns a desired config value for a specified autologging integration.
|
|
|
|
Returns `None` if specified `flavor_name` has no recorded configs.
|
|
If `config_key` is not set on the config object, default value is returned.
|
|
|
|
Args:
|
|
flavor_name: An autologging integration flavor name.
|
|
config_key: The key for the desired config value.
|
|
default_value: The default_value to return.
|
|
"""
|
|
config = AUTOLOGGING_INTEGRATIONS.get(flavor_name)
|
|
if config is not None:
|
|
return config.get(config_key, default_value)
|
|
else:
|
|
return default_value
|
|
|
|
|
|
def autologging_is_disabled(integration_name):
|
|
"""Returns a boolean flag of whether the autologging integration is disabled.
|
|
|
|
Args:
|
|
integration_name: An autologging integration flavor name.
|
|
|
|
"""
|
|
explicit_disabled = get_autologging_config(integration_name, "disable", True)
|
|
if explicit_disabled:
|
|
return True
|
|
|
|
if (
|
|
integration_name in FLAVOR_TO_MODULE_NAME
|
|
and get_autologging_config(integration_name, "disable_for_unsupported_versions", False)
|
|
and not is_flavor_supported_for_associated_package_versions(integration_name)
|
|
):
|
|
return True
|
|
|
|
return False
|
|
|
|
|
|
def is_autolog_supported(integration_name: str) -> bool:
|
|
"""
|
|
Whether the specified autologging integration is supported by the current environment.
|
|
|
|
Args:
|
|
integration_name: An autologging integration flavor name.
|
|
"""
|
|
# NB: We don't check for the presence of autolog() function as it requires importing
|
|
# the flavor module, which may cause import error or overhead.
|
|
return "autologging" in _ML_PACKAGE_VERSIONS.get(integration_name, {})
|
|
|
|
|
|
def get_autolog_function(integration_name: str) -> Callable[..., Any] | None:
|
|
"""
|
|
Get the autolog() function for the specified integration.
|
|
Returns None if the flavor does not have an autolog() function.
|
|
"""
|
|
flavor_module = importlib.import_module(f"mlflow.{integration_name}")
|
|
return getattr(flavor_module, "autolog", None)
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def disable_autologging():
|
|
"""
|
|
Context manager that temporarily disables autologging globally for all integrations upon
|
|
entry and restores the previous autologging configuration upon exit.
|
|
"""
|
|
global _AUTOLOGGING_GLOBALLY_DISABLED
|
|
_AUTOLOGGING_GLOBALLY_DISABLED = True
|
|
try:
|
|
yield
|
|
finally:
|
|
_AUTOLOGGING_GLOBALLY_DISABLED = False
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def disable_discrete_autologging(flavors_to_disable: list[str]) -> None:
|
|
"""
|
|
Context manager for disabling specific autologging integrations temporarily while another
|
|
flavor's autologging is activated. This context wrapper is useful in the event that, for
|
|
example, a particular library calls upon another library within a training API that has a
|
|
current MLflow autologging integration.
|
|
For instance, the transformers library's Trainer class, when running metric scoring,
|
|
builds a sklearn model and runs evaluations as part of its accuracy scoring. Without this
|
|
temporary autologging disabling, a new run will be generated that contains a sklearn model
|
|
that holds no use for tracking purposes as it is only used during the metric evaluation phase
|
|
of training.
|
|
|
|
Args:
|
|
flavors_to_disable: A list of flavors that need to be temporarily disabled while
|
|
executing another flavor's autologging to prevent spurious run
|
|
logging of unrelated models, metrics, and parameters.
|
|
"""
|
|
enabled_flavors = []
|
|
for flavor in flavors_to_disable:
|
|
if not autologging_is_disabled(flavor):
|
|
enabled_flavors.append(flavor)
|
|
autolog_func = getattr(mlflow, flavor)
|
|
autolog_func.autolog(disable=True)
|
|
yield
|
|
for flavor in enabled_flavors:
|
|
autolog_func = getattr(mlflow, flavor)
|
|
autolog_func.autolog(disable=False)
|
|
|
|
|
|
_training_sessions = []
|
|
|
|
|
|
def _get_new_training_session_class():
|
|
"""
|
|
Returns a session manager class for nested autologging runs.
|
|
|
|
Examples
|
|
--------
|
|
>>> class Parent:
|
|
... pass
|
|
>>> class Child:
|
|
... pass
|
|
>>> class Grandchild:
|
|
... pass
|
|
>>>
|
|
>>> _TrainingSession = _get_new_training_session_class()
|
|
>>> with _TrainingSession(Parent, False) as p:
|
|
... with _SklearnTrainingSession(Child, True) as c:
|
|
... with _SklearnTrainingSession(Grandchild, True) as g:
|
|
... print(p.should_log(), c.should_log(), g.should_log())
|
|
True False False
|
|
>>>
|
|
>>> with _TrainingSession(Parent, True) as p:
|
|
... with _TrainingSession(Child, False) as c:
|
|
... with _TrainingSession(Grandchild, True) as g:
|
|
... print(p.should_log(), c.should_log(), g.should_log())
|
|
True True False
|
|
>>>
|
|
>>> with _TrainingSession(Child, True) as c1:
|
|
... with _TrainingSession(Child, True) as c2:
|
|
... print(c1.should_log(), c2.should_log())
|
|
True False
|
|
"""
|
|
|
|
# NOTE: The current implementation doesn't guarantee thread-safety, but that's okay for now
|
|
# because:
|
|
# 1. We don't currently have any use cases for allow_children=True.
|
|
# 2. The list append & pop operations are thread-safe, so we will always clear the session stack
|
|
# once all _TrainingSessions exit.
|
|
class _TrainingSession:
|
|
_session_stack = []
|
|
|
|
def __init__(self, estimator, allow_children=True):
|
|
"""A session manager for nested autologging runs.
|
|
|
|
Args:
|
|
estimator: An estimator that this session originates from.
|
|
allow_children: If True, allows autologging in child sessions.
|
|
If False, disallows autologging in all descendant sessions.
|
|
|
|
"""
|
|
self.allow_children = allow_children
|
|
self.estimator = estimator
|
|
self._parent = None
|
|
|
|
def __enter__(self):
|
|
if len(_TrainingSession._session_stack) > 0:
|
|
self._parent = _TrainingSession._session_stack[-1]
|
|
self.allow_children = (
|
|
_TrainingSession._session_stack[-1].allow_children and self.allow_children
|
|
)
|
|
_TrainingSession._session_stack.append(self)
|
|
return self
|
|
|
|
def __exit__(self, tp, val, traceback):
|
|
_TrainingSession._session_stack.pop()
|
|
|
|
def should_log(self):
|
|
"""
|
|
Returns True when at least one of the following conditions satisfies:
|
|
|
|
1. This session is the root session.
|
|
2. The parent session allows autologging and its estimator differs from this session's
|
|
estimator.
|
|
"""
|
|
for training_session in _TrainingSession._session_stack:
|
|
if training_session is self:
|
|
break
|
|
elif training_session.estimator is self.estimator:
|
|
return False
|
|
|
|
return self._parent is None or self._parent.allow_children
|
|
|
|
@staticmethod
|
|
def is_active():
|
|
return len(_TrainingSession._session_stack) != 0
|
|
|
|
@staticmethod
|
|
def get_current_session():
|
|
if _TrainingSession.is_active():
|
|
return _TrainingSession._session_stack[-1]
|
|
return None
|
|
|
|
_training_sessions.append(_TrainingSession)
|
|
return _TrainingSession
|
|
|
|
|
|
def _has_active_training_session():
|
|
return any(s.is_active() for s in _training_sessions)
|
|
|
|
|
|
def get_instance_method_first_arg_value(method, call_pos_args, call_kwargs):
|
|
"""Get instance method first argument value (exclude the `self` argument).
|
|
|
|
Args:
|
|
method: A `cls.method` object which includes the `self` argument.
|
|
call_pos_args: positional arguments excluding the first `self` argument.
|
|
call_kwargs: keywords arguments.
|
|
"""
|
|
if len(call_pos_args) >= 1:
|
|
return call_pos_args[0]
|
|
else:
|
|
param_sig = inspect.signature(method).parameters
|
|
first_arg_name = list(param_sig.keys())[1]
|
|
assert param_sig[first_arg_name].kind not in [
|
|
inspect.Parameter.VAR_KEYWORD,
|
|
inspect.Parameter.VAR_POSITIONAL,
|
|
]
|
|
return call_kwargs.get(first_arg_name)
|
|
|
|
|
|
def get_method_call_arg_value(arg_index, arg_name, default_value, call_pos_args, call_kwargs):
|
|
"""Get argument value for a method call.
|
|
|
|
Args:
|
|
arg_index: The argument index in the function signature. Start from 0.
|
|
arg_name: The argument name in the function signature.
|
|
default_value: Default argument value.
|
|
call_pos_args: The positional argument values in the method call.
|
|
call_kwargs: The keyword argument values in the method call.
|
|
"""
|
|
if arg_name in call_kwargs:
|
|
return call_kwargs[arg_name]
|
|
elif arg_index < len(call_pos_args):
|
|
return call_pos_args[arg_index]
|
|
else:
|
|
return default_value
|