Fixing stage name

getindata · Jun 7, 2023 · 6daf8d2 · 6daf8d2
1 parent 83dcd45
commit 6daf8d2
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 3 deletions.
diff --git a/kedro_snowflake/config.py b/kedro_snowflake/config.py
@@ -92,7 +92,7 @@ class SnowflakeMLflowConfig(BaseModel):
     experiment_name: Optional[str]
     functions: MLflowFunctionsConfig
     run_id: Optional[str]
-    stage_name: Optional[str]
+    stage: Optional[str]
 
 
 class SnowflakeConfig(BaseModel):
@@ -169,14 +169,15 @@ class KedroSnowflakeConfig(BaseModel):
       - more-itertools
       - openpyxl
       - backoff
+      - pydantic
     # Optionally provide mapping for user-friendly pipeline names
     pipeline_name_mapping:
      __default__: default
   # EXPERIMENTAL: Either MLflow experiment name to enable MLflow tracking
   # or leave empty
   mlflow:
     experiment_name: ~
-    stage_name: ~
+    stage: ~
     # Snowflake external functions needed for calling MLflow instance
     functions:
         experiment_get_by_name: mlflow_experiment_get_by_name

diff --git a/kedro_snowflake/starters/snowflights/{{ cookiecutter.repo_name }}/conf/base/snowflake.yml b/kedro_snowflake/starters/snowflights/{{ cookiecutter.repo_name }}/conf/base/snowflake.yml
@@ -68,6 +68,7 @@ snowflake:
   # or leave empty
   mlflow:
     experiment_name: ~
+    stage: "MLFLOW_STAGE"
     # Snowflake external functions needed for calling MLflow instance
     functions:
       experiment_get_by_name: mlflow_experiment_get_by_name

diff --git a/...okiecutter.repo_name }}/src/{{ cookiecutter.python_package }}/pipelines/mlflow_helpers.py b/...okiecutter.repo_name }}/src/{{ cookiecutter.python_package }}/pipelines/mlflow_helpers.py
@@ -68,6 +68,6 @@ def log_model(model: RandomForestRegressor, X_train: pd.DataFrame, y_train: pd.S
         model.fit(X_train, y_train)
         session = _get_current_session()
         session.file.put(f"/tmp/mlruns/{local_exp_id}/{local_run_id}/artifacts/model/*",
-                         f"@{mlflow_config.stage_name}/{exp_id}/{run_id}/artifacts/model/",
+                         f"{mlflow_config.stage}/{exp_id}/{run_id}/artifacts/model/",
                          auto_compress=False)
         return model