Use conda channel trulens packages by default. (#1570)

truera · Oct 21, 2024 · 7e673f8 · 7e673f8
1 parent b148930
commit 7e673f8
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 23 deletions.
diff --git a/src/connectors/snowflake/trulens/connectors/snowflake/connector.py b/src/connectors/snowflake/trulens/connectors/snowflake/connector.py
@@ -43,6 +43,7 @@ def __init__(
         role: Optional[str] = None,
         snowpark_session: Optional[Session] = None,
         init_server_side: bool = False,
+        init_server_side_with_staged_packages: bool = False,
         database_redact_keys: bool = False,
         database_prefix: Optional[str] = None,
         database_args: Optional[Dict[str, Any]] = None,
@@ -116,6 +117,7 @@ def __init__(
         self._init_with_snowpark_session(
             snowpark_session,
             init_server_side,
+            init_server_side_with_staged_packages,
             database_redact_keys,
             database_prefix,
             database_args,
@@ -127,6 +129,7 @@ def _init_with_snowpark_session(
         self,
         snowpark_session: Session,
         init_server_side: bool,
+        init_server_side_with_staged_packages: bool,
         database_redact_keys: bool,
         database_prefix: Optional[str],
         database_args: Optional[Dict[str, Any]],
@@ -176,7 +179,9 @@ def _init_with_snowpark_session(
 
         if init_server_side:
             ServerSideEvaluationArtifacts(
-                snowpark_session, database_args["database_prefix"]
+                snowpark_session,
+                database_args["database_prefix"],
+                init_server_side_with_staged_packages,
             ).set_up_all()
 
     @classmethod

diff --git a/...nnectors/snowflake/trulens/connectors/snowflake/utils/server_side_evaluation_artifacts.py b/...nnectors/snowflake/trulens/connectors/snowflake/utils/server_side_evaluation_artifacts.py
@@ -13,12 +13,11 @@
     "server_side_evaluation_stored_procedure.py",
 )
 
-_ZIPS_TO_UPLOAD = [
-    "snowflake_sqlalchemy.zip",
-    "trulens_connectors_snowflake.zip",
-    "trulens_core.zip",
-    "trulens_feedback.zip",
-    "trulens_providers_cortex.zip",
+_TRULENS_PACKAGES = [
+    "trulens-connectors-snowflake",
+    "trulens-core",
+    "trulens-feedback",
+    "trulens-providers-cortex",
 ]
 
 
@@ -29,15 +28,18 @@ def __init__(
         self,
         session: Session,
         database_prefix: str,
+        use_staged_packages: bool,
     ) -> None:
         self._session = session
         self._database = session.get_current_database()
         self._schema = session.get_current_schema()
         self._warehouse = session.get_current_warehouse()
         self._database_prefix = database_prefix
+        self._use_staged_packages = use_staged_packages
 
     def set_up_all(self) -> None:
-        self._set_up_stage()
+        if self._use_staged_packages:
+            self._set_up_stage()
         self._set_up_stream()
         self._set_up_stored_procedure()
         self._set_up_wrapper_stored_procedure()
@@ -51,9 +53,9 @@ def _set_up_stage(self) -> None:
         data_directory = os.path.join(
             os.path.dirname(__file__), "../../../data/snowflake_stage_zips"
         )
-        for zip_to_upload in _ZIPS_TO_UPLOAD:
+        for trulens_package in _TRULENS_PACKAGES:
             self._session.file.put(
-                os.path.join(data_directory, zip_to_upload),
+                os.path.join(data_directory, f"{trulens_package}.zip"),
                 f"@{_STAGE_NAME}",
             )
 
@@ -67,9 +69,17 @@ def _set_up_stream(self) -> None:
         )
 
     def _set_up_stored_procedure(self) -> None:
-        imports = ",".join([
-            f"'@{_STAGE_NAME}/{curr}'" for curr in _ZIPS_TO_UPLOAD
-        ])
+        if self._use_staged_packages:
+            import_packages = ",".join([
+                f"'@{_STAGE_NAME}/{curr}.zip'" for curr in _TRULENS_PACKAGES
+            ])
+            import_statement = f"IMPORTS = ({import_packages})"
+            trulens_packages = ""
+        else:
+            import_statement = ""
+            trulens_packages = "".join([
+                f"'{curr}'," for curr in _TRULENS_PACKAGES
+            ])
         with open(_PYTHON_STORED_PROCEDURE_CODE_FILENAME, "r") as fh:
             python_code = fh.read()
         self._run_query(
@@ -79,6 +89,7 @@ def _set_up_stored_procedure(self) -> None:
                 LANGUAGE PYTHON
                 RUNTIME_VERSION = '3.11'
                 PACKAGES = (
+                    {trulens_packages}
                     -- TODO(dkurokawa): get these package versions automatically.
                     'alembic',
                     'dill',
@@ -96,11 +107,12 @@ def _set_up_stored_procedure(self) -> None:
                     'scikit-learn',
                     'scipy',
                     'snowflake-snowpark-python',
+                    'snowflake-sqlalchemy',
                     'sqlalchemy',
                     'tqdm',
                     'typing_extensions'
                 )
-                IMPORTS = ({imports})
+                {import_statement}
                 HANDLER = 'run'
             AS
                 $$\n{python_code}$$

diff --git a/tests/util/snowflake_test_case.py b/tests/util/snowflake_test_case.py
@@ -37,10 +37,10 @@ def tearDown(self):
         # [HACK!] Clean up any instances of `TruSession` so tests don't interfere with each other.
         for key in [
             curr
-            for curr in core_session.TruSession._instances
-            if curr[0] == "TruSession"
+            for curr in core_session.TruSession._singleton_instances
+            if curr[0] == "trulens.core.session.TruSession"
         ]:
-            del core_session.TruSession._instances[key]
+            del core_session.TruSession._singleton_instances[key]
         # Clean up any Snowflake schemas.
         schemas_not_deleted = []
         for curr in self._snowflake_schemas_to_delete:
@@ -98,6 +98,7 @@ def get_session(
             schema=self._schema,
             **self._snowflake_connection_parameters,
             init_server_side=True,
+            init_server_side_with_staged_packages=True,
         )
         session = core_session.TruSession(connector=connector)
         self.assertIn(self._schema, self.list_schemas())

diff --git a/zip_wheels.sh b/zip_wheels.sh
@@ -37,9 +37,7 @@ zip_wheel() {
 # Zip wheels.
 rm -rf ${OUTPUT_DIRECTORY}
 mkdir -p ${OUTPUT_DIRECTORY}
-zip_wheel trulens-connectors-snowflake ${OUTPUT_DIRECTORY}/trulens_connectors_snowflake.zip
-zip_wheel trulens-core ${OUTPUT_DIRECTORY}/trulens_core.zip
-zip_wheel trulens-feedback ${OUTPUT_DIRECTORY}/trulens_feedback.zip
-zip_wheel trulens-providers-cortex ${OUTPUT_DIRECTORY}/trulens_providers_cortex.zip
-snow snowpark package create snowflake-sqlalchemy==1.6.1
-mv snowflake_sqlalchemy.zip ${OUTPUT_DIRECTORY}
+zip_wheel trulens-connectors-snowflake ${OUTPUT_DIRECTORY}/trulens-connectors-snowflake.zip
+zip_wheel trulens-core ${OUTPUT_DIRECTORY}/trulens-core.zip
+zip_wheel trulens-feedback ${OUTPUT_DIRECTORY}/trulens-feedback.zip
+zip_wheel trulens-providers-cortex ${OUTPUT_DIRECTORY}/trulens-providers-cortex.zip