Support additional fanout logging to Snowflake Experiments

docs/guides/operations/sf-experiment-tracking.md

@@ -0,0 +1,36 @@
+# Snowflake Experiment Tracking
+
+Training recipes can log metrics and parameters to Snowflake's native experiment
+tracking. Results are viewable in Snowsight under **AI & ML > Experiments**.
+
+## Prerequisites
+
+Install `snowflake-ml-python` (>= 1.19.0):
+
+```bash
+uv pip install "snowflake-ml-python>=1.19.0"
+```
+
+## Usage
+
+Set `sf_tracking=True` on the train command:
+
+```bash
+python -m recipes.sft.conversational.train \
+  config=/path/to/config.json sf_tracking=True
+```
+
+## How it works
+
+When a recipe creates a Cortex Training job, the server associates it with a
+Snowflake experiment and run. The recipe retrieves the experiment and run names,
+opens a Snowpark session using the same client credentials, and logs training
+hyperparameters and per-step metrics.
+
+## Viewing results
+
+When a run ends, URLs for the experiment and run are printed to stdout. You can
+also browse experiments in Snowsight: **AI & ML > Experiments**.
+
+Metrics are also logged locally (and to Weights & Biases if `wandb_project` is
+set), so all backends receive the same data.

recipes/logging.py

@@ -0,0 +1,87 @@
+# Copyright 2025 Snowflake Inc.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unified training-metrics logging for Cortex Training recipes."""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+class SnowflakeExperimentLogger:
+    """Logs metrics to Snowflake experiment tracking.
+
+    Same ``log_metrics`` / ``close`` interface as ``ml_log`` so it can be
+    composed via :class:`_CompositeLogger`.
+    """
+
+    def __init__(self, session: Any, experiment_name: str, run_name: str) -> None:
+        from snowflake.ml.experiment import ExperimentTracking
+
+        self._exp = ExperimentTracking(session=session)
+        self._exp.set_experiment(experiment_name)
+        self._exp.start_run(run_name)
+
+    def log_params(self, params: dict[str, Any]) -> None:
+        self._exp.log_params(params)
+
+    def log_metrics(self, metrics: dict[str, float], step: int = 0, **kwargs: Any) -> None:
+        self._exp.log_metrics(metrics, step=step)
+
+    def close(self) -> None:
+        pass
+
+
+class _CompositeLogger:
+    """Fans out ``log_metrics`` / ``close`` to multiple loggers."""
+
+    def __init__(self, loggers: list[Any]) -> None:
+        self._loggers = loggers
+
+    def log_metrics(self, metrics: dict[str, float], step: int = 0, **kwargs: Any) -> None:
+        for lg in self._loggers:
+            lg.log_metrics(metrics=metrics, step=step, **kwargs)
+
+    def close(self) -> None:
+        for lg in self._loggers:
+            lg.close()
+
+
+def setup_logging(
+    config: Any,
+    *,
+    client: Any | None = None,
+    job_id: str | None = None,
+) -> Any:
+    """Create a training-metrics logger, optionally with Snowflake experiment tracking."""
+    from tinker_cookbook.utils import ml_log
+
+    ml_logger = ml_log.setup_logging(
+        log_dir=config.log_path,
+        wandb_project=config.wandb_project,
+        wandb_name=config.wandb_name,
+        config=config,
+        do_configure_logging_module=True,
+    )
+    if not getattr(config, "sf_tracking", False) or client is None or job_id is None:
+        return ml_logger
+    run_info = client.get_experiment_run(job_id)
+    sf_logger = SnowflakeExperimentLogger(
+        client.create_snowpark_session(),
+        run_info["experiment_name"],
+        run_info["experiment_run_name"],
+    )
+    sf_logger.log_params(vars(config))
+    return _CompositeLogger([ml_logger, sf_logger])

recipes/rl/math_grpo/train.py

@@ -52,7 +52,7 @@
 from recipes.utils import sequence_from_rollout
 from recipes.utils import stop_params_for
 from recipes.utils import sync_weights
-from tinker_cookbook.utils import ml_log
+from recipes.logging import setup_logging
 
 from cortex_training.client import DEBUG_OPTIONS_ENV
 
@@ -237,6 +237,7 @@ class Config:
     log_path: str = "/tmp/cortex-training-examples/rl-loop"
     wandb_project: str | None = None
     wandb_name: str | None = None
+    sf_tracking: bool = False
 
     # Loaded as the colocated sampling + training create-job body.
     job_config: str = "configs/qwen3_8b_lora.json"
@@ -270,21 +271,12 @@ def main(config: Config):
         os.environ[DEBUG_OPTIONS_ENV] = "1"
         logger.info("Using debug image_tag=%s", config.debug_image_tag)
 
-    ml_logger = ml_log.setup_logging(
-        log_dir=config.log_path,
-        wandb_project=config.wandb_project,
-        wandb_name=config.wandb_name,
-        config=config,
-        do_configure_logging_module=True,
-    )
-
-    _train(config, ml_logger)
+    _train(config)
 
-    ml_logger.close()
     logger.info("Training completed")
 
 
-def _train(config: Config, ml_logger: Any) -> None:
+def _train(config: Config) -> None:
     body = job_body(config)
     subs = {sub.get("job_type"): sub for sub in body.get("sub_job_configs") or ()}
     training_sub = subs.get("training") or {}
@@ -362,6 +354,7 @@ def _train(config: Config, ml_logger: Any) -> None:
     client = make_client(config.config)
 
     with running_job(client, body, job_id=config.job_id) as job_id:
+        ml_logger = setup_logging(config, client=client, job_id=job_id)
         sampling_job_id: str | None = None
         if router_replay:
             logger.info("Bootstrapping router replay for job %s", job_id)
@@ -516,6 +509,8 @@ def _train(config: Config, ml_logger: Any) -> None:
             max_examples=config.n_test,
         )
 
+    ml_logger.close()
+
 
 if __name__ == "__main__":
     chz.nested_entrypoint(main)

recipes/sft/conversational/train.py

@@ -44,8 +44,8 @@
 from recipes.utils import save_recipe_checkpoints
 from recipes.utils import sequence_from_conversation
 from recipes.utils import use_next_token_labels
+from recipes.logging import setup_logging
 from tinker_cookbook import renderers
-from tinker_cookbook.utils import ml_log
 
 from cortex_training.client import DEBUG_OPTIONS_ENV
 
@@ -85,6 +85,7 @@ class Config:
     log_path: str = "/tmp/cortex-training-examples/sft-loop"
     wandb_project: str | None = None
     wandb_name: str | None = None
+    sf_tracking: bool = False
 
     # Loaded as the training-only create-job body.
     job_config: str = "configs/qwen3_8b_full.json"
@@ -175,14 +176,6 @@ def main(config: Config):
     model_name = training_sub.get("model_name")
     chunked_logprob_loss = _uses_chunked_logprob_loss(training)
 
-    ml_logger = ml_log.setup_logging(
-        log_dir=config.log_path,
-        wandb_project=config.wandb_project,
-        wandb_name=config.wandb_name,
-        config=config,
-        do_configure_logging_module=True,
-    )
-
     tokenizer, renderer, renderer_name = build_renderer(
         model_name,
         renderer_name=config.renderer_name,
@@ -215,6 +208,7 @@ def main(config: Config):
     client = make_client(config.config)
 
     with running_job(client, body, job_id=config.job_id) as job_id:
+        ml_logger = setup_logging(config, client=client, job_id=job_id)
         for step in range(total_steps):
             start_time = time.time()
             metrics: dict[str, float] = {}

src/cortex_training/client.py

@@ -1595,23 +1595,18 @@ def _experiment_run_uri(self, job_id: str) -> tuple[str, str]:
             run_name,
         )
 
-    def _open_experiment_artifact_connection(self) -> Any:
-        """Open a connector session for Snowflake experiment artifact LIST/GET."""
+    def _snowflake_connection_kwargs(self) -> dict[str, Any]:
+        """Return connection kwargs derived from this client's PAT credentials."""
         config = self._artifact_connection_config
         if config is None:
-            raise RuntimeError(
-                "experiment artifact download requires a PAT-authenticated client"
-            )
+            raise RuntimeError("requires a PAT-authenticated client")
         user, account, role = self._query_sql_row(
             "SELECT CURRENT_USER(), CURRENT_ACCOUNT_NAME(), CURRENT_ROLE()"
         )
         if not isinstance(user, str) or not user:
             raise ValueError("SQL identity response missing current user")
         if not isinstance(account, str) or not account:
             raise ValueError("SQL identity response missing current account")
-
-        import snowflake.connector
-
         kwargs: dict[str, Any] = {
             "host": config["host"],
             "account": account,
@@ -1623,7 +1618,19 @@ def _open_experiment_artifact_connection(self) -> Any:
         }
         if isinstance(role, str) and role:
             kwargs["role"] = role
-        return snowflake.connector.connect(**kwargs)
+        return kwargs
+
+    def _open_experiment_artifact_connection(self) -> Any:
+        """Open a connector session for Snowflake experiment artifact LIST/GET."""
+        import snowflake.connector
+
+        return snowflake.connector.connect(**self._snowflake_connection_kwargs())
+
+    def create_snowpark_session(self) -> Any:
+        """Create a Snowpark ``Session`` using this client's PAT credentials."""
+        from snowflake.snowpark import Session
+
+        return Session.builder.configs(self._snowflake_connection_kwargs()).create()
 
     @staticmethod
     def _list_experiment_artifacts(