From 9c672bb2e1428c406dc7ed57719f5ab92c1c4d34 Mon Sep 17 00:00:00 2001 From: Prudhvi Dharmana Date: Wed, 11 Sep 2024 15:33:17 -0700 Subject: [PATCH 1/4] v0 --- tools/snowflake/spcs_dashboard/Dockerfile | 10 ++ .../snowflake/spcs_dashboard/requirements.txt | 7 ++ .../snowflake/spcs_dashboard/run_container.py | 103 ++++++++++++++++++ .../snowflake/spcs_dashboard/run_dashboard.py | 20 ++++ 4 files changed, 140 insertions(+) create mode 100644 tools/snowflake/spcs_dashboard/Dockerfile create mode 100644 tools/snowflake/spcs_dashboard/requirements.txt create mode 100644 tools/snowflake/spcs_dashboard/run_container.py create mode 100644 tools/snowflake/spcs_dashboard/run_dashboard.py diff --git a/tools/snowflake/spcs_dashboard/Dockerfile b/tools/snowflake/spcs_dashboard/Dockerfile new file mode 100644 index 000000000..40d3fcf91 --- /dev/null +++ b/tools/snowflake/spcs_dashboard/Dockerfile @@ -0,0 +1,10 @@ +ARG BASE_IMAGE=python:3.11.9-slim-bullseye +FROM $BASE_IMAGE + +COPY ./ /trulens_dashboard/ + +WORKDIR /trulens_dashboard + +RUN pip install -r requirements.txt + +CMD ["python", "run_dashboard.py"] diff --git a/tools/snowflake/spcs_dashboard/requirements.txt b/tools/snowflake/spcs_dashboard/requirements.txt new file mode 100644 index 000000000..a3b79ff99 --- /dev/null +++ b/tools/snowflake/spcs_dashboard/requirements.txt @@ -0,0 +1,7 @@ +python-dotenv +pydantic +snowflake[ml] +snowflake-connector-python +snowflake-sqlalchemy +trulens +trulens-connectors-snowflake \ No newline at end of file diff --git a/tools/snowflake/spcs_dashboard/run_container.py b/tools/snowflake/spcs_dashboard/run_container.py new file mode 100644 index 000000000..62a1d51a9 --- /dev/null +++ b/tools/snowflake/spcs_dashboard/run_container.py @@ -0,0 +1,103 @@ +import os +from snowflake.snowpark import Session +from snowflake.snowpark.functions import call_builtin + +# Read environment variables +account = os.getenv('SNOWFLAKE_ACCOUNT') +user = os.getenv('SNOWFLAKE_USER') +password = os.getenv('SNOWFLAKE_PASSWORD') +role = os.getenv('SNOWFLAKE_ROLE') +warehouse = os.getenv('SNOWFLAKE_WAREHOUSE') +database = os.getenv('SNOWFLAKE_DATABASE') +schema = os.getenv('SNOWFLAKE_SCHEMA') + +# Define Snowflake connection parameters +connection_parameters = { + "account": account, + "user": user, + "password": password, + "role": role, + "warehouse": warehouse, + "database": database, + "schema": schema +} + +# Create a Snowflake session +session = Session.builder.configs(connection_parameters).create() + +# Create compute pool if it does not exist +compute_pool_name = input("Enter compute pool name: ") +compute_pools = session.sql("SHOW COMPUTE POOLS").collect() +compute_pool_exists = any(pool['name'] == compute_pool_name.upper() for pool in compute_pools) +if compute_pool_exists: + print(f"Compute pool {compute_pool_name} already exists") +else: + session.sql(f"CREATE COMPUTE POOL {compute_pool_name} MIN_NODES = 1 MAX_NODES = 1 INSTANCE_FAMILY = CPU_X64_M").collect() +session.sql(f"DESCRIBE COMPUTE POOL {compute_pool_name}").collect() + +# Create image repository +image_repository_name = f"trulens_image_repository" +session.sql(f"CREATE IMAGE REPOSITORY {image_repository_name}").collect() +session.sql("SHOW IMAGE REPOSITORIES").collect() + +# Create network rule +network_rule_name = f"{compute_pool_name}_allow_all_network_rule" +session.sql(f"CREATE OR REPLACE NETWORK RULE {network_rule_name} TYPE = 'HOST_PORT' MODE = 'EGRESS' VALUE_LIST = ('0.0.0.0:443','0.0.0.0:80')").collect() +session.sql("SHOW NETWORK RULES").collect() + +# Create external access integration +access_integration_name = f"{compute_pool_name}_access_integration" +session.sql(f"CREATE OR REPLACE EXTERNAL ACCESS INTEGRATION {access_integration_name} ALLOWED_NETWORK_RULES = ({network_rule_name}) ENABLED = true").collect() +session.sql("SHOW EXTERNAL ACCESS INTEGRATIONS").collect() + +app_name = "trulens_dashboard" +secret_name = f"{schema}.{app_name}_login_credentials" +session.sql(f"CREATE SECRET {secret_name} TYPE=password USERNAME={user} PASSWORD='{password}'").collect() + +service_name=compute_pool_name+"_trulens_dashboard" +session.sql(""" +CREATE SERVICE {service_name} + IN COMPUTE POOL {compute_pool_name} + EXTERNAL_ACCESS_INTEGRATIONS = ({access_integration_name}) + FROM SPECIFICATION $$ + spec: + containers: + - name: {container_name} + image: /{database}/{schema}/{container_name}/{app_name}:latest + env: + SNOWFLAKE_ACCOUNT: "{account}" + SNOWFLAKE_DATABASE: "{database}" + SNOWFLAKE_SCHEMA: "{schema}" + SNOWFLAKE_WAREHOUSE: "{warehouse}" + SNOWFLAKE_ROLE: "{role}" + RUN_DASHBOARD: "1" + secrets: + - snowflakeSecret: {secret_name} + secretKeyRef: username + envVarName: SNOWFLAKE_USER + - snowflakeSecret: {secret_name} + secretKeyRef: password + envVarName: SNOWFLAKE_PASSWORD + endpoints: + - name: trulens-demo-dashboard-endpoint + port: 8484 + public: true + $$ +""".format(service_name=service_name, + compute_pool_name=compute_pool_name, + access_integration_name=access_integration_name, + container_name=app_name+"_container", + account=account, + database=database, + schema=schema, + warehouse=warehouse, + role=role, + app_name=app_name)).collect() + +# Show services and get their status +session.sql(f"SHOW ENDPOINTS IN SERVICE {service_name}").collect() +session.sql("CALL SYSTEM$GET_SERVICE_STATUS('dkurokawa_trulens_demo_app')").collect() +session.sql("CALL SYSTEM$GET_SERVICE_STATUS('dkurokawa_trulens_demo_dashboard')").collect() + +# Close the session +session.close() \ No newline at end of file diff --git a/tools/snowflake/spcs_dashboard/run_dashboard.py b/tools/snowflake/spcs_dashboard/run_dashboard.py new file mode 100644 index 000000000..ef4345795 --- /dev/null +++ b/tools/snowflake/spcs_dashboard/run_dashboard.py @@ -0,0 +1,20 @@ +from trulens.dashboard import run_dashboard +from trulens.core import TruSession +from trulens.connectors.snowflake import SnowflakeConnector +import os + +connection_params = { + "account": os.environ.get("SNOWFLAKE_ACCOUNT"), + "user": os.environ.get("SNOWFLAKE_USER"), + "password": os.environ.get("SNOWFLAKE_PASSWORD"), + "database": os.environ.get("SNOWFLAKE_DATABASE"), + "schema": os.environ.get("SNOWFLAKE_SCHEMA"), + "warehouse": os.environ.get("SNOWFLAKE_WAREHOUSE"), + "role": os.environ.get("SNOWFLAKE_ROLE"), + "init_server_side": False, +} + +connector = SnowflakeConnector(**connection_params) +session = TruSession(connector=connector, init_server_side=False) + +run_dashboard(session, port=8484) \ No newline at end of file From d853b100430106ad051e41c11fd44945a94ac7e4 Mon Sep 17 00:00:00 2001 From: Prudhvi Dharmana Date: Tue, 8 Oct 2024 18:42:34 -0700 Subject: [PATCH 2/4] safe commit --- .../trulens/connectors/snowflake/connector.py | 4 ++ src/dashboard/trulens/dashboard/run.py | 4 +- .../trulens/dashboard/streamlit_utils.py | 51 ++++++++++++- tools/snowflake/spcs_dashboard/Dockerfile | 2 + .../snowflake/spcs_dashboard/requirements.txt | 2 +- .../snowflake/spcs_dashboard/run_dashboard.py | 71 ++++++++++++++++--- 6 files changed, 118 insertions(+), 16 deletions(-) diff --git a/src/connectors/snowflake/trulens/connectors/snowflake/connector.py b/src/connectors/snowflake/trulens/connectors/snowflake/connector.py index 67e9597f7..63658b794 100644 --- a/src/connectors/snowflake/trulens/connectors/snowflake/connector.py +++ b/src/connectors/snowflake/trulens/connectors/snowflake/connector.py @@ -41,6 +41,7 @@ def __init__( schema: str, warehouse: str, role: str, + host: Optional[str] = None, init_server_side: bool = True, database_redact_keys: bool = False, database_prefix: Optional[str] = None, @@ -91,6 +92,7 @@ def __init__( warehouse, role, database_args["database_prefix"], + host, ) def _initialize_snowflake_server_side_feedback_evaluations( @@ -103,6 +105,7 @@ def _initialize_snowflake_server_side_feedback_evaluations( warehouse: str, role: str, database_prefix: str, + host: Optional[str] = None, ): connection_parameters = { "account": account, @@ -112,6 +115,7 @@ def _initialize_snowflake_server_side_feedback_evaluations( "schema": schema, "warehouse": warehouse, "role": role, + **({"host": host} if host else {}) } with Session.builder.configs(connection_parameters).create() as session: ServerSideEvaluationArtifacts( diff --git a/src/dashboard/trulens/dashboard/run.py b/src/dashboard/trulens/dashboard/run.py index 66930f9dc..40e6b5cce 100644 --- a/src/dashboard/trulens/dashboard/run.py +++ b/src/dashboard/trulens/dashboard/run.py @@ -36,6 +36,7 @@ def run_dashboard( address: Optional[str] = None, force: bool = False, _dev: Optional[Path] = None, + spcs_runtime: Optional[bool] = False, ) -> Process: """Run a streamlit dashboard to view logged results and apps. @@ -125,7 +126,8 @@ def run_dashboard( "--database-prefix", session.connector.db.table_prefix, ] - + if spcs_runtime: + args.append("--spcs-runtime") proc = subprocess.Popen( args, stdout=subprocess.PIPE, diff --git a/src/dashboard/trulens/dashboard/streamlit_utils.py b/src/dashboard/trulens/dashboard/streamlit_utils.py index b9adbe718..6cd0cb423 100644 --- a/src/dashboard/trulens/dashboard/streamlit_utils.py +++ b/src/dashboard/trulens/dashboard/streamlit_utils.py @@ -1,8 +1,10 @@ import argparse import sys - +import os from trulens.core import TruSession from trulens.core.database import base as mod_db +from snowflake.snowpark import Session +from snowflake.sqlalchemy import URL def init_from_args(): @@ -16,6 +18,7 @@ def init_from_args(): parser.add_argument( "--database-prefix", default=mod_db.DEFAULT_DATABASE_PREFIX ) + parser.add_argument("--spcs-runtime", default=False) try: args = parser.parse_args() @@ -27,6 +30,48 @@ def init_from_args(): # so we have to do a hard exit. sys.exit(e.code) - TruSession( - database_url=args.database_url, database_prefix=args.database_prefix + def get_login_token(): + """ + Read the login token supplied automatically by Snowflake. These tokens + are short lived and should always be read right before creating any new connection. + """ + with open("/snowflake/session/token", "r") as f: + return f.read() + + connection_params = { + "account": os.environ.get("SNOWFLAKE_ACCOUNT"), + "host": os.getenv("SNOWFLAKE_HOST"), + "authenticator": "oauth", + "token": get_login_token(), + "warehouse": os.environ.get("SNOWFLAKE_WAREHOUSE"), + "database": os.environ.get("SNOWFLAKE_DATABASE"), + "schema": os.environ.get("SNOWFLAKE_SCHEMA"), + } + snowpark_session = Session.builder.configs(connection_params).create() + + # Set up sqlalchemy engine parameters. + conn = snowpark_session.connection + engine_params = {} + engine_params["paramstyle"] = "qmark" + engine_params["creator"] = lambda: conn + database_args = {"engine_params": engine_params} + db_url = URL( + account=snowpark_session.get_current_account(), + user=snowpark_session.get_current_user(), + password="password", + database=snowpark_session.get_current_database(), + schema=snowpark_session.get_current_schema(), + warehouse=snowpark_session.get_current_warehouse(), + role=snowpark_session.get_current_role(), ) + if args.spcs_runtime: + TruSession( + database_url=db_url, + database_check_revision=False, + database_args=database_args, + database_prefix=args.database_prefix + ) + else: + TruSession( + database_url=args.database_url, database_prefix=args.database_prefix + ) diff --git a/tools/snowflake/spcs_dashboard/Dockerfile b/tools/snowflake/spcs_dashboard/Dockerfile index 40d3fcf91..c6a3a1b0e 100644 --- a/tools/snowflake/spcs_dashboard/Dockerfile +++ b/tools/snowflake/spcs_dashboard/Dockerfile @@ -6,5 +6,7 @@ COPY ./ /trulens_dashboard/ WORKDIR /trulens_dashboard RUN pip install -r requirements.txt +RUN pip install trulens_connectors_snowflake-1.0.1-py3-none-any.whl +RUN pip install trulens_dashboard-1.0.1-py3-none-any.whl CMD ["python", "run_dashboard.py"] diff --git a/tools/snowflake/spcs_dashboard/requirements.txt b/tools/snowflake/spcs_dashboard/requirements.txt index a3b79ff99..12eea22b3 100644 --- a/tools/snowflake/spcs_dashboard/requirements.txt +++ b/tools/snowflake/spcs_dashboard/requirements.txt @@ -4,4 +4,4 @@ snowflake[ml] snowflake-connector-python snowflake-sqlalchemy trulens -trulens-connectors-snowflake \ No newline at end of file +# trulens-connectors-snowflake \ No newline at end of file diff --git a/tools/snowflake/spcs_dashboard/run_dashboard.py b/tools/snowflake/spcs_dashboard/run_dashboard.py index ef4345795..4983a405a 100644 --- a/tools/snowflake/spcs_dashboard/run_dashboard.py +++ b/tools/snowflake/spcs_dashboard/run_dashboard.py @@ -1,20 +1,69 @@ from trulens.dashboard import run_dashboard from trulens.core import TruSession +from snowflake.snowpark import Session from trulens.connectors.snowflake import SnowflakeConnector +from snowflake.sqlalchemy import URL import os +# connection_params = { +# "account": os.environ.get("SNOWFLAKE_ACCOUNT"), +# "host": os.getenv("SNOWFLAKE_HOST"), +# "user": os.environ.get("SNOWFLAKE_USER"), +# "password": os.environ.get("SNOWFLAKE_PASSWORD"), +# "database": os.environ.get("SNOWFLAKE_DATABASE"), +# "schema": os.environ.get("SNOWFLAKE_SCHEMA"), +# "warehouse": os.environ.get("SNOWFLAKE_WAREHOUSE"), +# "role": os.environ.get("SNOWFLAKE_ROLE"), +# "init_server_side": False, +# } + +# connector = SnowflakeConnector(**connection_params) +# session = TruSession(connector=connector, init_server_side=False) + +def get_login_token(): + """ + Read the login token supplied automatically by Snowflake. These tokens + are short lived and should always be read right before creating any new connection. + """ + with open("/snowflake/session/token", "r") as f: + return f.read() + connection_params = { - "account": os.environ.get("SNOWFLAKE_ACCOUNT"), - "user": os.environ.get("SNOWFLAKE_USER"), - "password": os.environ.get("SNOWFLAKE_PASSWORD"), - "database": os.environ.get("SNOWFLAKE_DATABASE"), - "schema": os.environ.get("SNOWFLAKE_SCHEMA"), - "warehouse": os.environ.get("SNOWFLAKE_WAREHOUSE"), - "role": os.environ.get("SNOWFLAKE_ROLE"), - "init_server_side": False, + "account": os.environ.get("SNOWFLAKE_ACCOUNT"), + "host": os.getenv("SNOWFLAKE_HOST"), + "authenticator": "oauth", + "token": get_login_token(), + "warehouse": os.environ.get("SNOWFLAKE_WAREHOUSE"), + "database": os.environ.get("SNOWFLAKE_DATABASE"), + "schema": os.environ.get("SNOWFLAKE_SCHEMA"), } +snowpark_session = Session.builder.configs(connection_params).create() -connector = SnowflakeConnector(**connection_params) -session = TruSession(connector=connector, init_server_side=False) +# Set up sqlalchemy engine parameters. +conn = snowpark_session.connection +engine_params = {} +engine_params["paramstyle"] = "qmark" +engine_params["creator"] = lambda: conn +database_args = {"engine_params": engine_params} +# # Ensure any Cortex provider uses the only Snowflake connection allowed in this stored procedure. +# trulens.providers.cortex.provider._SNOWFLAKE_STORED_PROCEDURE_CONNECTION = ( +# conn +# ) +# Run deferred feedback evaluator. +db_url = URL( + account=snowpark_session.get_current_account(), + user=snowpark_session.get_current_user(), + password="password", + database=snowpark_session.get_current_database(), + schema=snowpark_session.get_current_schema(), + warehouse=snowpark_session.get_current_warehouse(), + role=snowpark_session.get_current_role(), +) +tru_session = TruSession( + database_url=db_url, + database_check_revision=False, # TODO: check revision in the future? + database_args=database_args, +) +tru_session.get_records_and_feedback() -run_dashboard(session, port=8484) \ No newline at end of file +run_dashboard(tru_session, port=8484, spcs_runtime=True) \ No newline at end of file From 95c5c6d88adeb47eac3727f6cd97713f800d7ece Mon Sep 17 00:00:00 2001 From: Prudhvi Dharmana Date: Mon, 21 Oct 2024 23:29:54 -0700 Subject: [PATCH 3/4] working version --- .../trulens/connectors/snowflake/connector.py | 2 +- .../snowflake/spcs_dashboard/requirements.txt | 4 +- .../snowflake/spcs_dashboard/run_container.py | 170 +++++++++++------- .../snowflake/spcs_dashboard/run_dashboard.py | 78 +++----- 4 files changed, 131 insertions(+), 123 deletions(-) diff --git a/src/connectors/snowflake/trulens/connectors/snowflake/connector.py b/src/connectors/snowflake/trulens/connectors/snowflake/connector.py index 63658b794..c7ecdb54a 100644 --- a/src/connectors/snowflake/trulens/connectors/snowflake/connector.py +++ b/src/connectors/snowflake/trulens/connectors/snowflake/connector.py @@ -115,7 +115,7 @@ def _initialize_snowflake_server_side_feedback_evaluations( "schema": schema, "warehouse": warehouse, "role": role, - **({"host": host} if host else {}) + **({"host": host} if host else {}), } with Session.builder.configs(connection_parameters).create() as session: ServerSideEvaluationArtifacts( diff --git a/tools/snowflake/spcs_dashboard/requirements.txt b/tools/snowflake/spcs_dashboard/requirements.txt index 12eea22b3..a0b0d392e 100644 --- a/tools/snowflake/spcs_dashboard/requirements.txt +++ b/tools/snowflake/spcs_dashboard/requirements.txt @@ -4,4 +4,6 @@ snowflake[ml] snowflake-connector-python snowflake-sqlalchemy trulens -# trulens-connectors-snowflake \ No newline at end of file +trulens-connectors-snowflake +# trulens-dashboard +# trulens-feedback diff --git a/tools/snowflake/spcs_dashboard/run_container.py b/tools/snowflake/spcs_dashboard/run_container.py index 62a1d51a9..3a424126f 100644 --- a/tools/snowflake/spcs_dashboard/run_container.py +++ b/tools/snowflake/spcs_dashboard/run_container.py @@ -1,69 +1,115 @@ -import os +from argparse import ArgumentParser + from snowflake.snowpark import Session -from snowflake.snowpark.functions import call_builtin - -# Read environment variables -account = os.getenv('SNOWFLAKE_ACCOUNT') -user = os.getenv('SNOWFLAKE_USER') -password = os.getenv('SNOWFLAKE_PASSWORD') -role = os.getenv('SNOWFLAKE_ROLE') -warehouse = os.getenv('SNOWFLAKE_WAREHOUSE') -database = os.getenv('SNOWFLAKE_DATABASE') -schema = os.getenv('SNOWFLAKE_SCHEMA') - -# Define Snowflake connection parameters -connection_parameters = { - "account": account, - "user": user, - "password": password, - "role": role, - "warehouse": warehouse, - "database": database, - "schema": schema -} - -# Create a Snowflake session -session = Session.builder.configs(connection_parameters).create() + +# get args from command line +parser = ArgumentParser(description="Run container script") +parser.add_argument( + "--build-docker", + action="store_true", + help="Build and push the Docker container", +) +args = parser.parse_args() + +session = Session.builder.create() +account = session.get_current_account() +user = session.get_current_user() +database = session.get_current_database() +schema = session.get_current_schema() +warehouse = session.get_current_warehouse() +role = session.get_current_role() + + +def run_sql_command(command: str): + print(f"Running SQL command: {command}") + result = session.sql(command).collect() + print(f"Result: {result}") + return result + + +# Check if the image repository exists, if not create it +repository_name = "TRULENS_REPOSITORY" +images = session.sql("SHOW IMAGE REPOSITORIES").collect() +repository_exists = any(image["name"] == repository_name for image in images) + +if not repository_exists: + session.sql(f"CREATE IMAGE REPOSITORY {repository_name}").collect() + print(f"Image repository {repository_name} created.") +else: + print(f"Image repository {repository_name} already exists.") + +# Retrieve the repository URL +repository_url = ( + session.sql(f"SHOW IMAGE REPOSITORIES LIKE '{repository_name}'") + .select('"repository_url"') + .collect()[0]["repository_url"] +) + +image_name = "trulens_dashboard" +image_tag = "latest" +app_name = "trulens_dashboard" +container_name = app_name + "_container" +if args.build_docker: + # local build, with docker + import subprocess + + subprocess.run( + [ + "docker", + "build", + "--platform", + "linux/amd64", + "-t", + f"{repository_url}/{image_name}:{image_tag}", + ".", + ], + check=True, + ) + subprocess.run( + ["docker", "push", f"{repository_url}/{image_name}:{image_tag}"], + check=True, + ) + # Create compute pool if it does not exist compute_pool_name = input("Enter compute pool name: ") compute_pools = session.sql("SHOW COMPUTE POOLS").collect() -compute_pool_exists = any(pool['name'] == compute_pool_name.upper() for pool in compute_pools) +compute_pool_exists = any( + pool["name"] == compute_pool_name.upper() for pool in compute_pools +) if compute_pool_exists: print(f"Compute pool {compute_pool_name} already exists") else: - session.sql(f"CREATE COMPUTE POOL {compute_pool_name} MIN_NODES = 1 MAX_NODES = 1 INSTANCE_FAMILY = CPU_X64_M").collect() + session.sql( + f"CREATE COMPUTE POOL {compute_pool_name} MIN_NODES = 1 MAX_NODES = 1 INSTANCE_FAMILY = CPU_X64_M" + ).collect() session.sql(f"DESCRIBE COMPUTE POOL {compute_pool_name}").collect() -# Create image repository -image_repository_name = f"trulens_image_repository" -session.sql(f"CREATE IMAGE REPOSITORY {image_repository_name}").collect() -session.sql("SHOW IMAGE REPOSITORIES").collect() - # Create network rule -network_rule_name = f"{compute_pool_name}_allow_all_network_rule" -session.sql(f"CREATE OR REPLACE NETWORK RULE {network_rule_name} TYPE = 'HOST_PORT' MODE = 'EGRESS' VALUE_LIST = ('0.0.0.0:443','0.0.0.0:80')").collect() +network_rule_name = f"{compute_pool_name}_allow_http_https" +session.sql( + f"CREATE OR REPLACE NETWORK RULE {network_rule_name} TYPE = 'HOST_PORT' MODE = 'EGRESS' VALUE_LIST = ('0.0.0.0:443','0.0.0.0:80')" +).collect() session.sql("SHOW NETWORK RULES").collect() # Create external access integration access_integration_name = f"{compute_pool_name}_access_integration" -session.sql(f"CREATE OR REPLACE EXTERNAL ACCESS INTEGRATION {access_integration_name} ALLOWED_NETWORK_RULES = ({network_rule_name}) ENABLED = true").collect() +session.sql( + f"CREATE OR REPLACE EXTERNAL ACCESS INTEGRATION {access_integration_name} ALLOWED_NETWORK_RULES = ({network_rule_name}) ENABLED = true" +).collect() session.sql("SHOW EXTERNAL ACCESS INTEGRATIONS").collect() -app_name = "trulens_dashboard" -secret_name = f"{schema}.{app_name}_login_credentials" -session.sql(f"CREATE SECRET {secret_name} TYPE=password USERNAME={user} PASSWORD='{password}'").collect() - -service_name=compute_pool_name+"_trulens_dashboard" -session.sql(""" +service_name = compute_pool_name + "_trulens_dashboard" +session.sql( + """ CREATE SERVICE {service_name} IN COMPUTE POOL {compute_pool_name} EXTERNAL_ACCESS_INTEGRATIONS = ({access_integration_name}) FROM SPECIFICATION $$ spec: containers: - - name: {container_name} - image: /{database}/{schema}/{container_name}/{app_name}:latest + - name: trulens-dashboard + image: /{database}/{schema}/{repository_name}/{app_name}:latest env: SNOWFLAKE_ACCOUNT: "{account}" SNOWFLAKE_DATABASE: "{database}" @@ -71,33 +117,29 @@ SNOWFLAKE_WAREHOUSE: "{warehouse}" SNOWFLAKE_ROLE: "{role}" RUN_DASHBOARD: "1" - secrets: - - snowflakeSecret: {secret_name} - secretKeyRef: username - envVarName: SNOWFLAKE_USER - - snowflakeSecret: {secret_name} - secretKeyRef: password - envVarName: SNOWFLAKE_PASSWORD endpoints: - name: trulens-demo-dashboard-endpoint port: 8484 public: true $$ -""".format(service_name=service_name, - compute_pool_name=compute_pool_name, - access_integration_name=access_integration_name, - container_name=app_name+"_container", - account=account, - database=database, - schema=schema, - warehouse=warehouse, - role=role, - app_name=app_name)).collect() +""".format( + service_name=service_name, + compute_pool_name=compute_pool_name, + access_integration_name=access_integration_name, + repository_name=repository_name, + account=account, + database=database, + schema=schema, + warehouse=warehouse, + role=role, + app_name=app_name, + ) +).collect() # Show services and get their status -session.sql(f"SHOW ENDPOINTS IN SERVICE {service_name}").collect() -session.sql("CALL SYSTEM$GET_SERVICE_STATUS('dkurokawa_trulens_demo_app')").collect() -session.sql("CALL SYSTEM$GET_SERVICE_STATUS('dkurokawa_trulens_demo_dashboard')").collect() +run_sql_command(f"SHOW ENDPOINTS IN SERVICE {service_name}") +run_sql_command(f"CALL SYSTEM$GET_SERVICE_STATUS('{service_name}')") +run_sql_command(f"CALL SYSTEM$GET_SERVICE_STATUS('{service_name}')") # Close the session -session.close() \ No newline at end of file +session.close() diff --git a/tools/snowflake/spcs_dashboard/run_dashboard.py b/tools/snowflake/spcs_dashboard/run_dashboard.py index 4983a405a..06e80348b 100644 --- a/tools/snowflake/spcs_dashboard/run_dashboard.py +++ b/tools/snowflake/spcs_dashboard/run_dashboard.py @@ -1,69 +1,33 @@ -from trulens.dashboard import run_dashboard -from trulens.core import TruSession -from snowflake.snowpark import Session -from trulens.connectors.snowflake import SnowflakeConnector -from snowflake.sqlalchemy import URL import os -# connection_params = { -# "account": os.environ.get("SNOWFLAKE_ACCOUNT"), -# "host": os.getenv("SNOWFLAKE_HOST"), -# "user": os.environ.get("SNOWFLAKE_USER"), -# "password": os.environ.get("SNOWFLAKE_PASSWORD"), -# "database": os.environ.get("SNOWFLAKE_DATABASE"), -# "schema": os.environ.get("SNOWFLAKE_SCHEMA"), -# "warehouse": os.environ.get("SNOWFLAKE_WAREHOUSE"), -# "role": os.environ.get("SNOWFLAKE_ROLE"), -# "init_server_side": False, -# } +from snowflake.snowpark import Session +from trulens.connectors.snowflake import SnowflakeConnector +from trulens.core import TruSession +from trulens.dashboard import run_dashboard -# connector = SnowflakeConnector(**connection_params) -# session = TruSession(connector=connector, init_server_side=False) def get_login_token(): - """ - Read the login token supplied automatically by Snowflake. These tokens - are short lived and should always be read right before creating any new connection. - """ - with open("/snowflake/session/token", "r") as f: - return f.read() + """ + Read the login token supplied automatically by Snowflake. These tokens + are short lived and should always be read right before creating any new connection. + """ + with open("/snowflake/session/token", "r") as f: + return f.read() + connection_params = { - "account": os.environ.get("SNOWFLAKE_ACCOUNT"), - "host": os.getenv("SNOWFLAKE_HOST"), - "authenticator": "oauth", - "token": get_login_token(), - "warehouse": os.environ.get("SNOWFLAKE_WAREHOUSE"), - "database": os.environ.get("SNOWFLAKE_DATABASE"), - "schema": os.environ.get("SNOWFLAKE_SCHEMA"), + "account": os.environ.get("SNOWFLAKE_ACCOUNT"), + "host": os.getenv("SNOWFLAKE_HOST"), + "authenticator": "oauth", + "token": get_login_token(), + "warehouse": os.environ.get("SNOWFLAKE_WAREHOUSE"), + "database": os.environ.get("SNOWFLAKE_DATABASE"), + "schema": os.environ.get("SNOWFLAKE_SCHEMA"), } snowpark_session = Session.builder.configs(connection_params).create() -# Set up sqlalchemy engine parameters. -conn = snowpark_session.connection -engine_params = {} -engine_params["paramstyle"] = "qmark" -engine_params["creator"] = lambda: conn -database_args = {"engine_params": engine_params} -# # Ensure any Cortex provider uses the only Snowflake connection allowed in this stored procedure. -# trulens.providers.cortex.provider._SNOWFLAKE_STORED_PROCEDURE_CONNECTION = ( -# conn -# ) -# Run deferred feedback evaluator. -db_url = URL( - account=snowpark_session.get_current_account(), - user=snowpark_session.get_current_user(), - password="password", - database=snowpark_session.get_current_database(), - schema=snowpark_session.get_current_schema(), - warehouse=snowpark_session.get_current_warehouse(), - role=snowpark_session.get_current_role(), -) -tru_session = TruSession( - database_url=db_url, - database_check_revision=False, # TODO: check revision in the future? - database_args=database_args, -) +connector = SnowflakeConnector(snowpark_session=snowpark_session) +tru_session = TruSession(connector=connector) tru_session.get_records_and_feedback() -run_dashboard(tru_session, port=8484, spcs_runtime=True) \ No newline at end of file +run_dashboard(tru_session, port=8484, spcs_runtime=True) From 272a194419cd5551e971ff912581805dbf6ec5ad Mon Sep 17 00:00:00 2001 From: Prudhvi Dharmana Date: Mon, 21 Oct 2024 23:42:27 -0700 Subject: [PATCH 4/4] fix imports --- src/dashboard/trulens/dashboard/streamlit.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/dashboard/trulens/dashboard/streamlit.py b/src/dashboard/trulens/dashboard/streamlit.py index 25f676ef5..05887f750 100644 --- a/src/dashboard/trulens/dashboard/streamlit.py +++ b/src/dashboard/trulens/dashboard/streamlit.py @@ -31,6 +31,7 @@ class FeedbackDisplay(BaseModel): calls: List[feedback_schema.FeedbackCall] icon: str + def get_spcs_login_token(): """ Read the login token supplied automatically by Snowflake. These tokens @@ -39,6 +40,7 @@ def get_spcs_login_token(): with open("/snowflake/session/token", "r") as f: return f.read() + def init_from_args(): """Parse command line arguments and initialize Tru with them. @@ -63,9 +65,11 @@ def init_from_args(): sys.exit(e.code) if args.spcs_runtime: + import os + from snowflake.snowpark import Session from trulens.connectors.snowflake import SnowflakeConnector - import os + connection_params = { "account": os.environ.get("SNOWFLAKE_ACCOUNT"), "host": os.getenv("SNOWFLAKE_HOST"),