Skip to content

Support additional fanout logging to Snowflake Experiments - #31

Open
sfc-gh-yuzhao wants to merge 13 commits into
mainfrom
yuzhao-snowflake-experiment-tracking
Open

Support additional fanout logging to Snowflake Experiments#31
sfc-gh-yuzhao wants to merge 13 commits into
mainfrom
yuzhao-snowflake-experiment-tracking

Conversation

@sfc-gh-yuzhao

@sfc-gh-yuzhao sfc-gh-yuzhao commented Sep 16, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Summary

  • Add optional Snowflake experiment tracking to the SFT and Math GRPO training recipes
  • Reuse the CortexTrainingClient session for Snowpark — no separate login
  • Experiment/run names are tied to the Cortex Training job (server-assigned)
  • New recipes/logging.py consolidates all logging setup (local, W&B, SF) behind a single setup_logging() call

Changes

  • src/cortex_training/client.py: Extract _snowflake_connection_kwargs() to DRY up credential resolution; add create_snowpark_session() public method
  • recipes/logging.py (new): SnowflakeExperimentLogger, _CompositeLogger, and setup_logging(config, client, job_id) that handles tinker_cookbook local/W&B + optional SF tracking
  • recipes/sft/conversational/train.py: Add sf_tracking config field, replace two-step logging setup with single setup_logging() call
  • recipes/rl/math_grpo/train.py: Same as SFT
  • docs/guides/operations/sf-experiment-tracking.md: New doc page

Usage

python -m recipes.sft.conversational.train \
  config=config.json sf_tracking=True

W&B still works alongside — set wandb_project and sf_tracking=True together for both.

Test plan

  • Smoke tested against a live Snowflake environment (QA6) — session creation, param logging, per-step metric logging, and end_run() all verified

sfc-gh-yuzhao and others added 5 commits September 16, 2026 15:51
@sfc-gh-yuzhao @claude
Add Snowflake experiment tracking support to training recipes
2c53528 
Replace the direct tinker_cookbook ml_log dependency with a unified
setup_logging() factory in recipes/utils.py. When sf_experiment is set,
metrics and parameters are logged to Snowflake's native experiment
tracking instead of Weights & Biases. The SnowflakeExperimentLogger
class exposes the same log_metrics/close interface so training loops
are unchanged.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sfc-gh-yuzhao @claude
Support logging to both Snowflake and W&B simultaneously
570cf98 
Add _CompositeLogger to fan out log_metrics/close to multiple backends
when both sf_experiment and wandb_project are set.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sfc-gh-yuzhao @claude
Fix Snowpark session to pass user field, use valid SQL identifiers in…
5146c1a 
… docs

Testing revealed two issues:
- _create_snowpark_session was missing the user field, causing auth failures
- Doc examples used hyphens in experiment names which are invalid SQL identifiers

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sfc-gh-yuzhao @claude
Clean up Snowpark session config and param logging
65061ec 
- Pass user field inline instead of conditional append
- Drop unnecessary _ prefix filter in _log_config_params
- List user in doc prerequisites

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sfc-gh-yuzhao @claude
Remove redundant validation in _create_snowpark_session
897e8fe 
KeyError from missing host/pat is clear enough; Snowpark session
creation also fails with its own error on bad config.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sfc-gh-halu

sfc-gh-halu commented Sep 17, 2026

Copy link
Copy Markdown
Collaborator

thanks for doing that.
some comments:

  1. for session, let's use same session from CortexTrainingClient, could expose a public method to retrieve same session so we do not need extra login code.
  2. for experiment_name , each job associated with snowflake experiment/run, so we should explicit set it when creating job, or calling get_experiment_run to get experimetn and run info.

sfc-gh-yuzhao and others added 6 commits September 17, 2026 11:06
@sfc-gh-yuzhao @claude
Tie experiment tracking to job and reuse client session
143ba6e 
Address PR feedback:
- Add create_snowpark_session() to CortexTrainingClient so recipes
  reuse the client's PAT credentials instead of separate login code
- Pass experiment_name into the job body via running_job() so the
  experiment is tied to the Cortex Training job
- Retrieve server-assigned run name via get_experiment_run(job_id)
- Remove sf_run_name config field (server owns run naming)
- Delete _create_snowpark_session() and setup_logging() factory

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sfc-gh-yuzhao @claude
Extract shared connection kwargs in CortexTrainingClient
b81f266 
DRY up _open_experiment_artifact_connection and create_snowpark_session
by factoring credential resolution into _snowflake_connection_kwargs.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sfc-gh-yuzhao @claude
Always log to Snowflake experiment tracking
baa2f6a 
Remove sf_experiment config field and sf_tracking conditional — every
training run now logs to Snowflake automatically. The server assigns
experiment/run names per job, so no user configuration is needed.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sfc-gh-yuzhao @claude
Extract setup_sf_logging helper to reduce recipe duplication
dc4a0bd 
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sfc-gh-yuzhao @claude
Restore ml_log import to module level
d38dcae 
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sfc-gh-yuzhao @claude
Remove stray blank lines in train.py files
2c99012 
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sfc-gh-halu

sfc-gh-halu commented Sep 17, 2026

Copy link
Copy Markdown
Collaborator

Also, could we update the README at recipe folder level about snowflake experiment tracking integration, it's easier for people to figure out there

sfc-gh-srajbhandari
sfc-gh-srajbhandari reviewed Sep 18, 2026
View reviewed changes

@sfc-gh-srajbhandari sfc-gh-srajbhandari left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Lets not force SF logging. Please make it optional. Lot of folks are going to want to use wandb.

sfc-gh-yuzhao and others added 2 commits September 18, 2026 12:07
@sfc-gh-yuzhao @claude
Add sf_tracking toggle and remove end_run from close
127a00a 
SF experiment logging is now opt-in via sf_tracking=True. The logger
no longer calls end_run on close — the server manages run lifecycle.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sfc-gh-yuzhao @claude
Extract recipes/logging.py for unified training-metrics logging
d6bb004 
Move Snowflake experiment tracking classes and setup out of utils.py
into a new logging module. Consolidate the two-step ml_log + SF setup
into a single setup_logging(config, client, job_id) call.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@sfc-gh-halu

sfc-gh-halu commented Sep 18, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

could we also update /recipe/README.md to flag both options available for logging metrics? especially given snowml package installation is needed, feel free outline and then link to reference page you created

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

1 more reviewer
@sfc-gh-srajbhandari sfc-gh-srajbhandari sfc-gh-srajbhandari left review comments
Reviewers whose approvals may not affect merge requirements

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@sfc-gh-yuzhao @sfc-gh-halu @sfc-gh-srajbhandari