def cleanup_on_error(self): """Method ends mlflow run with correct exit status for failed runs. Note that this method does not work when a pipeline running in dagit fails, it seems that in this case a different process runs the pipeline and when it fails the stack trace is therefore not available. For this case we can use the cleanup_on_failure hook defined below. """ any_error = sys.exc_info() if any_error[1]: if isinstance(any_error[1], KeyboardInterrupt): mlflow.end_run(status=RunStatus.to_string(RunStatus.KILLED)) else: mlflow.end_run(status=RunStatus.to_string(RunStatus.FAILED))
def fit_mlflow(self, clazz, func_name, *args, **kwargs): """ Autologging function that performs model training by executing the training method referred to be `func_name` on the instance of `clazz` referred to by `self` & records MLflow parameters, metrics, tags, and artifacts to a corresponding MLflow Run. """ should_start_run = mlflow.active_run() is None if should_start_run: try_mlflow_log(mlflow.start_run) _log_pretraining_metadata(self, *args, **kwargs) original_fit = gorilla.get_original_attribute(clazz, func_name) try: fit_output = original_fit(self, *args, **kwargs) except Exception as e: if should_start_run: try_mlflow_log(mlflow.end_run, RunStatus.to_string(RunStatus.FAILED)) raise e _log_posttraining_metadata(self, *args, **kwargs) if should_start_run: try_mlflow_log(mlflow.end_run) return fit_output
def end_mlflow_run_on_pipeline_finished(context, event_list): for event in event_list: if event.is_step_success: _cleanup_on_success(context) elif event.is_step_failure: mlf = context.resources.mlflow mlf.end_run(status=RunStatus.to_string(RunStatus.FAILED))
def get_status(self): """Gets the human-readable status of the MLflow run from the tracking server.""" if not self._active_run: eprint( "Can't get MLflow run status; the run's status has not been " "persisted to an accessible tracking server.") return None return RunStatus.to_string(self._active_run.get_run().info.status)
def patch_with_managed_run(original, *args, **kwargs): managed_run = None if not mlflow.active_run(): managed_run = try_mlflow_log(mlflow.start_run) try: result = patch_function(original, *args, **kwargs) except: if managed_run: try_mlflow_log(mlflow.end_run, RunStatus.to_string(RunStatus.FAILED)) raise else: if managed_run: try_mlflow_log(mlflow.end_run, RunStatus.to_string(RunStatus.FINISHED)) return result
def patch_with_managed_run(original, *args, **kwargs): managed_run = None if not mlflow.active_run(): managed_run = try_mlflow_log(create_managed_run) try: result = patch_function(original, *args, **kwargs) except (Exception, KeyboardInterrupt): # In addition to standard Python exceptions, handle keyboard interrupts to ensure # that runs are terminated if a user prematurely interrupts training execution # (e.g. via sigint / ctrl-c) if managed_run: try_mlflow_log(mlflow.end_run, RunStatus.to_string(RunStatus.FAILED)) raise else: if managed_run: try_mlflow_log(mlflow.end_run, RunStatus.to_string(RunStatus.FINISHED)) return result
def _hook(context, event_list): for event in event_list: if event.is_step_success: _cleanup_on_success(context) elif event.is_step_failure: mlf = context.resources.mlflow mlf.end_run(status=RunStatus.to_string(RunStatus.FAILED)) return HookExecutionResult(hook_name=name, is_skipped=False)
def fit_mlflow(self, func_name, *args, **kwargs): should_start_run = mlflow.active_run() is None if should_start_run: try_mlflow_log(mlflow.start_run) # TODO: We should not log nested estimator parameters for # parameter search estimators (GridSearchCV, RandomizedSearchCV) # Chunk and truncate model parameters to avoid hitting the log_batch API limit for chunk in _chunk_dict(self.get_params(deep=True), chunk_size=MAX_PARAMS_TAGS_PER_BATCH): truncated = _truncate_dict(chunk, MAX_ENTITY_KEY_LENGTH, MAX_PARAM_VAL_LENGTH) try_mlflow_log(mlflow.log_params, truncated) try_mlflow_log( mlflow.set_tags, { "estimator_name": self.__class__.__name__, "estimator_class": self.__class__.__module__ + "." + self.__class__.__name__, }, ) original_fit = gorilla.get_original_attribute(self, func_name) try: fit_output = original_fit(*args, **kwargs) except Exception as e: if should_start_run: try_mlflow_log(mlflow.end_run, RunStatus.to_string(RunStatus.FAILED)) raise e if hasattr(self, "score"): try: score_args = _get_args_for_score(self.score, self.fit, args, kwargs) training_score = self.score(*score_args) except Exception as e: # pylint: disable=broad-except msg = ( self.score.__qualname__ + " failed. The 'training_score' metric will not be recorded. Scoring error: " + str(e)) _logger.warning(msg) else: try_mlflow_log(mlflow.log_metric, "training_score", training_score) try_mlflow_log(log_model, self, artifact_path="model") if should_start_run: try_mlflow_log(mlflow.end_run) return fit_output
def from_proto(cls, proto): end_time = proto.end_time # The proto2 default scalar value of zero indicates that the run's end time is absent. # An absent end time is represented with a NoneType in the `RunInfo` class if end_time == 0: end_time = None return cls(run_uuid=proto.run_uuid, run_id=proto.run_id, experiment_id=proto.experiment_id, user_id=proto.user_id, status=RunStatus.to_string(proto.status), start_time=proto.start_time, end_time=end_time, lifecycle_stage=proto.lifecycle_stage, artifact_uri=proto.artifact_uri)
def _patch_implementation(self, original, *args, **kwargs): if not mlflow.active_run(): self.managed_run = try_mlflow_log(create_managed_run) result = super(PatchWithManagedRun, self)._patch_implementation( original, *args, **kwargs ) if self.managed_run: try_mlflow_log(mlflow.end_run, RunStatus.to_string(RunStatus.FINISHED)) return result
def fit_mlflow(self, func_name, *args, **kwargs): should_start_run = mlflow.active_run() is None if should_start_run: try_mlflow_log(mlflow.start_run) _log_pretraining_metadata(self, *args, **kwargs) original_fit = gorilla.get_original_attribute(self, func_name) try: fit_output = original_fit(*args, **kwargs) except Exception as e: if should_start_run: try_mlflow_log(mlflow.end_run, RunStatus.to_string(RunStatus.FAILED)) raise e _log_posttraining_metadata(self, *args, **kwargs) if should_start_run: try_mlflow_log(mlflow.end_run) return fit_output
def _on_exception(self, e): if self.managed_run: try_mlflow_log(mlflow.end_run, RunStatus.to_string(RunStatus.FAILED)) super(PatchWithManagedRun, self)._on_exception(e)
def get_status(self): return RunStatus.to_string(self._get_status())