Snowflake ML Jobs¶

Run a Python function as Snowflake ML Job¶

Use Function Dispatch to run individual Python functions remotely on Snowflake’s compute resources with the @remote decorator.

Using the @remote decorator, you can do the following:

• Serializate the function and its dependencies.

• Upload it to a specified Snowflake stage.

• Execute it within the Container Runtime for ML.

The following example Python code uses the @remote decorator to submit a Snowflake ML Job. Note that a Snowpark Session is required, see Prerequisites.

from snowflake.ml.jobs import remote

@remote("MY_COMPUTE_POOL", stage_name="payload_stage", session=session)
def train_model(data_table: str):
  # Provide your ML code here, including imports and function calls
  ...

job = train_model("my_training_data")

Invoking a @remote decorated function returns a Snowflake MLJob object that can be used to manage and monitor the job execution. For more information, see Snowflake ML Job management.

Run a Python file as a Snowflake ML Job¶

Run Python files or project directories on Snowflake compute resources. This is useful when:

• You have complex ML projects with multiple modules and dependencies.

• You want to maintain separation between local development and production code.

• You need to run scripts that use command-line arguments.

• You’re working with existing ML projects that weren’t specifically designed for execution on Snowflake compute.

The Snowflake Job API offers two main methods:

• submit_file(): For running single Python files

• submit_directory(): For running entire project directories with multiple files and resources

Both methods support:

• Command-line argument passing

• Environment variable configuration

• Custom dependency specification

• Project asset management through Snowflake stages

File Dispatch is particularly useful for productionizing existing ML workflows and maintaining clear separation between development and execution environments.

The following Python code submits a file to a Snowflake ML Job:

from snowflake.ml.jobs import submit_file

# Run a single file
job1 = submit_file(
  "train.py",
  "MY_COMPUTE_POOL",
  stage_name="payload_stage",
  args=["--data-table", "my_training_data"],
  session=session,
)

The following Python code submits a directory to a Snowflake ML Job:

from snowflake.ml.jobs import submit_directory

# Run from a directory
job2 = submit_directory(
  "./ml_project/",
  "MY_COMPUTE_POOL",
  entrypoint="train.py",
  stage_name="payload_stage",
  session=session,
)

Submitting a file or directory returns a Snowflake MLJob object that can be used to manage and monitor the job execution. For more information, see Snowflake ML Job management.

Private package feeds¶

Snowflake ML Jobs also support loading packages from private feeds such as JFrog Artifactory and Sonatype Nexus Repository. These feeds are commonly used to distribute internal and proprietary packages, maintain control over dependency versions, and ensure security/compliance.

To install packages from a private feed, you must do the following:

1. Create a Network Rule to allow access to the private feed’s URL.

   1. For sources which use basic authentication, you can simply create a network rule.

      CREATE OR REPLACE NETWORK RULE private_feed_nr
      MODE = EGRESS
      TYPE = HOST_PORT
      VALUE_LIST = ('<your-repo>.jfrog.io');
      

      Copy

   2. To configure access to a source using private connectivity (i.e. Private Link), follow the steps in Network egress using private connectivity.

2. Create an External Access Integration using the network rule. Grant permission to use the EAI to the role that will be submitting jobs.

   CREATE OR REPLACE EXTERNAL ACCESS INTEGRATION private_feed_eai
   ALLOWED_NETWORK_RULES = (PRIVATE_FEED_NR)
   ENABLED = true;
   
   GRANT USAGE ON INTEGRATION private_feed_eai TO ROLE <role_name>;
   

   Copy

3. Specify the private feed URL, External Access Integration, and package(s) when submitting the job

   # Option 1: Specify private feed URL in pip_requirements
   job = submit_file(
     "/path/to/script.py",
     compute_pool="MY_COMPUTE_POOL",
     stage_name="payload_stage",
     pip_requirements=[
       "--index-url=https://your.private.feed.url",
       "internal-package==1.2.3"
     ],
     external_access_integrations=["PRIVATE_FEED_EAI"]
   )
   

   Copy

   # Option 2: Specify private feed URL by environment variable
   job = submit_directory(
     "/path/to/code/",
     compute_pool="MY_COMPUTE_POOL",
     entrypoint="script.py",
     stage_name="payload_stage",
     pip_requirements=["internal-package==1.2.3"],
     external_access_integrations=["PRIVATE_FEED_EAI"],
     env_vars={'PIP_INDEX_URL': 'https://your.private.feed.url'},
   )
   

   Copy

If your private feed URL contains sensitive information like authentication tokens, manage the URL by creating a Snowflake Secret. Use the CREATE SECRET to create a secret. Configure secrets during job submission with the :spec_overrides argument.

# Create secret for private feed URL with embedded auth token
feed_url = "<your-repo>.jfrog.io/artifactory/api/pypi/test-pypi/simple"
user = "<auth_user>"
token = "<auth_token>"
session.sql(f"""
CREATE SECRET IF NOT EXISTS PRIVATE_FEED_URL_SECRET
 TYPE = GENERIC_STRING
 SECRET_STRING = 'https://{auth_user}:{auth_token}@{feed_url}'
""").collect()

# Prepare service spec override for mounting secret into job execution
spec_overrides = {
 "spec": {
  "containers": [
    {
     "name": "main",  # Primary container name is always "main"
     "secrets": [
      {
        "snowflakeSecret": "PRIVATE_FEED_URL_SECRET",
        "envVarName": "PIP_INDEX_URL",
        "secretKeyRef": "secret_string"
      },
     ],
    }
  ]
 }
}

# Load private feed URL from secret (e.g. if URL includes auth token)
job = submit_file(
  "/path/to/script.py",
  compute_pool="MY_COMPUTE_POOL",
  stage_name="payload_stage",
  pip_requirements=[
    "internal-package==1.2.3"
  ],
  external_access_integrations=["PRIVATE_FEED_EAI"],
  spec_overrides=spec_overrides,
)

For more information about the container.secrets, see containers.secrets field.