def test_with_managed_run_sets_specified_run_tags(): client = MlflowClient() tags_to_set = { "foo": "bar", "num_layers": "7", } patch_function_1 = with_managed_run( lambda original, *args, **kwargs: mlflow.active_run(), tags=tags_to_set) run1 = patch_function_1(lambda: "foo") assert tags_to_set.items() <= client.get_run( run1.info.run_id).data.tags.items() class PatchFunction2(PatchFunction): def _patch_implementation(self, original, *args, **kwargs): return mlflow.active_run() def _on_exception(self, exception): pass patch_function_2 = with_managed_run(PatchFunction2, tags=tags_to_set) run2 = patch_function_2.call(lambda: "foo") assert tags_to_set.items() <= client.get_run( run2.info.run_id).data.tags.items()
def test_with_managed_run_with_throwing_class_exhibits_expected_behavior(): client = MlflowClient() patch_function_active_run = None @with_managed_run class TestPatch(PatchFunction): def _patch_implementation(self, original, *args, **kwargs): nonlocal patch_function_active_run patch_function_active_run = mlflow.active_run() raise Exception("bad implementation") def _on_exception(self, exception): pass with pytest.raises(Exception): TestPatch.call(lambda: "foo") assert patch_function_active_run is not None status1 = client.get_run(patch_function_active_run.info.run_id).info.status assert RunStatus.from_string(status1) == RunStatus.FAILED with mlflow.start_run() as active_run, pytest.raises(Exception): TestPatch.call(lambda: "foo") assert patch_function_active_run == active_run # `with_managed_run` should not terminate a preexisting MLflow run, # even if the patch function throws status2 = client.get_run(active_run.info.run_id).info.status assert RunStatus.from_string(status2) == RunStatus.FINISHED
def _set_run_tag(run_id, path, version, data_format): client = MlflowClient() table_info_string = _get_table_info_string(path, version, data_format) existing_run = client.get_run(run_id) existing_tag = existing_run.data.tags.get(_SPARK_TABLE_INFO_TAG_NAME) new_table_info = _merge_tag_lines(existing_tag, table_info_string) client.set_tag(run_id, _SPARK_TABLE_INFO_TAG_NAME, new_table_info)
def test_safe_patch_manages_run_if_specified_and_sets_expected_run_tags( patch_destination, test_autologging_integration ): client = MlflowClient() active_run = None def patch_impl(original, *args, **kwargs): nonlocal active_run active_run = mlflow.active_run() return original(*args, **kwargs) with mock.patch( "mlflow.utils.autologging_utils.safety.with_managed_run", wraps=with_managed_run ) as managed_run_mock: safe_patch( test_autologging_integration, patch_destination, "fn", patch_impl, manage_run=True ) patch_destination.fn() assert managed_run_mock.call_count == 1 assert active_run is not None assert active_run.info.run_id is not None assert ( client.get_run(active_run.info.run_id).data.tags[MLFLOW_AUTOLOGGING] == "test_integration" )
def test_with_managed_run_with_non_throwing_function_exhibits_expected_behavior(): client = MlflowClient() @with_managed_run def patch_function(original, *args, **kwargs): return mlflow.active_run() run1 = patch_function(lambda: "foo") run1_status = client.get_run(run1.info.run_id).info.status assert RunStatus.from_string(run1_status) == RunStatus.FINISHED with mlflow.start_run() as active_run: run2 = patch_function(lambda: "foo") assert run2 == active_run run2_status = client.get_run(run2.info.run_id).info.status assert RunStatus.from_string(run2_status) == RunStatus.FINISHED
def get_best_model(**kwargs): """ For all the models we logged, determine the best performing run """ ids = [r for ids in kwargs["ti"].xcom_pull(task_ids=["model_lr", "model_rf"], key="run_id") for r in ids] client = MlflowClient() runs = [client.get_run(run_id) for run_id in ids] run_r2 = [run.data.metrics["r2"] for run in runs] best_run = runs[np.argmax(run_r2)] kwargs["ti"].xcom_push(key="best_model_run_id", value=best_run.info.run_id)
def test_with_managed_run_with_non_throwing_class_exhibits_expected_behavior(): client = MlflowClient() @with_managed_run class TestPatch(PatchFunction): def _patch_implementation(self, original, *args, **kwargs): return mlflow.active_run() def _on_exception(self, exception): pass run1 = TestPatch.call(lambda: "foo") run1_status = client.get_run(run1.info.run_id).info.status assert RunStatus.from_string(run1_status) == RunStatus.FINISHED with mlflow.start_run() as active_run: run2 = TestPatch.call(lambda: "foo") assert run2 == active_run run2_status = client.get_run(run2.info.run_id).info.status assert RunStatus.from_string(run2_status) == RunStatus.FINISHED
def extract_run_meta_metrics(experiment, float_cols=[], filter_unfinished=True): client = MlflowClient() # pull runs for experiment experiment_id = client.get_experiment_by_name(experiment).experiment_id run_infos = client.list_run_infos(experiment_id) runs = [client.get_run(run_info.run_id) for run_info in run_infos] # extract run metrics and metadata from runs run_metrics = [] run_metas = [] for run in runs: # extract metadata from metrics run_id = run.info.run_id run_status = run.info.status # filter out unfinished runs if specified if filter_unfinished and run_status != "FINISHED": continue time_ms = (run.info.end_time - run.info.start_time if run.info.status == "FINISHED" else np.nan) meta = { "artifact_uri": run.info.artifact_uri, "time_ms": time_ms, "run_id": run_id, "status": run.info.status, } meta.update(run.data.params) meta.update(run.data.tags) meta.update(run.data.metrics) run_metas.append(meta) # metrics history have to be extracted separately as they have # a separate epoch dimension metric_names = run.data.metrics.keys() histories = [client.get_metric_history(run_id, name) for name in metric_names] for history in histories: for metric in history: metric_dict = { "run_id": run_id, "metric": metric.key, "step": metric.step, "value": metric.value, } run_metrics.append(metric_dict) run_metrics.append(metric_dict) meta_df = pd.DataFrame(run_metas).astype({ col: "float64" for col in float_cols }).set_index("run_id") metrics_df = pd.DataFrame(run_metrics) return meta_df, metrics_df
def _poll_once( client: MlflowClient, prev_run_id: Optional[str] = None ) -> Optional[str]: lst = client.get_latest_versions(MLFLOW_MODEL_NAME, stages=[MLFLOW_MODEL_STAGE]) if not lst: logger.info(f"No model name={MLFLOW_MODEL_NAME} in stage={MLFLOW_MODEL_STAGE}") logger.info(f"Deleting existing model deployments.") subprocess.run( f"kubectl -n {K8S_NS} delete --ignore-not-found cm {MS_CONFIGMAP_NAME}", shell=True, check=True, ) subprocess.run( f"kubectl -n {K8S_NS} delete --ignore-not-found deployment {MS_DEPLOYMENT_NAME}", shell=True, check=True, ) return None assert len(lst) == 1, (len(lst), lst) m: ModelVersion = lst[0] if m.run_id == prev_run_id: logger.info(f"Already deployed model: run_id={prev_run_id} version={m.version}") else: logger.info(f"Found a model to deploy! run_id={m.run_id} version={m.version}") logger.info(f"Model details:\n{m.to_proto()}") r: Run = client.get_run(m.run_id) logger.info(f"Model run details:\n{yaml.dump(r.to_dictionary())}") logger.info(f"Deploying model run_id={m.run_id} version={m.version}") k = f"kubectl -n {K8S_NS}" configmap_path = _render_configmap(MS_CONFIGMAP_TEMPLATE, run_id=m.run_id) # logger.debug(f"Rendered configmap:\n{configmap_path.read_text()}") subprocess.run(f"{k} apply -f {configmap_path}", shell=True, check=True) # No zero-downtime :( subprocess.run( f"{k} delete --ignore-not-found deployment {MS_DEPLOYMENT_NAME}", shell=True, check=True, ) subprocess.run( f"envsubst '$$GIT_BRANCH' < {MS_DEPLOYMENT_PATH} | {k} apply -f -", shell=True, check=True, ) logger.info(f"") logger.info( f"[+] Successfully patched model deployment to run_id={m.run_id} version={m.version}" ) logger.info(f"") return m.run_id
def test_with_managed_run_ends_run_on_keyboard_interrupt(): client = MlflowClient() run = None def original(): nonlocal run run = mlflow.active_run() raise KeyboardInterrupt patch_function_1 = with_managed_run( "test_integration", lambda original, *args, **kwargs: original(*args, **kwargs) ) with pytest.raises(KeyboardInterrupt): patch_function_1(original) assert not mlflow.active_run() run_status_1 = client.get_run(run.info.run_id).info.status assert RunStatus.from_string(run_status_1) == RunStatus.FAILED class PatchFunction2(PatchFunction): def _patch_implementation(self, original, *args, **kwargs): return original(*args, **kwargs) def _on_exception(self, exception): pass patch_function_2 = with_managed_run("test_integration", PatchFunction2) with pytest.raises(KeyboardInterrupt): patch_function_2.call(original) assert not mlflow.active_run() run_status_2 = client.get_run(run.info.run_id).info.status assert RunStatus.from_string(run_status_2) == RunStatus.FAILED
def _validate_autologging_run(autologging_integration, run_id): """ For testing purposes, verifies that an MLflow run produced by an `autologging_integration` satisfies the following properties: - The run has an autologging tag whose value is the name of the autologging integration - The run has a terminal status (e.g., KILLED, FAILED, FINISHED) """ client = MlflowClient() run = client.get_run(run_id) autologging_tag_value = run.data.tags.get(MLFLOW_AUTOLOGGING) assert autologging_tag_value == autologging_integration, ( "Autologging run with id {} failed to set autologging tag with expected value. Expected: " "'{}', Actual: '{}'".format(run_id, autologging_integration, autologging_tag_value) ) assert RunStatus.is_terminated( RunStatus.from_string(run.info.status) ), "Autologging run with id {} has a non-terminal status '{}'".format(run_id, run.info.status)