Distribute ALS Recommendation Engine Training With Ease Using Snowflake HPO API

Snowflake HPO API

Alternating Least Squares (ALS) remains a go-to collaborative filtering algorithm, but scaling it efficiently has always been a challenge. This blog combines the popular **implicit library with Snowflake's Hyperparameter Optimization (HPO) API and its multi-node distributed compute** capabilities to run ALS faster — without leaving Snowflake.

📁Jump straight to Code? *Here it is!*

We use the open-source **MovieLens 1 million dataset**, loaded directly into a Snowflake table. The oldest 80% of data is used as training data (800,168 ratings) and newest 20% as test (200,041 ratings). Here’s a quick look at the sample of data with user_id, item_id, rating given by the user and timestamp.

Movie Lens Sample Data

1. Data Preparation and Index Mapping

ALS requires a sparse user-item matrix with 0-based integer indices rather than raw user/item IDs. To ensure compatibility with real-world datasets where IDs may be non-contiguous or strings, we implement user_to_idx and item_to_idx mapping dictionaries. These mappings are then embedded directly into the DataFrames as USER_IDX and PRODUCT_IDX columns, so each HPO trial reads them instantly rather than recomputing them in every trial.

# --- Precompute index mappings (optimization) ---
all_users = pd.concat([train_df['USER'], test_df['USER']]).unique()
all_items = pd.concat([train_df['PRODUCT'], test_df['PRODUCT']]).unique()
user_to_idx = {user: idx for idx, user in enumerate(sorted(all_users))}
item_to_idx = {item: idx for idx, item in enumerate(sorted(all_items))}

# Embed the precomputed indices directly into the DataFrames
train_df['USER_IDX'] = train_df['USER'].map(user_to_idx)
train_df['PRODUCT_IDX'] = train_df['PRODUCT'].map(item_to_idx)
test_df['USER_IDX'] = test_df['USER'].map(user_to_idx)
test_df['PRODUCT_IDX'] = test_df['PRODUCT'].map(item_to_idx)

DataConnector is how the Snowflake Tuner distributes your data across nodes during hyperparameter trials. Each trial gets the same preprocessed data without re-running the split or index mapping logic.

dataset_map = {
    "train": DataConnector.from_dataframe(session.create_dataframe(train_data)),
    "test":  DataConnector.from_dataframe(session.create_dataframe(test_data)),
}

2. Model Training

Let’s define the training function Snowflake’s HPO Tuner calls for every hyperparameter trial. Each trial receives its own unique hyperparameter config and the shared dataset_map from create_data_connectors() — no data re-splitting per trial.

The precomputed USER_IDX, PRODUCT_IDX, N_USERS, and N_ITEMS columns are read directly, skipping redundant ID→index mapping on every trial. These are then used to build a CSR sparse matrix in users × items format — exactly what the implicit ALS library expects.

def train_func():

    tuner_context = get_tuner_context()
    config = tuner_context.get_hyper_params()
    dm = tuner_context.get_dataset_map()

    train_pdf = dm["train"].to_pandas()
    test_pdf = dm["test"].to_pandas()

    ...

    # Build user-item sparse matrix (users x items)
    user_item_train = csr_matrix(
        (train_ratings, (train_user_idx, train_item_idx)),
        shape=(n_users, n_items)
    )
    ...

    factors = config['RANK']
    regularization = config['REGPARAM']
    iterations = config['MAXITER']
    alpha = config['ALPHA']

    als_model = implicit.als.AlternatingLeastSquares(
        factors=factors,
        regularization=regularization,
        iterations=iterations,
        random_state=42
    )

    confidence_matrix = (user_item_train * alpha).astype('double')
    als_model.fit(confidence_matrix)
    ...

    tuner_context.report(
        metrics={"precision_at_10": p_at_k, "map_at_10": map_at_k, "ndcg_at_10": ndcg_at_k},
        model=als_model
    )

TunerConfig — the HPO API’s Brain

Snowflake’s Tuner takes the dataset_map (your preprocessed DataConnector objects), distributes data to each trial, calls train_func 81 times with different hyperparameter configs as defined below, and tracks all results.

tuner_single = tune.Tuner(
    train_func=train_func,
    search_space={
        "RANK": [20, 50, 100],
        "REGPARAM": [0.01, 0.05, 0.1],
        "MAXITER": [10, 20, 30],
        "ALPHA": [1.0, 10.0, 40.0]
    },
    tuner_config=tune.TunerConfig(
        metric="map_at_10",  # what to optimize
        mode="max",         # maximize it (higher accuracy = better)
        search_alg=GridSearch(),
        num_trials=81,
        max_concurrent_trials=1, # sequential (single-node baseline)
    ),
)

start_single = time.time()
tuner_results_single = tuner_single.run(dataset_map=dataset_map)

Scaling Across Multiple Nodes

We can increase the number of active nodes in a compute cluster and the parallelised HPO API will handle all the complexities of distributing the training across multiple resources.

#Trails with parallel execution across multi-nodes
target_nodes = 3
scale_cluster(expected_cluster_size=target_nodes) #Scale nodes 
print(f"81 trials, max_concurrent_trials=8 (parallel across {actual_nodes} nodes)")

tuner_multi = tune.Tuner(
    train_func=train_func,
    search_space={
        "RANK": [20, 50, 100],
        "REGPARAM": [0.01, 0.05, 0.1],
        "MAXITER": [10, 20, 30],
        "ALPHA": [1.0, 10.0, 40.0]
    },
    tuner_config=tune.TunerConfig(
        metric="map_at_10",
        mode="max",
        search_alg=GridSearch(),
        num_trials=81,
        max_concurrent_trials=8, #Increase the concurrent trials 
    ),
)

And all these experiments are tracked under the experiment tracking feature in Snowflake.

Experiment Tracking In Snowflake

3. Inference

With the best trial identified, we register the trained ALS model in the Snowflake Model Registry as a bring your own custom model (BYOM). The custom model class encapsulates the ALS model alongside the ID-to-index mappings so that inference is fully self-contained — callers pass real user IDs and receive real item IDs back.

class ALSRecommendationModel(custom_model.CustomModel):
    """CustomModel that accepts real user IDs and returns real item IDs.
    Embeds the ALS model + ID mappings so inference is fully self-contained.
    """
    def __init__(self, context: custom_model.ModelContext) -> None:
        super().__init__(context)
        with open(self.context['als_model_path'], 'rb') as f:
            self.als_model = pickle.load(f)
        with open(self.context['mappings_path'], 'rb') as f:
            self.mappings = pickle.load(f)
        self.user_to_idx = self.mappings['user_to_idx']
        self.idx_to_item = self.mappings['idx_to_item']

    @custom_model.inference_api
    def recommend(self, input_df: pd.DataFrame) -> pd.DataFrame:
        results = []
        for _, row in input_df.iterrows():
            user_id = int(row['USER_ID'])
            user_idx = self.user_to_idx.get(user_id)

            if user_idx is None:
                results.append({
                    'USER_ID': user_id,
                    'RECOMMENDED_ITEMS': 'UNKNOWN_USER',
                    'SCORES': ''
                })
                continue

            item_indices, scores = self.als_model.recommend(
                user_idx, None, N=10, filter_already_liked_items=False
            )
            real_item_ids = [str(self.idx_to_item.get(int(idx), f'UNK_{idx}')) for idx in item_indices]
            results.append({
                'USER_ID': user_id,
                'RECOMMENDED_ITEMS': ','.join(real_item_ids),
                'SCORES': ','.join(f'{s:.4f}' for s in scores)
            })
        return pd.DataFrame(results)

model_context = custom_model.ModelContext(
    als_model_path=als_model_path,
    mappings_path=mappings_path
)

Once registered, you can call the model directly from the registry to generate recommendations at scale.

#Load the model from the registry
loaded_model = mv.load()

#Use real user IDs from the MovieLens 1M dataset
sample_users = pd.DataFrame({'USER_ID': [2000, 3500, 5000, 1500, 4200]})
#Call the recommendation custom function to get the recommendations
recommendations = loaded_model.recommend(sample_users)

In this post, we walked through how to scale ALS hyperparameter search with Snowflake’s HPO API — from preprocessing and index mapping, through distributed multi-node trials, to registering and serving the best model — all within a single Snowflake environment. The same pattern generalises to any training function you bring, making it straightforward to benchmark and productionalize recommendation models without managing external infrastructure.

📁Code? *Here it is!*