bump timeouts for mlflow

zenml-io · Oct 12, 2023 · 23f9358 · 23f9358
1 parent 1377a8d
commit 23f9358
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 1 deletion.
diff --git a/template/pipelines/batch_inference.py b/template/pipelines/batch_inference.py
@@ -67,6 +67,7 @@ def {{product_name}}_batch_inference():
         registry_model_name=get_pipeline_context().extra["mlflow_model_name"],
         registry_model_version=registry_model_version,
         replace_existing=False,
+        timeout=180,
     )
     inference_predict(
         deployment_service=deployment_service,

diff --git a/template/pipelines/training.py b/template/pipelines/training.py
@@ -150,6 +150,7 @@ def {{product_name}}_training(
         registry_model_name=pipeline_extra["mlflow_model_name"],
         registry_model_version=latest_version,
         replace_existing=False,
+        timeout=180,
     )
     latest_metric = promote_get_metric(
         id="get_metrics_latest_model_version",
@@ -162,6 +163,7 @@ def {{product_name}}_training(
         registry_model_name=pipeline_extra["mlflow_model_name"],
         registry_model_version=current_version,
         replace_existing=False,
+        timeout=180,
         after=["get_metrics_latest_model_version"],
     )
     current_metric = promote_get_metric(

diff --git a/tests/test_template.py b/tests/test_template.py
@@ -92,7 +92,7 @@ def generate_and_run_project(
         except subprocess.CalledProcessError as e:
             raise RuntimeError(
                 f"Failed to run project generated with parameters: {answers}\n"
-                f"{e.output}"
+                f"{e.output.decode()}"
             ) from e
 
         # check the pipeline run is successful