From 8303e977b6f4b1968b7c36f400c55ed4457c0641 Mon Sep 17 00:00:00 2001 From: Miles <187469313+miles2542@users.noreply.github.com> Date: Sun, 19 Apr 2026 11:27:22 +0700 Subject: [PATCH] feat: add retry logic to model artifact upload, fall back to local if timeout - more robust on flaky networks --- src/rossmann_ops/train_model.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/src/rossmann_ops/train_model.py b/src/rossmann_ops/train_model.py index e309c50..784f593 100644 --- a/src/rossmann_ops/train_model.py +++ b/src/rossmann_ops/train_model.py @@ -208,8 +208,16 @@ def train_production_model() -> None: plt.savefig(shap_tmp) if False else shutil.copy(shap_out_path, shap_tmp) mlflow.log_artifact(shap_tmp) - # 13. Log Model to MLflow Registry - mlflow.sklearn.log_model(model, artifact_path="production_model") + # 13. Log Model to MLflow Registry (wrapped in try/except with fast timeout for network resilience) + os.environ["MLFLOW_HTTP_REQUEST_MAX_RETRIES"] = "1" + try: + mlflow.sklearn.log_model(model, artifact_path="production_model") + except Exception as e: + logger.warning( + "Failed to upload model to MLflow remote Registry due to network issue (%s). " + "Falling back to local DVC artifact.", + e, + ) # 14. Save Model Locally for Docker/CI builds local_model_dir = project_root / config["model"]["save_path"]