test: move Claude3 tests to load test (#997)

* test: move Claude3 tests to load test

* add conftest

Author: GarrettWu

tests/system/conftest.py

@@ -146,16 +146,6 @@ def session() -> Generator[bigframes.Session, None, None]:
     session.close()  # close generated session at cleanup time
 
 
-@pytest.fixture(scope="session")
-def session_us_east5() -> Generator[bigframes.Session, None, None]:
-    context = bigframes.BigQueryOptions(
-        location="us-east5",
-    )
-    session = bigframes.Session(context=context)
-    yield session
-    session.close()  # close generated session at cleanup time
-
-
 @pytest.fixture(scope="session")
 def session_load() -> Generator[bigframes.Session, None, None]:
     context = bigframes.BigQueryOptions(location="US", project="bigframes-load-testing")
@@ -188,6 +178,11 @@ def session_tokyo(tokyo_location: str) -> Generator[bigframes.Session, None, Non
     session.close()  # close generated session at cleanup type
 
 
+@pytest.fixture(scope="session")
+def bq_connection(bigquery_client: bigquery.Client) -> str:
+    return f"{bigquery_client.project}.{bigquery_client.location}.bigframes-rf-conn"
+
+
 @pytest.fixture(scope="session", autouse=True)
 def cleanup_datasets(bigquery_client: bigquery.Client) -> None:
     """Cleanup any datasets that were created but not cleaned up."""
@@ -728,6 +723,25 @@ def new_penguins_df(session, new_penguins_pandas_df):
     return session.read_pandas(new_penguins_pandas_df)
 
 
+@pytest.fixture(scope="session")
+def llm_text_pandas_df():
+    """Additional data matching the penguins dataset, with a new index"""
+    return pd.DataFrame(
+        {
+            "prompt": [
+                "What is BigQuery?",
+                "What is BQML?",
+                "What is BigQuery DataFrame?",
+            ],
+        }
+    )
+
+
+@pytest.fixture(scope="session")
+def llm_text_df(session, llm_text_pandas_df):
+    return session.read_pandas(llm_text_pandas_df)
+
+
 @pytest.fixture(scope="session")
 def penguins_linear_model_name(
     session: bigframes.Session, dataset_id_permanent, penguins_table_id

tests/system/load/conftest.py

@@ -0,0 +1,39 @@
+# Copyright 2024 Google LLC
+#
+# 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.
+
+
+from typing import Generator
+
+import pytest
+
+import bigframes
+
+
+# Override the session to target at bigframes-load-testing at all load tests. That allows to run load tests locally with authentic env.
+@pytest.fixture(scope="session")
+def session() -> Generator[bigframes.Session, None, None]:
+    context = bigframes.BigQueryOptions(location="US", project="bigframes-load-testing")
+    session = bigframes.Session(context=context)
+    yield session
+    session.close()  # close generated session at cleanup time
+
+
+@pytest.fixture(scope="session")
+def session_us_east5() -> Generator[bigframes.Session, None, None]:
+    context = bigframes.BigQueryOptions(
+        location="us-east5", project="bigframes-load-testing"
+    )
+    session = bigframes.Session(context=context)
+    yield session
+    session.close()  # close generated session at cleanup time

tests/system/load/test_llm.py

@@ -90,3 +90,69 @@ def test_llm_gemini_configure_fit(llm_fine_tune_df_default_index, llm_remote_tex
         index=3,
     )
     # TODO(ashleyxu b/335492787): After bqml rolled out version control: save, load, check parameters to ensure configuration was kept
+
+
+# (b/366290533): Claude models are of extremely low capacity. The tests should reside in small tests. Moving these here just to protect BQML's shared capacity(as load test only runs once per day.) and make sure we still have minimum coverage.
+@pytest.mark.parametrize(
+    "model_name",
+    ("claude-3-sonnet", "claude-3-haiku", "claude-3-5-sonnet", "claude-3-opus"),
+)
+@pytest.mark.flaky(retries=3, delay=120)
+def test_claude3_text_generator_create_load(
+    dataset_id, model_name, session, session_us_east5, bq_connection
+):
+    if model_name in ("claude-3-5-sonnet", "claude-3-opus"):
+        session = session_us_east5
+    claude3_text_generator_model = llm.Claude3TextGenerator(
+        model_name=model_name, connection_name=bq_connection, session=session
+    )
+    assert claude3_text_generator_model is not None
+    assert claude3_text_generator_model._bqml_model is not None
+
+    # save, load to ensure configuration was kept
+    reloaded_model = claude3_text_generator_model.to_gbq(
+        f"{dataset_id}.temp_text_model", replace=True
+    )
+    assert f"{dataset_id}.temp_text_model" == reloaded_model._bqml_model.model_name
+    assert reloaded_model.connection_name == bq_connection
+    assert reloaded_model.model_name == model_name
+
+
+@pytest.mark.parametrize(
+    "model_name",
+    ("claude-3-sonnet", "claude-3-haiku", "claude-3-5-sonnet", "claude-3-opus"),
+)
+@pytest.mark.flaky(retries=3, delay=120)
+def test_claude3_text_generator_predict_default_params_success(
+    llm_text_df, model_name, session, session_us_east5, bq_connection
+):
+    if model_name in ("claude-3-5-sonnet", "claude-3-opus"):
+        session = session_us_east5
+    claude3_text_generator_model = llm.Claude3TextGenerator(
+        model_name=model_name, connection_name=bq_connection, session=session
+    )
+    df = claude3_text_generator_model.predict(llm_text_df).to_pandas()
+    utils.check_pandas_df_schema_and_index(
+        df, columns=utils.ML_GENERATE_TEXT_OUTPUT, index=3, col_exact=False
+    )
+
+
+@pytest.mark.parametrize(
+    "model_name",
+    ("claude-3-sonnet", "claude-3-haiku", "claude-3-5-sonnet", "claude-3-opus"),
+)
+@pytest.mark.flaky(retries=3, delay=120)
+def test_claude3_text_generator_predict_with_params_success(
+    llm_text_df, model_name, session, session_us_east5, bq_connection
+):
+    if model_name in ("claude-3-5-sonnet", "claude-3-opus"):
+        session = session_us_east5
+    claude3_text_generator_model = llm.Claude3TextGenerator(
+        model_name=model_name, connection_name=bq_connection, session=session
+    )
+    df = claude3_text_generator_model.predict(
+        llm_text_df, max_output_tokens=100, top_k=20, top_p=0.5
+    ).to_pandas()
+    utils.check_pandas_df_schema_and_index(
+        df, columns=utils.ML_GENERATE_TEXT_OUTPUT, index=3, col_exact=False
+    )

tests/system/small/ml/conftest.py

@@ -34,11 +34,6 @@
 )
 
 
-@pytest.fixture(scope="session")
-def bq_connection(bigquery_client) -> str:
-    return f"{bigquery_client.project}.us.bigframes-rf-conn"
-
-
 @pytest.fixture(scope="session")
 def penguins_bqml_linear_model(session, penguins_linear_model_name) -> core.BqmlModel:
     model = session.bqclient.get_model(penguins_linear_model_name)
@@ -157,20 +152,6 @@ def penguins_pca_model(
     )
 
 
-@pytest.fixture(scope="session")
-def llm_text_pandas_df():
-    """Additional data matching the penguins dataset, with a new index"""
-    return pd.DataFrame(
-        {
-            "prompt": [
-                "What is BigQuery?",
-                "What is BQML?",
-                "What is BigQuery DataFrame?",
-            ],
-        }
-    )
-
-
 @pytest.fixture(scope="session")
 def onnx_iris_pandas_df():
     """Data matching the iris dataset."""
@@ -212,11 +193,6 @@ def xgboost_iris_df(session, xgboost_iris_pandas_df):
     return session.read_pandas(xgboost_iris_pandas_df)
 
 
-@pytest.fixture(scope="session")
-def llm_text_df(session, llm_text_pandas_df):
-    return session.read_pandas(llm_text_pandas_df)
-
-
 @pytest.fixture(scope="session")
 def bqml_palm2_text_generator_model(session, bq_connection) -> core.BqmlModel:
     options = {

tests/system/small/ml/test_llm.py

@@ -414,73 +414,6 @@ def test_gemini_text_generator_predict_with_params_success(
     )
 
 
-@pytest.mark.parametrize(
-    "model_name",
-    ("claude-3-sonnet", "claude-3-haiku", "claude-3-5-sonnet", "claude-3-opus"),
-)
-@pytest.mark.flaky(retries=3, delay=120)
-def test_claude3_text_generator_create_load(
-    dataset_id, model_name, session, session_us_east5, bq_connection
-):
-    if model_name in ("claude-3-5-sonnet", "claude-3-opus"):
-        session = session_us_east5
-    claude3_text_generator_model = llm.Claude3TextGenerator(
-        model_name=model_name, connection_name=bq_connection, session=session
-    )
-    assert claude3_text_generator_model is not None
-    assert claude3_text_generator_model._bqml_model is not None
-
-    # save, load to ensure configuration was kept
-    reloaded_model = claude3_text_generator_model.to_gbq(
-        f"{dataset_id}.temp_text_model", replace=True
-    )
-    assert f"{dataset_id}.temp_text_model" == reloaded_model._bqml_model.model_name
-    assert reloaded_model.connection_name == bq_connection
-    assert reloaded_model.model_name == model_name
-
-
-@pytest.mark.skip("b/366290533 too many requests are exhausting bqml capacity")
-@pytest.mark.parametrize(
-    "model_name",
-    ("claude-3-sonnet", "claude-3-haiku", "claude-3-5-sonnet", "claude-3-opus"),
-)
-@pytest.mark.flaky(retries=3, delay=120)
-def test_claude3_text_generator_predict_default_params_success(
-    llm_text_df, model_name, session, session_us_east5, bq_connection
-):
-    if model_name in ("claude-3-5-sonnet", "claude-3-opus"):
-        session = session_us_east5
-    claude3_text_generator_model = llm.Claude3TextGenerator(
-        model_name=model_name, connection_name=bq_connection, session=session
-    )
-    df = claude3_text_generator_model.predict(llm_text_df).to_pandas()
-    utils.check_pandas_df_schema_and_index(
-        df, columns=utils.ML_GENERATE_TEXT_OUTPUT, index=3, col_exact=False
-    )
-
-
-@pytest.mark.skip("b/366290533 too many requests are exhausting bqml capacity")
-@pytest.mark.parametrize(
-    "model_name",
-    ("claude-3-sonnet", "claude-3-haiku", "claude-3-5-sonnet", "claude-3-opus"),
-)
-@pytest.mark.flaky(retries=3, delay=120)
-def test_claude3_text_generator_predict_with_params_success(
-    llm_text_df, model_name, session, session_us_east5, bq_connection
-):
-    if model_name in ("claude-3-5-sonnet", "claude-3-opus"):
-        session = session_us_east5
-    claude3_text_generator_model = llm.Claude3TextGenerator(
-        model_name=model_name, connection_name=bq_connection, session=session
-    )
-    df = claude3_text_generator_model.predict(
-        llm_text_df, max_output_tokens=100, top_k=20, top_p=0.5
-    ).to_pandas()
-    utils.check_pandas_df_schema_and_index(
-        df, columns=utils.ML_GENERATE_TEXT_OUTPUT, index=3, col_exact=False
-    )
-
-
 @pytest.mark.flaky(retries=2)
 def test_llm_palm_score(llm_fine_tune_df_default_index):
     model = llm.PaLM2TextGenerator(model_name="text-bison")