Snowpark环境添加learning rate参数致XGBoost训练函数失败

问题概述 #

在Snowflake环境中基于XGBoost执行回归任务时，train_model函数在不添加learning_rate参数时可正常运行，但添加该超参数后，创建永久存储过程时触发报错。

核心代码 #

train_model函数 #

from typing import Tuple
import numpy as np
import os
import snowflake.snowpark.types as T
from sklearn.pipeline import Pipeline
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.compose import ColumnTransformer
from xgboost import XGBRegressor
from sklearn.metrics import mean_squared_error
import joblib

"""
Trains an XGBoost model using the provided Snowflake session, table, features, and target variable.

Args:
    session (snowflake.snowpark.Session): Snowflake session object for connecting to Snowflake.
    table (str): Name of the table in Snowflake containing the training data.
    features (list): List of feature column names to be used for training.
    target_variable (str): Name of the target variable column.
    cat_cols (list): List of categorical column names in the feature set.
    num_cols (list): List of numerical column names in the feature set.

Returns:
    float: Root mean squared error (RMSE) of the trained XGBoost model on the validation set.
"""

def train_model(session: snowflake.snowpark.Session, 
                table: str, 
                features: list, 
                target_variable: str,
                cat_cols: list,
                num_cols: list) -> T.Variant:

    # Load the Snowflake table
    snowdf = session.table(table)

    # Split the data into training and validation sets
    snowdf_train, snowdf_valid = snowdf.random_split([0.75, 0.25], seed=123)

    # Save the train and validation sets in Snowflake
    snowdf_train.write.mode("overwrite").save_as_table("lapse_data_train")
    snowdf_valid.write.mode("overwrite").save_as_table("lapse_data_valid")

    # Prepare the training and validation data
    train_x = snowdf_train[features].to_pandas()
    train_y = snowdf_train.select(target_variable).to_pandas()
    valid_x = snowdf_valid[features].to_pandas()
    valid_y = snowdf_valid.select(target_variable).to_pandas()

    # Define preprocessing pipelines
    num_pipeline = Pipeline([
        ('imputer', SimpleImputer(strategy="median")),
        ('std_scaler', StandardScaler()),
    ])

    preprocessor = ColumnTransformer(
        transformers=[
            ('num', num_pipeline, num_cols),
            ('encoder', OneHotEncoder(handle_unknown="ignore"), cat_cols),
        ]
    )

    # Build pipeline with XGBoost (添加learning_rate参数后触发问题)
    pipe = Pipeline([
        ('preprocessor', preprocessor),
        ('xgboost', XGBRegressor(learning_rate=0.01)),
    ])

    # Train model
    pipe.fit(train_x, train_y)

    # Evaluate on validation set
    valid_preds = pipe.predict(valid_x)
    rmse = mean_squared_error(valid_y, valid_preds, squared=False)
    
    # Save model to stage
    model_file = os.path.join('/tmp', 'model.joblib')
    joblib.dump(pipe, model_file)
    session.file.put(model_file, "@SANDBOX_SGATE", overwrite=True)

    return rmse

创建存储过程代码 #

# 创建永久存储过程并绑定到Snowflake
train_model_sp = F.sproc(train_model, 
                         session=session, 
                         replace=True,
                         is_permanent=True, 
                         name="xgboost_sproc", 
                         stage_location="@SANDBOX_SGATE")

报错信息 #

ProgrammingError: 091003 (22000): Failure using stage area. Cause: [SANDBOX_SGATE GET and PUT commands are not supported with external stage]

已完成的排查 #

• 确认XGBoost的learning_rate参数符合官方文档定义，参数本身合法
• 验证Snowpark Python Connector、scikit-learn、XGBoost等依赖库已正确安装并适配版本
• 尝试多种方式传递learning_rate参数（如默认值、变量传入），问题依旧

解决方案 #

1. 切换为Snowflake内部阶段（核心修复） #

报错本质原因是外部阶段不支持PUT/GET命令，而创建永久存储过程时，Snowflake需要将代码包上传到指定阶段，此时外部阶段的限制会触发报错。

• 创建内部阶段（若不存在）：

CREATE OR REPLACE STAGE SANDBOX_SGATE_INTERNAL
COMMENT = 'Internal stage for storing permanent stored procedures';

• 修改存储过程创建代码，指向内部阶段：

train_model_sp = F.sproc(train_model, 
                         session=session, 
                         replace=True,
                         is_permanent=True, 
                         name="xgboost_sproc", 
                         stage_location="@SANDBOX_SGATE_INTERNAL")

2. 同步修复函数内的模型存储阶段 #

函数内使用session.file.put上传模型文件到@SANDBOX_SGATE，同样需要确保该阶段为内部阶段。若必须使用外部阶段存储模型，需改用对应云存储的SDK直接上传，而非Snowpark的file.put方法。

3. 确认权限配置 #

确保执行存储过程创建的用户拥有内部阶段的USAGE和WRITE权限：

GRANT USAGE, WRITE ON STAGE SANDBOX_SGATE_INTERNAL TO YOUR_USER;

关键说明 #

• 不添加learning_rate时函数能运行，是因为此时函数作为本地代码执行，未触发存储过程的阶段上传逻辑；创建永久存储过程时才会触发代码包上传，外部阶段的限制才会暴露。
• XGBoost的learning_rate参数本身无问题，报错是阶段类型不兼容导致的，与参数无关。

内容的提问来源于stack exchange，提问作者Rafid Sarker