Gestion des coûts des fonctions CortexAI avec Utilisation du compte¶

SELECT
    DATE_TRUNC('day', START_TIME) AS usage_date,
    FUNCTION_NAME,
    MODEL_NAME,
    SUM(CREDITS) AS total_credits,
    COUNT(DISTINCT QUERY_ID) AS query_count
FROM SNOWFLAKE.ACCOUNT_USAGE.CORTEX_AI_FUNCTIONS_USAGE_HISTORY
WHERE START_TIME >= DATEADD('day', -30, CURRENT_TIMESTAMP())
GROUP BY 1, 2, 3
ORDER BY usage_date DESC, total_credits DESC;

SELECT
    DATE_TRUNC('month', h.START_TIME) AS usage_month,
    u.NAME AS user_name,
    u.EMAIL,
    u.DEFAULT_ROLE,
    SUM(h.CREDITS) AS total_credits,
    COUNT(DISTINCT h.QUERY_ID) AS query_count
FROM SNOWFLAKE.ACCOUNT_USAGE.CORTEX_AI_FUNCTIONS_USAGE_HISTORY h
JOIN SNOWFLAKE.ACCOUNT_USAGE.USERS u
    ON h.USER_ID = u.USER_ID
WHERE h.START_TIME >= DATEADD('month', -3, CURRENT_TIMESTAMP())
GROUP BY 1, 2, 3, 4
ORDER BY usage_month DESC, total_credits DESC;

CREATE OR REPLACE NOTIFICATION INTEGRATION ai_cost_alerts
    TYPE = EMAIL
    ENABLED = TRUE
    ALLOWED_RECIPIENTS = ('admin@company.com', 'finops@company.com')

CREATE TABLE IF NOT EXISTS AI_FUNCTIONS_ALERT_STATE (
    ALERT_NAME VARCHAR NOT NULL,
    ALERT_MONTH DATE NOT NULL,
    SENT_AT TIMESTAMP_LTZ DEFAULT CURRENT_TIMESTAMP(),
    CREDITS_AT_ALERT NUMBER(38,6),
    PRIMARY KEY (ALERT_NAME, ALERT_MONTH)
);

CREATE OR REPLACE ALERT ai_functions_monthly_spend_alert
    WAREHOUSE = <your_warehouse>
    SCHEDULE = 'USING CRON 0 * * * * UTC'  -- Runs every hour
    IF (EXISTS (
        SELECT 1
        FROM SNOWFLAKE.ACCOUNT_USAGE.CORTEX_AI_FUNCTIONS_USAGE_HISTORY
        WHERE START_TIME >= DATE_TRUNC('month', CURRENT_TIMESTAMP())
        HAVING SUM(CREDITS) > 1000  -- adjust the limit accordingly
    ))
    THEN
        CALL SEND_MONTHLY_SPEND_ALERT(1000);  -- please adjust the limit accordingly

-- enable the alert
ALTER ALERT ai_functions_monthly_spend_alert RESUME;

-- Make sure alert exists
SHOW ALERTS LIKE 'ai_functions_monthly_spend_alert';

-- Check alert history
SELECT *
FROM TABLE(INFORMATION_SCHEMA.ALERT_HISTORY(
    SCHEDULED_TIME_RANGE_START => DATEADD('day', -1, CURRENT_TIMESTAMP()),
    ALERT_NAME => 'ai_functions_monthly_spend_alert'
))
ORDER BY SCHEDULED_TIME DESC;

-- Check which months have had alerts sent
SELECT * FROM AI_FUNCTIONS_ALERT_STATE ORDER BY ALERT_MONTH DESC;

-- Create a role specifically for AI Function access
CREATE ROLE IF NOT EXISTS AI_FUNCTIONS_USER_ROLE;

-- Grant necessary privileges to the role
GRANT DATABASE ROLE SNOWFLAKE.CORTEX_USER TO ROLE AI_FUNCTIONS_USER_ROLE;

-- Grant usage on warehouse
GRANT USAGE ON WAREHOUSE AI_FUNCTIONS_WAREHOUSE TO ROLE AI_FUNCTIONS_USER_ROLE;

CREATE TABLE IF NOT EXISTS AI_FUNCTIONS_ACCESS_CONTROL (
    USER_NAME VARCHAR NOT NULL,
    USER_ID NUMBER,
    GRANTED_AT TIMESTAMP_LTZ DEFAULT CURRENT_TIMESTAMP(),
    MONTHLY_CREDIT_LIMIT NUMBER(38,6) DEFAULT 100,  -- adjust the limit accordingly
    IS_ACTIVE BOOLEAN DEFAULT TRUE,
    REVOKED_AT TIMESTAMP_LTZ,
    REVOCATION_REASON VARCHAR,
    PRIMARY KEY (USER_NAME)
);

CREATE OR REPLACE PROCEDURE GRANT_AI_FUNCTIONS_ACCESS(
    P_USER_NAME VARCHAR,
    P_MONTHLY_LIMIT NUMBER(38,6) DEFAULT 100  -- adjust the limit accordingly
)
RETURNS VARCHAR
LANGUAGE SQL
AS
$$
DECLARE
    v_user_id NUMBER;
BEGIN
    -- Look up USER_ID from account usage
    SELECT USER_ID INTO :v_user_id
    FROM SNOWFLAKE.ACCOUNT_USAGE.USERS
    WHERE NAME = :P_USER_NAME
    LIMIT 1;

    -- Grant the AI Functions role to the user
    EXECUTE IMMEDIATE 'GRANT ROLE AI_FUNCTIONS_USER_ROLE TO USER ' || P_USER_NAME;

    -- Register or update the user in the access control table
    MERGE INTO AI_FUNCTIONS_ACCESS_CONTROL tgt
    USING (SELECT :P_USER_NAME AS USER_NAME) src
    ON tgt.USER_NAME = src.USER_NAME
    WHEN MATCHED THEN
        UPDATE SET
            USER_ID = :v_user_id,
            IS_ACTIVE = TRUE,
            MONTHLY_CREDIT_LIMIT = :P_MONTHLY_LIMIT,
            GRANTED_AT = CURRENT_TIMESTAMP(),
            REVOKED_AT = NULL,
            REVOCATION_REASON = NULL
    WHEN NOT MATCHED THEN
        INSERT (USER_NAME, USER_ID, MONTHLY_CREDIT_LIMIT, IS_ACTIVE)
        VALUES (:P_USER_NAME, :v_user_id, :P_MONTHLY_LIMIT, TRUE);

    RETURN 'Access granted to ' || P_USER_NAME || ' with monthly limit of ' || P_MONTHLY_LIMIT || ' credits';
END;
$$;

CALL GRANT_AI_FUNCTIONS_ACCESS('ALICE', 1000);  -- grants access to user ALICE with a monthly limit of 100 credits
CALL GRANT_AI_FUNCTIONS_ACCESS('BOB', 2000);    -- grants access to user BOB with a monthly limit of 200 credits

-- Create a procedure to re-grant access to all entitled users
CREATE OR REPLACE PROCEDURE GRANT_ALL_ENTITLED_USERS()
RETURNS TABLE (USER_NAME VARCHAR, CREDIT_LIMIT NUMBER, ACTION VARCHAR)
LANGUAGE SQL
AS
$$
DECLARE
    result RESULTSET;
BEGIN
    result := (
        SELECT
            USER_NAME,
            MONTHLY_CREDIT_LIMIT AS CREDIT_LIMIT,
            'GRANTED' AS ACTION
        FROM AI_FUNCTIONS_ACCESS_CONTROL
    );

    -- Re-grant access for each entitled user
    FOR rec IN result DO
        CALL GRANT_AI_FUNCTIONS_ACCESS(rec.USER_NAME, rec.CREDIT_LIMIT);
    END FOR;

    RETURN TABLE(result);
END;
$$;

-- Create a task to run on the 1st of each month at midnight UTC
CREATE OR REPLACE TASK MONTHLY_AI_FUNCTIONS_ACCESS_REFRESH
    WAREHOUSE = <your_warehouse>
    SCHEDULE = 'USING CRON 0 0 1 * * UTC'  -- 1st day of each month at 00:00 UTC
AS
    CALL GRANT_ALL_ENTITLED_USERS();

-- Enable the task
ALTER TASK MONTHLY_AI_FUNCTIONS_ACCESS_REFRESH RESUME;

-- Run once initially to populate grantees
CALL GRANT_ALL_ENTITLED_USERS();

-- Verify task status
SHOW TASKS LIKE 'MONTHLY_AI_FUNCTIONS_ACCESS_REFRESH';

-- Create a procedure to re-grant access to all entitled users
CREATE OR REPLACE PROCEDURE GRANT_ALL_ENTITLED_USERS()
RETURNS TABLE (USER_NAME VARCHAR, CREDIT_LIMIT NUMBER, ACTION VARCHAR)
LANGUAGE SQL
AS
$$
DECLARE
    result RESULTSET;
BEGIN
    result := (
        SELECT
            USER_NAME,
            MONTHLY_CREDIT_LIMIT AS CREDIT_LIMIT,
            'GRANTED' AS ACTION
        FROM AI_FUNCTIONS_ACCESS_CONTROL
    );

    -- Re-grant access for each entitled user
    FOR rec IN result DO
        CALL GRANT_AI_FUNCTIONS_ACCESS(rec.USER_NAME, rec.CREDIT_LIMIT);
    END FOR;

    RETURN TABLE(result);
END;
$$;

-- Create a task to run on the 1st of each month at midnight UTC
CREATE OR REPLACE TASK MONTHLY_AI_FUNCTIONS_ACCESS_REFRESH
    WAREHOUSE = <your_warehouse>
    SCHEDULE = 'USING CRON 0 0 1 * * UTC'  -- 1st day of each month at 00:00 UTC
AS
    CALL GRANT_ALL_ENTITLED_USERS();

-- Enable the task
ALTER TASK MONTHLY_AI_FUNCTIONS_ACCESS_REFRESH RESUME;

-- Run once initially to populate grantees
CALL GRANT_ALL_ENTITLED_USERS();

-- Verify task status
SHOW TASKS LIKE 'MONTHLY_AI_FUNCTIONS_ACCESS_REFRESH';

-- Create a procedure to detect and cancel expensive runaway queries
CREATE OR REPLACE PROCEDURE MONITOR_AND_CANCEL_RUNAWAY_QUERIES(
    P_CREDIT_THRESHOLD NUMBER DEFAULT 50  -- adjust the limit accordingly
)
RETURNS TABLE (
    QUERY_ID VARCHAR,
    USER_NAME VARCHAR,
    FUNCTION_NAME VARCHAR,
    MODEL_NAME VARCHAR,
    CREDITS NUMBER,
    START_TIME TIMESTAMP_LTZ,
    ACTION VARCHAR
)
LANGUAGE SQL
AS
$$
DECLARE
    result RESULTSET;
BEGIN
    -- Find queries from the last 48 hours that exceed the threshold and are still running
    result := (
        SELECT
            h.QUERY_ID,
            u.NAME AS USER_NAME,
            h.FUNCTION_NAME,
            h.MODEL_NAME,
            h.CREDITS,
            h.START_TIME,
            h.ROLE_NAMES,
            h.QUERY_TAG,
            h.WAREHOUSE_ID,
            'CANCELLED' AS ACTION
        FROM SNOWFLAKE.ACCOUNT_USAGE.CORTEX_AI_FUNCTIONS_USAGE_HISTORY h
        LEFT JOIN SNOWFLAKE.ACCOUNT_USAGE.USERS u
            ON h.USER_ID = u.USER_ID
        WHERE h.START_TIME >= DATEADD('hour', -48, CURRENT_TIMESTAMP())
        AND h.CREDITS > :P_CREDIT_THRESHOLD
        AND h.IS_COMPLETED = FALSE
    );

    -- Cancel each runaway query and send alert
    FOR rec IN result DO
        -- Attempt to cancel the query
        BEGIN
            EXECUTE IMMEDIATE 'SELECT SYSTEM$CANCEL_QUERY(''' || rec.QUERY_ID || ''')';
        EXCEPTION
            WHEN OTHER THEN
                NULL;  -- Query may have already completed
        END;

        -- Send alert with query details
        CALL SYSTEM$SEND_EMAIL(
            'ai_cost_alerts',
            'admin@company.com',
            'Runaway AI Query Cancelled - ' || rec.QUERY_ID,
            'A runaway AI Function query has been cancelled due to excessive cost.\n\n' ||
            'Query Details:\n' ||
            '- Query ID: ' || rec.QUERY_ID || '\n' ||
            '- User: ' || COALESCE(rec.USER_NAME, 'Unknown') || '\n' ||
            '- Function: ' || rec.FUNCTION_NAME || '\n' ||
            '- Model: ' || rec.MODEL_NAME || '\n' ||
            '- Credits Used: ' || rec.CREDITS::VARCHAR || '\n' ||
            '- Threshold: ' || :P_CREDIT_THRESHOLD::VARCHAR || '\n' ||
            '- Start Time: ' || rec.START_TIME::VARCHAR || '\n' ||
            '- Roles: ' || COALESCE(rec.ROLE_NAMES::VARCHAR, 'N/A') || '\n' ||
            '- Query Tag: ' || COALESCE(rec.QUERY_TAG, 'N/A') || '\n' ||
            '- Warehouse ID: ' || COALESCE(rec.WAREHOUSE_ID::VARCHAR, 'N/A') || '\n\n' ||
            'Please investigate this query and take appropriate action.'
        );
    END FOR;

    RETURN TABLE(result);
END;
$$;

-- Create a task to monitor and cancel runaway queries every hour
CREATE OR REPLACE TASK MONITOR_RUNAWAY_AI_QUERIES
    WAREHOUSE = <your_warehouse>
    SCHEDULE = 'USING CRON 0 * * * * UTC'  -- Every hour
AS
    CALL MONITOR_AND_CANCEL_RUNAWAY_QUERIES(50);  -- adjust the limit accordingly

-- Enable the task
ALTER TASK MONITOR_RUNAWAY_AI_QUERIES RESUME;

-- Verify task status
SHOW TASKS LIKE 'MONITOR_RUNAWAY_AI_QUERIES';

-- Check task execution history
SELECT *
FROM TABLE(INFORMATION_SCHEMA.TASK_HISTORY(
    SCHEDULED_TIME_RANGE_START => DATEADD('day', -1, CURRENT_TIMESTAMP()),
    TASK_NAME => 'MONITOR_RUNAWAY_AI_QUERIES'
))
ORDER BY SCHEDULED_TIME DESC;