From b9d567dd2779125d610805836d3da685c3df13eb Mon Sep 17 00:00:00 2001 From: Snowflake Authors Date: Tue, 30 Jun 2026 17:47:43 -0700 Subject: [PATCH] Project import generated by Copybara. GitOrigin-RevId: 35231831e1be58aa4ad02929024253216badab57 --- CHANGELOG.md | 74 +++- MODULE.bazel | 1 + .../conda-optional-dependency-llm.yml | 2 +- .../conda-optional-dependency-ml.yml | 2 +- .../conda-optional-dependency-torch.yml | 2 +- ci/conda_recipe/meta.yaml | 4 +- codegen/sklearn_wrapper_generator.py | 1 + ...nsformer_autogen_test_template.py_template | 10 - requirements.yml | 2 +- snowflake/ml/_internal/file_utils.py | 51 ++- .../_client/experiment_tracking_sql_client.py | 17 +- .../experiment_tracking_sql_client_test.py | 32 ++ snowflake/ml/feature_store/BUILD.bazel | 12 + .../ml/feature_store/_internal/__init__.py | 0 snowflake/ml/feature_store/feature_store.py | 177 +++++++--- .../feature_store_entity_test.py | 12 +- ...ure_store_initialization_warehouse_test.py | 89 +++++ .../feature_store_oft_online_config_test.py | 228 ++++++++++++- .../feature_store_realtime_test.py | 11 +- .../feature_store_snapshot_validation_test.py | 14 +- snowflake/ml/feature_store/feature_view.py | 88 +++++ .../ml/feature_store/feature_view_test.py | 81 +++-- .../ml/feature_store/realtime_registration.py | 6 +- snowflake/ml/feature_store/spec/builder.py | 41 ++- .../ml/feature_store/spec/builder_test.py | 51 ++- .../feature_store/streaming_registration.py | 16 +- .../streaming_registration_test.py | 50 +++ snowflake/ml/model/_client/model/BUILD.bazel | 5 +- .../_client/model/batch_inference_specs.py | 59 ++++ .../model/batch_inference_specs_test.py | 33 ++ .../model/_client/model/model_version_impl.py | 105 ++++++ .../_client/model/model_version_impl_test.py | 92 +++++ snowflake/ml/model/_client/ops/BUILD.bazel | 1 + snowflake/ml/model/_client/ops/service_ops.py | 158 ++++++++- .../ml/model/_client/ops/service_ops_test.py | 192 ++++++++++- .../ml/model/_client/service/BUILD.bazel | 22 ++ .../service/inference_job_service_spec.py | 75 ++++ .../inference_job_service_spec_test.py | 126 +++++++ snowflake/ml/model/_client/sql/service.py | 60 ++++ .../ml/model/_client/sql/service_test.py | 164 +++++++++ .../ml/model/_model_composer/BUILD.bazel | 22 ++ .../huggingface_lazy_uploader.py | 244 +++++++++++++ .../huggingface_lazy_uploader_test.py | 211 ++++++++++++ .../model/_model_composer/model_composer.py | 8 + .../_model_composer/model_composer_test.py | 42 ++- .../_packager/model_handlers/BUILD.bazel | 2 + .../model_handlers/huggingface/_handler.py | 25 +- .../model_handlers/huggingface/_utils.py | 27 +- .../model_handlers/sentence_transformers.py | 96 +++++- .../_packager/model_handlers_test/BUILD.bazel | 12 + .../huggingface_pipeline_test.py | 45 +++ .../huggingface_utils_test.py | 41 +++ .../sentence_transformers_test.py | 322 +++++++++++++++++- .../ml/model/_packager/model_meta/BUILD.bazel | 1 + .../model/_packager/model_meta/model_meta.py | 3 + snowflake/ml/model/models/huggingface.py | 270 ++++++++++++++- snowflake/ml/model/models/huggingface_test.py | 115 +++++++ .../models/sentence_transformers_test.py | 83 +++++ .../snowpark_handlers.py | 12 +- .../snowpark_trainer.py | 27 +- snowflake/ml/version.py | 2 +- .../feature_store_feature_group_bundled.py | 132 +++++++ .../feature_store_online_management_test.py | 55 ++- .../feature_store_rollup_test.py | 49 +++ .../feature_store_streaming_fv_bundled.py | 46 +++ .../feature_store_streaming_fv_integ_base.py | 3 +- .../ml/feature_store/feature_store_test.py | 111 +++++- .../snowflake/ml/jobs/jobs_integ_test.py | 3 +- .../model/model_version_impl_integ_test.py | 112 +++--- ...registry_batch_inference_cse_stage_test.py | 1 + ...tom_multi_modality_batch_inference_test.py | 1 + ...ggingface_pipeline_batch_inference_test.py | 2 + .../registry_keras_batch_inference_test.py | 1 + ...ging_face_pipeline_batch_inference_test.py | 1 + ...istry_new_hf_tasks_batch_inference_test.py | 3 + .../registry_pip_only_batch_inference_test.py | 68 +--- ...istry_pip_only_gpu_batch_inference_test.py | 1 + .../registry_pytorch_batch_inference_test.py | 1 + ...tence_transformers_batch_inference_test.py | 45 +++ ...egistry_tensorflow_batch_inference_test.py | 1 + .../snowflake/ml/registry/model/BUILD.bazel | 2 +- ...egistry_huggingface_pipeline_model_test.py | 6 +- ...registry_pip_only_model_deployment_test.py | 49 +-- 83 files changed, 4123 insertions(+), 346 deletions(-) create mode 100644 snowflake/ml/feature_store/_internal/__init__.py create mode 100644 snowflake/ml/feature_store/feature_store_initialization_warehouse_test.py create mode 100644 snowflake/ml/model/_client/service/inference_job_service_spec.py create mode 100644 snowflake/ml/model/_client/service/inference_job_service_spec_test.py create mode 100644 snowflake/ml/model/_model_composer/huggingface_lazy_uploader.py create mode 100644 snowflake/ml/model/_model_composer/huggingface_lazy_uploader_test.py create mode 100644 snowflake/ml/model/_packager/model_handlers_test/huggingface_utils_test.py diff --git a/CHANGELOG.md b/CHANGELOG.md index ca0182c6..94a1cb76 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,78 @@ # Release History -## 1.44.0 +## 1.45.0 + +### New Features + +* Feature Store: `FeatureView` now supports an `initialization_warehouse` that is used for the initial build and any + subsequent reinitializations of the backing dynamic table (a full scan of the source data), while `warehouse` + continues to drive the lighter incremental refreshes. This mirrors the dynamic table `INITIALIZATION_WAREHOUSE` + knob, lets you pair a larger warehouse for initialization with a smaller one for steady-state refresh, and is also + used for the one-time backfill of streaming feature views. It can be set at registration, changed via + `update_feature_view(initialization_warehouse=...)`, and is surfaced by `list_feature_views(verbose=True)`. + +```python +draft_fv = FeatureView( + name="F_TRIP", + entities=[entity], + feature_df=feature_df, + refresh_freq="1d", + warehouse="SMALL_WH", # incremental refreshes + initialization_warehouse="LARGE_WH", # initial build / reinitialization +) +fv = fs.register_feature_view(draft_fv, version="1.0") +``` + +* Registry: LLM models deployed with the OpenAI chat signatures now support structured outputs + through an optional `response_format` param matching the OpenAI Chat Completions API + (`{"type": "json_schema", "json_schema": {"name": "...", "schema": {...}}}`), letting callers + constrain model output to a JSON Schema. + +```python +from pydantic import BaseModel +import pandas as pd + +class CityCountry(BaseModel): + city: str + country: str + +response_format = { + "type": "json_schema", + "json_schema": { + "name": "city_country", + "schema": CityCountry.model_json_schema(), + }, +} + +x_df = pd.DataFrame.from_records( + [ + { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is the capital of France?"}, + ], + }, + ], + } + ] +) + +mv.run( + X=x_df, + params={"response_format": response_format}, + service_name=..., +) +``` + +### Bug Fixes + +### Behavior Changes + +### Deprecations + +## 1.44.0 (2026-06-23) ### New Features diff --git a/MODULE.bazel b/MODULE.bazel index b9670145..b8da5e7f 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -73,6 +73,7 @@ use_repo( "com_github_apache_arrow_go_v18", "com_github_buger_jsonparser", "com_github_caarlos0_env_v11", + "com_github_goccy_go_json", "com_github_snowflakedb_gosnowflake", "in_gopkg_natefinch_lumberjack_v2", "in_gopkg_yaml_v2", diff --git a/bazel/environments/conda-optional-dependency-llm.yml b/bazel/environments/conda-optional-dependency-llm.yml index a276b974..195a98fa 100755 --- a/bazel/environments/conda-optional-dependency-llm.yml +++ b/bazel/environments/conda-optional-dependency-llm.yml @@ -6,7 +6,7 @@ channels: - https://repo.anaconda.com/pkgs/snowflake - nodefaults dependencies: - - mlflow>=2.16.0, <4 + - mlflow>=2.16.0, <3 - pytorch>=2.0.1,<3 - sentence-transformers>=2.7.0,<6 - sentencepiece>=0.1.95,<0.3 diff --git a/bazel/environments/conda-optional-dependency-ml.yml b/bazel/environments/conda-optional-dependency-ml.yml index 4cb510e0..c0c1829f 100755 --- a/bazel/environments/conda-optional-dependency-ml.yml +++ b/bazel/environments/conda-optional-dependency-ml.yml @@ -9,6 +9,6 @@ dependencies: - altair>=5,<7 - catboost>=1.2.0, <2 - lightgbm>=4.1.0, <5 - - mlflow>=2.16.0, <4 + - mlflow>=2.16.0, <3 - prophet>=1.1.0, <2 - streamlit>=1.30.0,<2 diff --git a/bazel/environments/conda-optional-dependency-torch.yml b/bazel/environments/conda-optional-dependency-torch.yml index 5ae7b73e..9b5ed06b 100755 --- a/bazel/environments/conda-optional-dependency-torch.yml +++ b/bazel/environments/conda-optional-dependency-torch.yml @@ -6,7 +6,7 @@ channels: - https://repo.anaconda.com/pkgs/snowflake - nodefaults dependencies: - - mlflow>=2.16.0, <4 + - mlflow>=2.16.0, <3 - pytorch>=2.0.1,<3 - sentence-transformers>=2.7.0,<6 - sentencepiece>=0.1.95,<0.3 diff --git a/ci/conda_recipe/meta.yaml b/ci/conda_recipe/meta.yaml index 30c43f53..7da9ca54 100644 --- a/ci/conda_recipe/meta.yaml +++ b/ci/conda_recipe/meta.yaml @@ -17,7 +17,7 @@ build: noarch: python package: name: snowflake-ml-python - version: 1.44.0 + version: 1.45.0 requirements: build: - python @@ -58,7 +58,7 @@ requirements: - catboost>=1.2.0, <2 - keras>=2.0.0,<4 - lightgbm>=4.1.0, <5 - - mlflow>=2.16.0, <4 + - mlflow>=2.16.0, <3 - prophet>=1.1.0, <2 - pytorch>=2.0.1,<3 - sentence-transformers>=2.7.0,<6 diff --git a/codegen/sklearn_wrapper_generator.py b/codegen/sklearn_wrapper_generator.py index f1ec128f..022d8751 100644 --- a/codegen/sklearn_wrapper_generator.py +++ b/codegen/sklearn_wrapper_generator.py @@ -292,6 +292,7 @@ def _is_deterministic(class_object: tuple[str, type]) -> bool: WrapperGeneratorFactory._is_class_of_type(class_object[1], "LinearDiscriminantAnalysis") or WrapperGeneratorFactory._is_class_of_type(class_object[1], "BernoulliRBM") or WrapperGeneratorFactory._is_class_of_type(class_object[1], "TSNE") + or WrapperGeneratorFactory._is_class_of_type(class_object[1], "MDS") ) @staticmethod diff --git a/codegen/transformer_autogen_test_template.py_template b/codegen/transformer_autogen_test_template.py_template index c61dd59b..29668452 100644 --- a/codegen/transformer_autogen_test_template.py_template +++ b/codegen/transformer_autogen_test_template.py_template @@ -85,16 +85,6 @@ class {transform.test_class_name}(TestCase): fit_with_sproc: bool = True, inference_with_udf: bool = True ) -> None: - # sklearn _parallel_pairwise bug (https://github.com/scikit-learn/scikit-learn/issues/33877): - # when effective_n_jobs(None) > 1 (warehouse environment), Y is chunked but Y_norm_squared - # is passed through unchanged, causing a dimension mismatch in Birch._predict. - # Fixed in https://github.com/scikit-learn/scikit-learn/pull/33876 (sklearn 1.9.0). - # TODO(SNOW-952252): Re-enable when test env sklearn is bumped to >= 1.9.0. - if "{transform.original_class_name}" == "Birch" and (fit_with_sproc or inference_with_udf): # type: ignore[comparison-overlap,unused-ignore] - import joblib - if joblib.effective_n_jobs(None) > 1: - self.skipTest("sklearn _parallel_pairwise bug causes dimension mismatch when n_jobs > 1") - input_df_pandas = {transform.test_dataset_func}(as_frame=True).frame cols = [inflection.parameterize(c, "_").upper() for c in input_df_pandas.columns if not c.startswith("target")] cols_half_1, cols_half_2 = cols[:int(len(cols)/2)], cols[int(len(cols)/2)+1:] diff --git a/requirements.yml b/requirements.yml index a2f5ed0d..910514b5 100644 --- a/requirements.yml +++ b/requirements.yml @@ -150,7 +150,7 @@ version_requirements: '>=0.46.0,<1' - name: mlflow dev_version: 2.16.2 - version_requirements: '>=2.16.0, <4' + version_requirements: '>=2.16.0, <3' requirements_extra_tags: - mlflow - llm diff --git a/snowflake/ml/_internal/file_utils.py b/snowflake/ml/_internal/file_utils.py index eb16d379..04463e9a 100644 --- a/snowflake/ml/_internal/file_utils.py +++ b/snowflake/ml/_internal/file_utils.py @@ -313,6 +313,38 @@ def _resolve_stage_dir_path( return str((stage_path / relative_path).parent) +def upload_file_to_stage( + session: snowpark.Session, + local_file_path: Union[str, pathlib.Path], + stage_dir_path: str, + *, + statement_params: Optional[dict[str, Any]] = None, +) -> None: + """Upload a single local file to a stage directory. + + Args: + session: Snowpark Session. + local_file_path: Local path of the file to upload. + stage_dir_path: Destination directory path in the stage. + statement_params: Statement Params. + """ + import retrying + + file_operation = snowpark.FileOperation(session=session) + retrying.retry( + retry_on_exception=_retry_on_sql_error, + stop_max_attempt_number=5, + wait_exponential_multiplier=100, + wait_exponential_max=10000, + )(file_operation.put)( + str(local_file_path), + str(stage_dir_path), + auto_compress=False, + overwrite=False, + statement_params=statement_params, + ) + + def upload_directory_to_stage( session: snowpark.Session, local_path: pathlib.Path, @@ -328,27 +360,16 @@ def upload_directory_to_stage( stage_path: Base path in the stage. statement_params: Statement Params. """ - import retrying - - file_operation = snowpark.FileOperation(session=session) - for root, _, filenames in os.walk(local_path): root_path = pathlib.Path(root) for filename in filenames: local_file_path = root_path / filename relative_path = pathlib.PurePosixPath(local_file_path.relative_to(local_path).as_posix()) stage_dir_path = _resolve_stage_dir_path(stage_path, relative_path) - - retrying.retry( - retry_on_exception=_retry_on_sql_error, - stop_max_attempt_number=5, - wait_exponential_multiplier=100, - wait_exponential_max=10000, - )(file_operation.put)( - str(local_file_path), - str(stage_dir_path), - auto_compress=False, - overwrite=False, + upload_file_to_stage( + session, + local_file_path, + stage_dir_path, statement_params=statement_params, ) diff --git a/snowflake/ml/experiment/_client/experiment_tracking_sql_client.py b/snowflake/ml/experiment/_client/experiment_tracking_sql_client.py index b21d569d..b64d659c 100644 --- a/snowflake/ml/experiment/_client/experiment_tracking_sql_client.py +++ b/snowflake/ml/experiment/_client/experiment_tracking_sql_client.py @@ -118,11 +118,20 @@ def drop_experiment( ).validate() @telemetry.send_api_usage_telemetry(project=telemetry.TelemetryProject.EXPERIMENT_TRACKING.value) - def add_run(self, *, experiment_name: sql_identifier.SqlIdentifier, run_name: sql_identifier.SqlIdentifier) -> None: + def add_run( + self, + *, + experiment_name: sql_identifier.SqlIdentifier, + run_name: sql_identifier.SqlIdentifier, + source_info: Optional[str] = None, + ) -> None: experiment_fqn = self.fully_qualified_object_name(self._database_name, self._schema_name, experiment_name) - query_result_checker.SqlResultValidator( - self._session, f"ALTER EXPERIMENT {experiment_fqn} ADD RUN {run_name}" - ).has_dimensions(expected_rows=1, expected_cols=1).validate() + query = f"ALTER EXPERIMENT {experiment_fqn} ADD RUN {run_name}" + if source_info: + query += f" WITH (SOURCE_INFO = $${source_info}$$)" + query_result_checker.SqlResultValidator(self._session, query).has_dimensions( + expected_rows=1, expected_cols=1 + ).validate() @telemetry.send_api_usage_telemetry(project=telemetry.TelemetryProject.EXPERIMENT_TRACKING.value) def commit_run( diff --git a/snowflake/ml/experiment/_client/experiment_tracking_sql_client_test.py b/snowflake/ml/experiment/_client/experiment_tracking_sql_client_test.py index 102801cf..5158249e 100644 --- a/snowflake/ml/experiment/_client/experiment_tracking_sql_client_test.py +++ b/snowflake/ml/experiment/_client/experiment_tracking_sql_client_test.py @@ -169,6 +169,38 @@ def test_add_run(self) -> None: ) self.client.add_run(experiment_name=experiment_name, run_name=run_name) + def test_add_run_with_source_info(self) -> None: + # The caller passes an already-serialized JSON string; the client embeds it verbatim. + experiment_name = sql_identifier.SqlIdentifier("TEST_EXPERIMENT") + run_name = sql_identifier.SqlIdentifier("TEST_RUN") + + payload = '{"entry_point": "train/main.py", "git": {"commit_hash": "abc"}}' + self.m_session.add_mock_sql( + "ALTER EXPERIMENT TEST_DB.TEST_SCHEMA.TEST_EXPERIMENT ADD RUN TEST_RUN" + f" WITH (SOURCE_INFO = $${payload}$$)", + self._create_mock_df(), + ) + self.client.add_run(experiment_name=experiment_name, run_name=run_name, source_info=payload) + + def test_add_run_with_none_source_info_omits_clause(self) -> None: + experiment_name = sql_identifier.SqlIdentifier("TEST_EXPERIMENT") + run_name = sql_identifier.SqlIdentifier("TEST_RUN") + + self.m_session.add_mock_sql( + "ALTER EXPERIMENT TEST_DB.TEST_SCHEMA.TEST_EXPERIMENT ADD RUN TEST_RUN", self._create_mock_df() + ) + self.client.add_run(experiment_name=experiment_name, run_name=run_name, source_info=None) + + def test_add_run_with_empty_source_info_omits_clause(self) -> None: + # An empty string is falsy and must be treated like "nothing to send". + experiment_name = sql_identifier.SqlIdentifier("TEST_EXPERIMENT") + run_name = sql_identifier.SqlIdentifier("TEST_RUN") + + self.m_session.add_mock_sql( + "ALTER EXPERIMENT TEST_DB.TEST_SCHEMA.TEST_EXPERIMENT ADD RUN TEST_RUN", self._create_mock_df() + ) + self.client.add_run(experiment_name=experiment_name, run_name=run_name, source_info="") + def test_commit_run(self) -> None: # Test committing a run experiment_name = sql_identifier.SqlIdentifier("TEST_EXPERIMENT") diff --git a/snowflake/ml/feature_store/BUILD.bazel b/snowflake/ml/feature_store/BUILD.bazel index d8a8b182..ac1623b5 100644 --- a/snowflake/ml/feature_store/BUILD.bazel +++ b/snowflake/ml/feature_store/BUILD.bazel @@ -274,6 +274,18 @@ py_test( ], ) +py_test( + name = "feature_store_initialization_warehouse_test", + srcs = ["feature_store_initialization_warehouse_test.py"], + tags = [ + "feature:feature_store", + "short_regress", + ], + deps = [ + ":feature_store_lib", + ], +) + py_test( name = "realtime_config_test", srcs = ["realtime_config_test.py"], diff --git a/snowflake/ml/feature_store/_internal/__init__.py b/snowflake/ml/feature_store/_internal/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/snowflake/ml/feature_store/feature_store.py b/snowflake/ml/feature_store/feature_store.py index 467f9b4e..a8258879 100644 --- a/snowflake/ml/feature_store/feature_store.py +++ b/snowflake/ml/feature_store/feature_store.py @@ -330,6 +330,27 @@ class _MaterializedResourceKind(Enum): flags=re.DOTALL | re.IGNORECASE | re.X, ) +# Sentinel for update_feature_view: distinguishes "leave unchanged" (default) +# from an explicit None, which clears (UNSETs) the initialization warehouse. +_KEEP_CURRENT: Any = object() + + +def _initialization_warehouse_clause(feature_view: FeatureView) -> str: + """Build the trailing ``INITIALIZATION_WAREHOUSE = ...`` fragment for CREATE DYNAMIC TABLE. + + Returns an empty string when the feature view has no initialization + warehouse, leaving the DDL identical to the single-warehouse behavior. + + Args: + feature_view: The feature view being materialized. + + Returns: + ``\\n INITIALIZATION_WAREHOUSE=`` or an empty string. + """ + iw = feature_view.initialization_warehouse + return f"\n INITIALIZATION_WAREHOUSE = {iw}" if iw is not None else "" + + # ``list_feature_views`` output schemas. Row builders always populate the verbose # schema; non-verbose callers receive _LIST_FEATURE_VIEW_SCHEMA after trimming # verbose-only trailing fields. @@ -356,6 +377,9 @@ class _MaterializedResourceKind(Enum): ] _LIST_FEATURE_VIEW_VERBOSE_EXTRA_FIELDS = [ + # Initialization warehouse used for the initial build / reinitialization of a + # managed FV's dynamic table. Verbose-only: ``None`` for FVs without one. + StructField("initialization_warehouse", StringType()), # JSON-encoded authored source-ref list, or ``None`` when no # ``FV_SOURCE_REFS`` metadata row was written for this FV. StructField("source_refs", StringType()), @@ -1261,6 +1285,7 @@ def update_feature_view( *, refresh_freq: Optional[str] = _UNSET, warehouse: Optional[str] = None, + initialization_warehouse: Optional[str] = _KEEP_CURRENT, desc: Optional[str] = None, online_config: Optional[fv_mod.OnlineConfig] = None, updated_feature_df: Optional[DataFrame] = None, @@ -1275,6 +1300,7 @@ def update_feature_view( *, refresh_freq: Optional[str] = _UNSET, warehouse: Optional[str] = None, + initialization_warehouse: Optional[str] = _KEEP_CURRENT, desc: Optional[str] = None, online_config: Optional[fv_mod.OnlineConfig] = None, updated_feature_df: Optional[DataFrame] = None, @@ -1289,6 +1315,7 @@ def update_feature_view( *, refresh_freq: Optional[str] = _UNSET, warehouse: Optional[str] = None, + initialization_warehouse: Optional[str] = _KEEP_CURRENT, desc: Optional[str] = None, online_config: Optional[fv_mod.OnlineConfig] = None, updated_feature_df: Optional[DataFrame] = None, @@ -1301,6 +1328,10 @@ def update_feature_view( version: Optional version of feature view. Must set when argument feature_view is a str. refresh_freq: updated refresh frequency. warehouse: updated warehouse. + initialization_warehouse: updated initialization warehouse, used for the initial build and + reinitializations of the backing dynamic table. Pass a warehouse name to set it, or ``None`` to clear + it (all refreshes then run on ``warehouse``). When omitted, the existing value is left unchanged. + Not supported for static feature views. desc: description of feature view. online_config: updated online configuration for the online feature table. If provided with enable=True, creates online feature table if absent. @@ -1380,6 +1411,8 @@ def update_feature_view( actual_refresh_freq: Optional[str] = refresh_freq if refresh_freq is not _UNSET else None new_desc = desc if desc is not None else feature_view.desc + init_wh_changed = initialization_warehouse is not _KEEP_CURRENT + if refresh_freq is not _UNSET: # Validate the prospective refresh_freq against the snapshot-config contract # without mutating the registered FV. Shallow copy is sufficient: the update @@ -1392,15 +1425,31 @@ def update_feature_view( feature_view_append_only_validation.validate_snapshot_config_for_update(validation_fv) # Validate static feature view constraints - if feature_view.status == FeatureViewStatus.STATIC and (actual_refresh_freq or warehouse): + if feature_view.status == FeatureViewStatus.STATIC and (actual_refresh_freq or warehouse or init_wh_changed): full_name = f"{feature_view.name}/{feature_view.version}" raise snowml_exceptions.SnowflakeMLException( error_code=error_codes.INVALID_ARGUMENT, original_exception=RuntimeError( - f"Static feature view '{full_name}' does not support refresh_freq and warehouse." + f"Static feature view '{full_name}' does not support refresh_freq, warehouse, " + "and initialization_warehouse." ), ) + # Fail fast on an online store type that cannot back this feature view (e.g. HYBRID_TABLE + # for a tiled FV) before any planning/resource work, so it surfaces as a clean + # INVALID_ARGUMENT rather than being wrapped by the update-failure rollback handler. The + # rule is a feature-view invariant; check the target config on a copy since the live rebuild + # path does not re-run __init__ validation. + if online_config is not None: + probe_fv = feature_view.copy() + probe_fv._online_config = online_config + try: + probe_fv._validate_online_store_supported() + except ValueError as e: + raise snowml_exceptions.SnowflakeMLException( + error_code=error_codes.INVALID_ARGUMENT, original_exception=e + ) from e + if updated_feature_df is not None: if not feature_view.append_only: raise snowml_exceptions.SnowflakeMLException( @@ -1482,7 +1531,12 @@ def update_feature_view( rollback_operations: list[Any] = [] try: operations, rollback_operations = self._plan_feature_view_update_operations( - feature_view, actual_refresh_freq, warehouse, new_desc, online_config + feature_view, + actual_refresh_freq, + warehouse, + initialization_warehouse, + new_desc, + online_config, ) # Step 3: Execute atomically @@ -1773,18 +1827,20 @@ def list_feature_views( Args: entity_name: Entity name. feature_view_name: FeatureView name. - verbose: When True, include the ``source_refs`` and ``backup_source`` columns - in the output. ``source_refs`` (string, nullable) is the JSON-encoded list - of authored source bindings captured at registration time (``None`` when no - source refs were recorded). ``backup_source`` (string, nullable) contains - the fully-qualified name of the historical snapshot table cloned at + verbose: When True, include the ``initialization_warehouse``, ``source_refs`` and + ``backup_source`` columns in the output. ``initialization_warehouse`` (string, + nullable) is the warehouse used for the initial build / reinitialization of the + backing dynamic table (``None`` when unset). ``source_refs`` (string, nullable) is + the JSON-encoded list of authored source bindings captured at registration time + (``None`` when no source refs were recorded). ``backup_source`` (string, nullable) + contains the fully-qualified name of the historical snapshot table cloned at registration time (only set for append-only feature views registered with a ``backup_source``; ``None`` otherwise). Defaults to False. Returns: FeatureViews information as a Snowpark DataFrame. Each row always includes - ``append_only`` (bool). ``source_refs`` and ``backup_source`` (both string, - nullable) are included only when ``verbose=True``. + ``append_only`` (bool). ``initialization_warehouse``, ``source_refs`` and + ``backup_source`` (all string, nullable) are included only when ``verbose=True``. Example:: @@ -3825,48 +3881,31 @@ def _rollback_created_resources(self, created_resources: list[tuple[_FeatureStor def _create_updated_feature_view( self, base_fv: FeatureView, online_config: Optional[fv_mod.OnlineConfig] = None ) -> FeatureView: - """Create an updated FeatureView with new online configuration.""" - assert base_fv.version is not None - assert base_fv.database is not None - assert base_fv.schema is not None - assert base_fv.feature_df is not None - - feature_descs_str: Optional[dict[str, str]] = ( - {k.identifier(): v for k, v in base_fv.feature_descs.items()} if base_fv.feature_descs is not None else None - ) - cluster_by_str: Optional[list[str]] = ( - [col.identifier() for col in base_fv.cluster_by] if base_fv.cluster_by is not None else None - ) + """Return a copy of ``base_fv`` with a new online configuration applied. - return FeatureView._construct_feature_view( - name=base_fv.name.identifier(), - entities=base_fv.entities, - feature_df=base_fv.feature_df, - timestamp_col=(base_fv.timestamp_col.identifier() if base_fv.timestamp_col is not None else None), - desc=base_fv.desc, - version=str(base_fv.version), - status=base_fv.status, - feature_descs=feature_descs_str or {}, - refresh_freq=base_fv.refresh_freq, - database=base_fv.database.identifier(), - schema=base_fv.schema.identifier(), - warehouse=(base_fv.warehouse.identifier() if base_fv.warehouse is not None else None), - refresh_mode=base_fv.refresh_mode, - refresh_mode_reason=base_fv.refresh_mode_reason, - initialize=base_fv.initialize, - owner=base_fv.owner, - infer_schema_df=base_fv._infer_schema_df, - session=self._session, - cluster_by=cluster_by_str, - online_config=online_config, - storage_config=base_fv.storage_config, - ) + A shallow copy preserves the full feature-view identity (tiled / streaming / realtime / + append-only / source-refs) automatically, so only the online config changes. Hand-copying + individual fields here would silently drop any field not listed and drift as new fields are + added (the same class of bug this avoids). Mirrors the ``copy() + _online_config`` pattern + used elsewhere in ``update_feature_view``. + + Args: + base_fv: The feature view to copy. + online_config: The online configuration to apply to the copy. + + Returns: + A shallow copy of ``base_fv`` with ``online_config`` applied. + """ + fv = base_fv.copy() + fv._online_config = online_config + return fv def _build_offline_update_queries( self, feature_view: FeatureView, refresh_freq: Optional[str], warehouse: Optional[str], + initialization_warehouse: Optional[str], desc: str, ) -> tuple[list[tuple[str, str]], list[tuple[str, str]]]: """Build offline update operations and their rollback operations. @@ -3897,6 +3936,9 @@ def _build_offline_update_queries( ``None`` to keep the existing value. warehouse: New warehouse identifier requested by the caller, or ``None`` to keep the existing value. + initialization_warehouse: New initialization warehouse. ``_KEEP_CURRENT`` + leaves it unchanged; ``None`` clears it (``UNSET``); any other + value sets it (``SET``). desc: New description to set on the DT/view. Returns: @@ -3992,6 +4034,20 @@ def alter_task_schedule_ops(op_type: str, *, cron_expr: str) -> list[tuple[str, elif not old_is_cron and new_is_cron: rollback_ops.append(drop_task_op("OFFLINE_ROLLBACK")) + # INITIALIZATION_WAREHOUSE is a standalone SET/UNSET. _KEEP_CURRENT means + # the caller didn't pass the argument; an explicit None clears it. + if initialization_warehouse is not _KEEP_CURRENT: + old_init_wh = feature_view.initialization_warehouse + + def set_or_unset_init_wh(op_type: str, value: Optional[SqlIdentifier]) -> tuple[str, str]: + if value is None: + return (op_type, f"ALTER DYNAMIC TABLE {fqn} UNSET INITIALIZATION_WAREHOUSE") + return (op_type, f"ALTER DYNAMIC TABLE {fqn} SET INITIALIZATION_WAREHOUSE = {value}") + + new_init_wh = SqlIdentifier(initialization_warehouse) if initialization_warehouse is not None else None + operations.append(set_or_unset_init_wh("OFFLINE_UPDATE", new_init_wh)) + rollback_ops.append(set_or_unset_init_wh("OFFLINE_ROLLBACK", old_init_wh)) + return operations, rollback_ops @dataclass(frozen=True) @@ -4063,6 +4119,7 @@ def _plan_online_enable( final_config = fv_mod.OnlineConfig( enable=True, target_lag=(online_config.target_lag if online_config.target_lag is not None else default_target_lag), + store_type=online_config.store_type, ) temp_fv = self._create_updated_feature_view(feature_view, final_config) @@ -4120,6 +4177,7 @@ def _plan_feature_view_update_operations( feature_view: FeatureView, refresh_freq: Optional[str], warehouse: Optional[str], + initialization_warehouse: Optional[str], desc: str, online_config: Optional[fv_mod.OnlineConfig], ) -> tuple[list[tuple[str, Union[str, FeatureView]]], list[tuple[str, Union[str, FeatureView]]],]: @@ -4129,7 +4187,7 @@ def _plan_feature_view_update_operations( # Plan offline updates offline_ops, offline_rollback_ops = self._build_offline_update_queries( - feature_view, refresh_freq, warehouse, desc + feature_view, refresh_freq, warehouse, initialization_warehouse, desc ) operations.extend(offline_ops) rollback_operations.extend(offline_rollback_ops) @@ -5360,6 +5418,7 @@ def _create_dynamic_table_query( # Include column definitions only if provided (skip for tiled feature views) column_clause = f" ({column_descs})" if column_descs else "" + init_wh_clause = _initialization_warehouse_clause(feature_view) storage_config = feature_view.storage_config if storage_config is not None and storage_config.format == StorageFormat.ICEBERG: # These should be validated by FeatureView constructor and _resolve_storage_config @@ -5371,7 +5430,7 @@ def _create_dynamic_table_query( TAG ( {tagging_clause} ) - WAREHOUSE = {warehouse} + WAREHOUSE = {warehouse}{init_wh_clause} REFRESH_MODE = {feature_view.refresh_mode} INITIALIZE = {feature_view.initialize} CATALOG = 'SNOWFLAKE' @@ -5385,7 +5444,7 @@ def _create_dynamic_table_query( TAG ( {tagging_clause} ) - WAREHOUSE = {warehouse} + WAREHOUSE = {warehouse}{init_wh_clause} REFRESH_MODE = {feature_view.refresh_mode} INITIALIZE = {feature_view.initialize} """ @@ -6025,7 +6084,7 @@ def _create_rollup_feature_view( TAG ( {tagging_clause_str} ) - WAREHOUSE = {warehouse} + WAREHOUSE = {warehouse}{_initialization_warehouse_clause(feature_view)} REFRESH_MODE = {feature_view.refresh_mode} INITIALIZE = {feature_view.initialize} """ @@ -7556,6 +7615,8 @@ def _extract_feature_view_info( values.append(row["scheduling_state"] if "scheduling_state" in row else None) values.append(row["warehouse"] if "warehouse" in row else None) values.append(json.dumps(self._extract_cluster_by_columns(row["cluster_by"])) if "cluster_by" in row else None) + # Verbose-only field, emitted into output_values_extra below. + initialization_warehouse = row["initialization_warehouse"] if "initialization_warehouse" in row else None online_config_json = self._determine_online_config_from_oft(name, version, include_online_service_metadata=True) values.append(online_config_json) @@ -7625,7 +7686,7 @@ def _extract_feature_view_info( append_only_meta = self._metadata_manager.get_append_only_metadata(name, version) if append_only_meta is not None: backup_source = AppendOnlyMetadata.from_dict(append_only_meta).backup_source - output_values_extra.append([source_refs_json, backup_source]) + output_values_extra.append([initialization_warehouse, source_refs_json, backup_source]) def _determine_online_config_from_oft( self, name: str, version: str, *, include_online_service_metadata: bool = False @@ -7810,6 +7871,12 @@ def find_and_compose_entity(name: str) -> Entity: timestamp_col = ts_col if ts_col not in _LEGACY_TIMESTAMP_COL_PLACEHOLDER_VALS else None re_initialize = re.match(_DT_INITIALIZE_PATTERN, row["text"]) initialize = re_initialize.group("initialize") if re_initialize is not None else "ON_CREATE" + # INITIALIZATION_WAREHOUSE is surfaced as a dedicated SHOW column (like WAREHOUSE), + # not echoed in the DDL text, so read it from the column. + init_wh_value = row["initialization_warehouse"] if "initialization_warehouse" in row else None + initialization_warehouse = ( + SqlIdentifier(init_wh_value, case_sensitive=True).identifier() if init_wh_value else None + ) # For tiled FVs, get descriptions from metadata table; otherwise from DT columns if is_tiled: @@ -7851,6 +7918,7 @@ def find_and_compose_entity(name: str) -> Entity: if len(row["warehouse"]) > 0 else None ), + initialization_warehouse=initialization_warehouse, refresh_mode=row["refresh_mode"], refresh_mode_reason=row["refresh_mode_reason"], initialize=initialize, @@ -8050,6 +8118,15 @@ def _create_online_feature_table( SnowflakeMLException: [ValueError] If OnlineConfig is required but not provided. SnowflakeMLException: If creating the online feature table fails. """ + # Defense-in-depth: also covers the registration path. Rejects store types that cannot back + # a tiled feature view before any DDL is issued. The rule lives on FeatureView. + try: + feature_view._validate_online_store_supported() + except ValueError as e: + raise snowml_exceptions.SnowflakeMLException( + error_code=error_codes.INVALID_ARGUMENT, original_exception=e + ) from e + online_table_name = FeatureView._get_online_table_name(feature_view_name) fully_qualified_online_name = self._get_fully_qualified_name(online_table_name) diff --git a/snowflake/ml/feature_store/feature_store_entity_test.py b/snowflake/ml/feature_store/feature_store_entity_test.py index b41feb82..55214cdb 100644 --- a/snowflake/ml/feature_store/feature_store_entity_test.py +++ b/snowflake/ml/feature_store/feature_store_entity_test.py @@ -214,7 +214,11 @@ def test_static_branch_escapes_desc(self) -> None: fv = _make_feature_view_mock(desc="ignored", status=FeatureViewStatus.STATIC, refresh_freq=None) ops, _rollback = fs._build_offline_update_queries( - feature_view=fv, refresh_freq=None, warehouse=None, desc=_EXPLOIT_DESC + feature_view=fv, + refresh_freq=None, + warehouse=None, + initialization_warehouse=fs_mod._KEEP_CURRENT, + desc=_EXPLOIT_DESC, ) self.assertEqual(len(ops), 1) @@ -228,7 +232,11 @@ def test_dt_branch_escapes_desc(self) -> None: fv = _make_feature_view_mock(desc=_EXPLOIT_DESC, refresh_freq="1 minute") ops, rollback = fs._build_offline_update_queries( - feature_view=fv, refresh_freq=None, warehouse=None, desc=_EXPLOIT_DESC + feature_view=fv, + refresh_freq=None, + warehouse=None, + initialization_warehouse=fs_mod._KEEP_CURRENT, + desc=_EXPLOIT_DESC, ) alter_dt_sqls = [sql for _op, sql in ops + rollback if "ALTER DYNAMIC TABLE" in sql] diff --git a/snowflake/ml/feature_store/feature_store_initialization_warehouse_test.py b/snowflake/ml/feature_store/feature_store_initialization_warehouse_test.py new file mode 100644 index 00000000..b1278c2f --- /dev/null +++ b/snowflake/ml/feature_store/feature_store_initialization_warehouse_test.py @@ -0,0 +1,89 @@ +"""Unit tests for the ``initialization_warehouse`` plumbing in feature_store.""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import MagicMock + +from absl.testing import absltest + +from snowflake.ml.feature_store import feature_store as fs_mod +from snowflake.ml.feature_store.entity import Entity +from snowflake.ml.feature_store.feature_view import ( + FeatureView, + StorageConfig, + StorageFormat, +) + + +def _make_fv(**kwargs: Any) -> FeatureView: + mock_df = MagicMock() + mock_df.columns = ["user_id", "amount"] + mock_df.queries = {"queries": ["SELECT * FROM source"]} + return FeatureView( + name="test_fv", + entities=[Entity(name="user", join_keys=["user_id"])], + feature_df=mock_df, + refresh_freq="1d", + warehouse="small_wh", + **kwargs, + ) + + +class InitializationWarehouseClauseTest(absltest.TestCase): + """The clause builder emits CREATE-time DDL; read-back uses the SHOW column.""" + + def test_clause_empty_when_unset(self) -> None: + fv = _make_fv() + self.assertEqual(fs_mod._initialization_warehouse_clause(fv), "") + + def test_clause_present_when_set(self) -> None: + fv = _make_fv(initialization_warehouse="large_wh") + clause = fs_mod._initialization_warehouse_clause(fv) + self.assertIn("INITIALIZATION_WAREHOUSE = LARGE_WH", clause) + + +class CreateDynamicTableQueryTest(absltest.TestCase): + """``_create_dynamic_table_query`` does not touch session state, so it can be + exercised on a bare FeatureStore instance.""" + + def setUp(self) -> None: + self._fs = object.__new__(fs_mod.FeatureStore) + + def _build(self, fv: FeatureView) -> str: + return self._fs._create_dynamic_table_query( + override_clause="", + table_name="DB.SCH.TEST_FV$V1", + column_descs="", + schedule_task=False, + feature_view=fv, + tagging_clause="TAG_A = 'x'", + warehouse="SMALL_WH", + ) + + def test_dynamic_table_emits_initialization_warehouse_when_set(self) -> None: + query = self._build(_make_fv(initialization_warehouse="large_wh")) + self.assertIn("WAREHOUSE = SMALL_WH", query) + self.assertIn("INITIALIZATION_WAREHOUSE = LARGE_WH", query) + + def test_dynamic_table_omits_clause_when_unset(self) -> None: + query = self._build(_make_fv()) + self.assertIn("WAREHOUSE = SMALL_WH", query) + self.assertNotIn("INITIALIZATION_WAREHOUSE", query) + + def test_iceberg_table_emits_initialization_warehouse_when_set(self) -> None: + fv = _make_fv( + initialization_warehouse="large_wh", + storage_config=StorageConfig( + format=StorageFormat.ICEBERG, + external_volume="MY_VOLUME", + base_location="feature_store/test_fv", + ), + ) + query = self._build(fv) + self.assertIn("DYNAMIC ICEBERG TABLE", query) + self.assertIn("INITIALIZATION_WAREHOUSE = LARGE_WH", query) + + +if __name__ == "__main__": + absltest.main() diff --git a/snowflake/ml/feature_store/feature_store_oft_online_config_test.py b/snowflake/ml/feature_store/feature_store_oft_online_config_test.py index ec96660c..80130a16 100644 --- a/snowflake/ml/feature_store/feature_store_oft_online_config_test.py +++ b/snowflake/ml/feature_store/feature_store_oft_online_config_test.py @@ -5,12 +5,56 @@ from typing import Optional from unittest.mock import MagicMock +import pandas as pd from absl.testing import absltest, parameterized +from snowflake.ml._internal.utils.sql_identifier import SqlIdentifier from snowflake.ml.feature_store import feature_store as fs_mod +from snowflake.ml.feature_store.entity import Entity +from snowflake.ml.feature_store.feature import Feature from snowflake.ml.feature_store.feature_store import FeatureStore -from snowflake.ml.feature_store.feature_view import OnlineConfig, OnlineStoreType +from snowflake.ml.feature_store.feature_view import ( + FeatureView, + FeatureViewStatus, + FeatureViewVersion, + OnlineConfig, + OnlineStoreType, +) +from snowflake.ml.feature_store.realtime_config import RealtimeConfig +from snowflake.ml.feature_store.spec.enums import FeatureAggregationMethod from snowflake.snowpark import Row +from snowflake.snowpark.types import ( + DoubleType, + StringType, + StructField, + StructType, + TimestampType, +) + +_RTFV_OUTPUT_SCHEMA = StructType([StructField("risk_score", DoubleType())]) + + +def _rtfv_compute_fn(txn: pd.DataFrame) -> pd.DataFrame: + return pd.DataFrame({"risk_score": txn["avg_amount"]}) + + +def _build_rtfv_upstream_fv() -> FeatureView: + """Build a registered-looking FV to use as an RTFV source.""" + schema = StructType([StructField("USER_ID", StringType()), StructField("avg_amount", DoubleType())]) + mock_df = MagicMock() + mock_df.columns = [f.name for f in schema.fields] + mock_df.schema = schema + mock_df.queries = {"queries": ["SELECT * FROM TXN_FV"]} + fv = FeatureView( + name="TXN_FV", + entities=[Entity(name="USER", join_keys=["USER_ID"])], + feature_df=mock_df, + online_config=OnlineConfig(enable=True, store_type=OnlineStoreType.POSTGRES), + ) + fv._version = FeatureViewVersion("v1") + fv._infer_schema_df = mock_df + fv._status = FeatureViewStatus.ACTIVE + return fv class StoreTypeFromOftShowRowTest(parameterized.TestCase): @@ -129,5 +173,187 @@ def fake_find_object(*, object_type: str, object_name: object) -> list[Row]: self.assertEqual(cfg.store_type, OnlineStoreType.POSTGRES) +class UpdateFeatureViewPreservesTiledIdentityTest(absltest.TestCase): + """``_create_updated_feature_view`` must preserve the tiled/streaming aggregation identity. + + ``update_feature_view`` (online enable/disable) rebuilds the FV via + ``_create_updated_feature_view`` and uses the result to create the Online Feature Table. + If the rebuild drops the aggregation config, the OFT is built from a non-tiled view of a + tiled FV. + """ + + def _make_mock_df(self) -> MagicMock: + df = MagicMock() + df.queries = {"queries": ["SELECT * FROM TBL"]} + df.columns = ["USER_ID", "EVENT_TIME", "AMOUNT"] + ts_field = MagicMock() + ts_field.datatype = TimestampType() + df.schema.__getitem__ = lambda _self, key: ts_field + return df + + def _make_tiled_streaming_base_fv(self) -> FeatureView: + """Build a reconstructed tiled streaming CONTINUOUS FV, as get_feature_view would.""" + return FeatureView._construct_feature_view( + name="TILED_FV", + entities=[Entity(name="user", join_keys=["USER_ID"])], + feature_df=self._make_mock_df(), + timestamp_col="EVENT_TIME", + desc="", + version="v1", + status=FeatureViewStatus.ACTIVE, + feature_descs={}, + refresh_freq="1 minute", + database="DB", + schema="SCH", + warehouse="WH", + refresh_mode="FULL", + refresh_mode_reason=None, + initialize="ON_CREATE", + owner=None, + infer_schema_df=None, + session=MagicMock(), + feature_granularity="1m", + aggregation_specs=[Feature.sum("AMOUNT", "30s").to_spec()], + feature_aggregation_method=FeatureAggregationMethod.CONTINUOUS, + is_streaming=True, + ) + + def test_create_updated_feature_view_preserves_tiled_streaming_identity(self) -> None: + base_fv = self._make_tiled_streaming_base_fv() + # Sanity: the base FV is tiled/streaming CONTINUOUS to begin with. + self.assertTrue(base_fv.is_tiled) + self.assertTrue(base_fv.is_streaming) + + fs = object.__new__(FeatureStore) + object.__setattr__(fs, "_session", MagicMock()) + + updated_fv = fs._create_updated_feature_view(base_fv, OnlineConfig(enable=True)) + + self.assertTrue(updated_fv.is_tiled, "updated FV lost tiled identity") + self.assertTrue(updated_fv.is_streaming, "updated FV lost streaming identity") + self.assertEqual(updated_fv.feature_granularity, "1m") + self.assertEqual(updated_fv.feature_aggregation_method, FeatureAggregationMethod.CONTINUOUS) + + def test_create_updated_feature_view_preserves_append_only(self) -> None: + base_fv = FeatureView._construct_feature_view( + name="APPEND_FV", + entities=[Entity(name="user", join_keys=["USER_ID"])], + feature_df=self._make_mock_df(), + timestamp_col="EVENT_TIME", + desc="", + version="v1", + status=FeatureViewStatus.ACTIVE, + feature_descs={}, + refresh_freq="0 0 * * * UTC", + database="DB", + schema="SCH", + warehouse="WH", + refresh_mode="FULL", + refresh_mode_reason=None, + initialize="ON_SCHEDULE", + owner=None, + infer_schema_df=None, + session=MagicMock(), + append_only=True, + backup_source="DB.SCH.HISTORY", + ) + self.assertTrue(base_fv.append_only) + + fs = object.__new__(FeatureStore) + object.__setattr__(fs, "_session", MagicMock()) + + updated_fv = fs._create_updated_feature_view(base_fv, OnlineConfig(enable=True)) + + self.assertTrue(updated_fv.append_only, "updated FV lost append_only identity") + self.assertEqual(updated_fv.backup_source, "DB.SCH.HISTORY") + + def test_create_updated_feature_view_preserves_realtime(self) -> None: + rtc = RealtimeConfig( + compute_fn=_rtfv_compute_fn, + sources=[_build_rtfv_upstream_fv()], + output_schema=_RTFV_OUTPUT_SCHEMA, + ) + base_fv = FeatureView._construct_feature_view( + name="RTFV", + entities=[Entity(name="user", join_keys=["USER_ID"])], + feature_df=None, + timestamp_col=None, + desc="", + version="v1", + status=FeatureViewStatus.ACTIVE, + feature_descs={}, + refresh_freq=None, + database="DB", + schema="SCH", + warehouse=None, + refresh_mode=None, + refresh_mode_reason=None, + initialize="ON_CREATE", + owner=None, + infer_schema_df=None, + session=MagicMock(), + is_realtime=True, + realtime_config=rtc, + ) + self.assertTrue(base_fv.is_realtime_feature_view) + + fs = object.__new__(FeatureStore) + object.__setattr__(fs, "_session", MagicMock()) + + updated_fv = fs._create_updated_feature_view(base_fv, OnlineConfig(enable=True)) + + self.assertTrue(updated_fv.is_realtime_feature_view, "updated FV lost realtime identity") + self.assertIsNotNone(updated_fv.realtime_config) + + def test_create_online_feature_table_rejects_hybrid_for_tiled(self) -> None: + """A HYBRID_TABLE OFT is not supported for a tiled FV and must be rejected clearly.""" + base_fv = self._make_tiled_streaming_base_fv() + base_fv._online_config = OnlineConfig(enable=True, store_type=OnlineStoreType.HYBRID_TABLE) + self.assertTrue(base_fv.is_tiled) + + fs = object.__new__(FeatureStore) + object.__setattr__(fs, "_session", MagicMock()) + + with self.assertRaisesRegex(Exception, "not supported for aggregation"): + fs._create_online_feature_table(base_fv, SqlIdentifier("TILED_FV"), "v1") + + def test_construct_tiled_fv_with_hybrid_online_rejected(self) -> None: + """The HYBRID-on-tiled rule is a feature-view invariant: rejected at construction time.""" + with self.assertRaisesRegex(ValueError, "not supported for aggregation"): + FeatureView( + name="TILED_FV", + entities=[Entity(name="user", join_keys=["USER_ID"])], + feature_df=self._make_mock_df(), + timestamp_col="EVENT_TIME", + refresh_freq="1h", + feature_granularity="1h", + features=[Feature.sum("AMOUNT", "2h")], + online_config=OnlineConfig(enable=True), # default HYBRID_TABLE + ) + + def test_plan_online_enable_forwards_store_type_for_tiled(self) -> None: + """_plan_online_enable must forward the requested store_type and rebuild a still-tiled FV. + + Regression: it previously dropped store_type (defaulting to HYBRID_TABLE), which combined + with the tiled guard left no way to enable POSTGRES online on a tiled FV. + """ + base_fv = self._make_tiled_streaming_base_fv() + self.assertTrue(base_fv.is_tiled) + + fs = object.__new__(FeatureStore) + object.__setattr__(fs, "_session", MagicMock()) + + strategy = fs._plan_online_enable(base_fv, OnlineConfig(enable=True, store_type=OnlineStoreType.POSTGRES)) + + assert strategy.final_config is not None + self.assertEqual(strategy.final_config.store_type, OnlineStoreType.POSTGRES) + op_type, temp_fv = strategy.operations[0] + self.assertEqual(op_type, "CREATE_ONLINE") + assert isinstance(temp_fv, FeatureView) + self.assertTrue(temp_fv.is_tiled, "rebuilt FV must remain tiled") + assert temp_fv.online_config is not None + self.assertEqual(temp_fv.online_config.store_type, OnlineStoreType.POSTGRES) + + if __name__ == "__main__": absltest.main() diff --git a/snowflake/ml/feature_store/feature_store_realtime_test.py b/snowflake/ml/feature_store/feature_store_realtime_test.py index a88a4301..fa4a0a3f 100644 --- a/snowflake/ml/feature_store/feature_store_realtime_test.py +++ b/snowflake/ml/feature_store/feature_store_realtime_test.py @@ -422,11 +422,12 @@ def _create_df(values: list[list[Any]], schema: Any) -> MagicMock: fs.list_feature_views(verbose=True) verbose_rows: list[list[Any]] = captured["values"] self.assertEqual(len(verbose_rows), 1) - # verbose=True: base + extra fields are merged (20 total). - # Extra fields: source_refs (18), backup_source (19). - self.assertEqual(len(verbose_rows[0]), 20) - self.assertIsNone(verbose_rows[0][18]) # source_refs — always None for RTFV - self.assertIsNone(verbose_rows[0][19]) # backup_source — always None for RTFV + # verbose=True: base + extra fields are merged (21 total). + # Extra fields: initialization_warehouse (18), source_refs (19), backup_source (20). + self.assertEqual(len(verbose_rows[0]), 21) + self.assertIsNone(verbose_rows[0][18]) # initialization_warehouse — always None for RTFV + self.assertIsNone(verbose_rows[0][19]) # source_refs — always None for RTFV + self.assertIsNone(verbose_rows[0][20]) # backup_source — always None for RTFV if __name__ == "__main__": diff --git a/snowflake/ml/feature_store/feature_store_snapshot_validation_test.py b/snowflake/ml/feature_store/feature_store_snapshot_validation_test.py index 6d7abc9e..702c4260 100644 --- a/snowflake/ml/feature_store/feature_store_snapshot_validation_test.py +++ b/snowflake/ml/feature_store/feature_store_snapshot_validation_test.py @@ -9,7 +9,10 @@ from absl.testing import absltest, parameterized from snowflake.ml._internal.utils.sql_identifier import SqlIdentifier -from snowflake.ml.feature_store import feature_view_append_only_validation +from snowflake.ml.feature_store import ( + feature_store as fs_mod, + feature_view_append_only_validation, +) from snowflake.ml.feature_store.entity import Entity from snowflake.ml.feature_store.feature_view import ( FeatureView, @@ -480,7 +483,8 @@ def test_omitted_refresh_freq_skips_snapshot_validation_on_append_only(self) -> # new desc — confirming the validation_fv branch was not taken. plan_args, _ = fs._plan_feature_view_update_operations.call_args self.assertIsNone(plan_args[1], "actual_refresh_freq must be None when caller omits refresh_freq") - self.assertEqual(plan_args[3], "new description") + # desc is at index 4 (initialization_warehouse occupies index 3). + self.assertEqual(plan_args[4], "new description") # Caller's FV is untouched — the existing cron schedule is preserved. self.assertEqual(fv.refresh_freq, original_refresh_freq) @@ -778,7 +782,7 @@ def test_cron_to_cron_update_rollback_references_old_cron(self) -> None: new_cron = "30 6 * * * UTC" operations, rollback_ops = fs._build_offline_update_queries( - fv, refresh_freq=new_cron, warehouse=None, desc="new desc" + fv, refresh_freq=new_cron, warehouse=None, initialization_warehouse=fs_mod._KEEP_CURRENT, desc="new desc" ) # Forward operations apply the new cron to the companion Task. @@ -815,7 +819,7 @@ def test_downstream_update_is_case_insensitive(self, new_freq: str) -> None: fv._desc = "original desc" operations, rollback_ops = fs._build_offline_update_queries( - fv, refresh_freq=new_freq, warehouse=None, desc="new desc" + fv, refresh_freq=new_freq, warehouse=None, initialization_warehouse=fs_mod._KEEP_CURRENT, desc="new desc" ) forward_sqls = [sql for _, sql in operations] @@ -856,7 +860,7 @@ def test_old_downstream_to_cron_update_is_case_insensitive(self, old_freq: str) new_cron = "30 6 * * * UTC" operations, _rollback_ops = fs._build_offline_update_queries( - fv, refresh_freq=new_cron, warehouse=None, desc="new desc" + fv, refresh_freq=new_cron, warehouse=None, initialization_warehouse=fs_mod._KEEP_CURRENT, desc="new desc" ) forward_sqls = [sql for _, sql in operations] diff --git a/snowflake/ml/feature_store/feature_view.py b/snowflake/ml/feature_store/feature_view.py index 560ecb42..041d9dba 100644 --- a/snowflake/ml/feature_store/feature_view.py +++ b/snowflake/ml/feature_store/feature_view.py @@ -81,6 +81,13 @@ class OnlineStoreType(Enum): POSTGRES = "postgres" +# Online store types that can back an aggregation (tiled) feature view. POSTGRES is spec-backed +# and tile-aware; HYBRID_TABLE builds the online table directly over the tile output (which exposes +# tile columns rather than the raw timestamp column) and cannot. New store types must opt in here +# explicitly so an unsupported store is rejected by default for tiled feature views. +_TILE_CAPABLE_ONLINE_STORES = frozenset({OnlineStoreType.POSTGRES}) + + class _FeatureViewSchemaNotReadyWarning(UserWarning): """Issued when feature view schema is unavailable because the backing dynamic table has not completed its initial refresh yet.""" @@ -670,6 +677,7 @@ def __init__( refresh_freq: Optional[str] = None, desc: str = "", warehouse: Optional[str] = None, + initialization_warehouse: Optional[str] = None, initialize: str = "ON_CREATE", refresh_mode: Optional[str] = "AUTO", cluster_by: Optional[list[str]] = None, @@ -712,6 +720,11 @@ def __init__( warehouse: The warehouse used to refresh this feature view. Not needed when ``refresh_freq`` is ``None``. This warehouse will overwrite the default warehouse of Feature Store if specified, otherwise the default warehouse will be used. + initialization_warehouse: The warehouse used for the initial build and any subsequent reinitializations of + this feature view, which perform a full scan of the source data and are typically more resource + intensive than incremental refreshes. Maps to ``INITIALIZATION_WAREHOUSE`` on the backing dynamic + table. When unset, ``warehouse`` is used for all refresh operations. For streaming feature views, this + warehouse is also used for the one-time backfill. Not effective for static feature views. initialize: Specifies the behavior of the initial refresh of feature view. This property cannot be altered after you register the feature view. It supports ON_CREATE (default) or ON_SCHEDULE. ON_CREATE refreshes the feature view synchronously at creation. ON_SCHEDULE refreshes the feature view at the next scheduled @@ -1019,6 +1032,9 @@ def __init__( self._schema: Optional[SqlIdentifier] = None self._initialize: str = initialize self._warehouse: Optional[SqlIdentifier] = SqlIdentifier(warehouse) if warehouse is not None else None + self._initialization_warehouse: Optional[SqlIdentifier] = ( + SqlIdentifier(initialization_warehouse) if initialization_warehouse is not None else None + ) self._refresh_mode: Optional[str] = refresh_mode self._refresh_mode_reason: Optional[str] = None self._owner: Optional[str] = None @@ -1822,6 +1838,42 @@ def warehouse(self, new_value: str) -> None: ) self._warehouse = SqlIdentifier(new_value) + @property + def initialization_warehouse(self) -> Optional[SqlIdentifier]: + return self._initialization_warehouse + + @initialization_warehouse.setter + def initialization_warehouse(self, new_value: str) -> None: + """Set the initialization warehouse of feature view. + + Args: + new_value: The new value of initialization warehouse. + + Example:: + + >>> fs = FeatureStore(...) + >>> e = fs.get_entity('TRIP_ID') + >>> draft_fv = FeatureView( + ... name='F_TRIP', + ... entities=[e], + ... feature_df=feature_df, + ... refresh_freq='1d', + ... warehouse='SMALL_WH', + ... initialization_warehouse='LARGE_WH', + ... ) + >>> fv_1 = fs.register_feature_view(draft_fv, version='1.0') + >>> print(fv_1.initialization_warehouse) + LARGE_WH + + """ + warnings.warn( + "You must call register_feature_view() to make it effective. " + "Or use update_feature_view(initialization_warehouse=).", + stacklevel=2, + category=UserWarning, + ) + self._initialization_warehouse = SqlIdentifier(new_value) + @property def initialize(self) -> str: return self._initialize @@ -2072,6 +2124,33 @@ def _validate(self) -> None: if self._refresh_freq is None: raise ValueError("Iceberg storage requires refresh_freq.") + # Validate the online store type can back this feature view (e.g. a tiled FV cannot use a + # HYBRID_TABLE online store). + self._validate_online_store_supported() + + def _validate_online_store_supported(self) -> None: + """Reject online store types that cannot back this feature view. + + An aggregation (tiled) feature view's online table is created over its tile output, which + only tile-aware (spec-backed) stores support. This is an invariant of the feature view, so + it is enforced here at validation time (construction / reconstruction) and re-checked before + creating the online table. + + Raises: + ValueError: If the configured online store cannot back a tiled feature view. + """ + if ( + self._online_config is not None + and self._online_config.enable + and self.is_tiled + and self._online_config.store_type not in _TILE_CAPABLE_ONLINE_STORES + ): + raise ValueError( + f"{self._online_config.store_type.name} online store is not supported for " + f"aggregation (tiled) feature views. Use " + f"OnlineConfig(store_type=OnlineStoreType.POSTGRES) to enable online storage." + ) + def _validate_window_offset_alignment(self) -> None: """Validate that window and offset are multiples of feature_granularity. @@ -2236,6 +2315,7 @@ def __eq__(self, other: object) -> bool: and str(self.status) == str(other.status) and self.database == other.database and self.warehouse == other.warehouse + and self.initialization_warehouse == other.initialization_warehouse and self.refresh_mode == other.refresh_mode and self.refresh_mode_reason == other.refresh_mode_reason and self._owner == other._owner @@ -2256,6 +2336,9 @@ def _to_dict(self) -> dict[str, str]: fv_dict["_database"] = str(self._database) if self._database is not None else None fv_dict["_schema"] = str(self._schema) if self._schema is not None else None fv_dict["_warehouse"] = str(self._warehouse) if self._warehouse is not None else None + fv_dict["_initialization_warehouse"] = ( + str(self._initialization_warehouse) if self._initialization_warehouse is not None else None + ) fv_dict["_timestamp_col"] = str(self._timestamp_col) if self._timestamp_col is not None else None fv_dict["_initialize"] = str(self._initialize) @@ -2374,6 +2457,7 @@ def from_json(cls, json_str: str, session: Session) -> FeatureView: database=json_dict["_database"], schema=json_dict["_schema"], warehouse=json_dict["_warehouse"], + initialization_warehouse=json_dict.get("_initialization_warehouse"), refresh_mode=json_dict["_refresh_mode"], refresh_mode_reason=json_dict["_refresh_mode_reason"], initialize=json_dict["_initialize"], @@ -2482,6 +2566,7 @@ def _construct_feature_view( owner: Optional[str], infer_schema_df: Optional[DataFrame], session: Session, + initialization_warehouse: Optional[str] = None, cluster_by: Optional[list[str]] = None, online_config: Optional[OnlineConfig] = None, feature_granularity: Optional[str] = None, @@ -2540,6 +2625,9 @@ def _construct_feature_view( fv._database = SqlIdentifier(database) if database is not None else None fv._schema = SqlIdentifier(schema) if schema is not None else None fv._warehouse = SqlIdentifier(warehouse) if warehouse is not None else None + fv._initialization_warehouse = ( + SqlIdentifier(initialization_warehouse) if initialization_warehouse is not None else None + ) fv._refresh_mode_reason = refresh_mode_reason fv._initialize = initialize fv._owner = owner diff --git a/snowflake/ml/feature_store/feature_view_test.py b/snowflake/ml/feature_store/feature_view_test.py index 4d1fe8b2..38823958 100644 --- a/snowflake/ml/feature_store/feature_view_test.py +++ b/snowflake/ml/feature_store/feature_view_test.py @@ -153,6 +153,52 @@ def test_duplicate_alias_case_insensitive(self) -> None: self.assertIn("Duplicate feature alias", str(cm.exception)) +class InitializationWarehouseTest(absltest.TestCase): + """Unit tests for the ``initialization_warehouse`` knob on FeatureView.""" + + def _make_fv(self, **kwargs: Any) -> FeatureView: + mock_df = MagicMock() + mock_df.columns = ["user_id", "amount"] + mock_df.queries = {"queries": ["SELECT * FROM source"]} + return FeatureView( + name="test_fv", + entities=[Entity(name="user", join_keys=["user_id"])], + feature_df=mock_df, + refresh_freq="1d", + **kwargs, + ) + + def test_defaults_to_none(self) -> None: + fv = self._make_fv() + self.assertIsNone(fv.initialization_warehouse) + + def test_stores_value_as_sql_identifier(self) -> None: + fv = self._make_fv(warehouse="small_wh", initialization_warehouse="large_wh") + self.assertEqual(fv.initialization_warehouse, SqlIdentifier("large_wh")) + self.assertEqual(fv.warehouse, SqlIdentifier("small_wh")) + + def test_setter_warns_and_sets(self) -> None: + fv = self._make_fv() + with self.assertWarnsRegex(UserWarning, "register_feature_view"): + fv.initialization_warehouse = "large_wh" + self.assertEqual(fv.initialization_warehouse, SqlIdentifier("large_wh")) + + def test_to_dict_includes_value(self) -> None: + fv = self._make_fv(initialization_warehouse="large_wh") + self.assertEqual(fv._to_dict()["_initialization_warehouse"], "LARGE_WH") + + def test_to_dict_none_when_unset(self) -> None: + fv = self._make_fv() + self.assertIsNone(fv._to_dict()["_initialization_warehouse"]) + + def test_equality_distinguishes_initialization_warehouse(self) -> None: + fv_a = self._make_fv(initialization_warehouse="large_wh") + fv_b = self._make_fv(initialization_warehouse="other_wh") + fv_c = self._make_fv(initialization_warehouse="large_wh") + self.assertNotEqual(fv_a, fv_b) + self.assertEqual(fv_a, fv_c) + + class SecondaryKeyFeatureViewTest(parameterized.TestCase): """Unit tests for FV-level ``aggregation_secondary_keys``.""" @@ -1481,26 +1527,23 @@ def test_postgres_pk_includes_aggregation_secondary_keys(self) -> None: parsed = json.loads(query[start:end]) self.assertEqual(parsed["spec"]["ordered_entity_column_names"], ["USER_ID", "AD_ID"]) - def test_hybrid_pk_excludes_secondary_keys(self) -> None: - """HYBRID_TABLE path: no spec JSON is sent and the DDL PK is unchanged when SKs are set.""" - fv = self._make_feature_view( - entity_keys=["USER_ID"], - columns=["USER_ID", "EVENT_TS", "AMOUNT", "AD_ID"], - column_types=[StringType(), TimestampType(TimestampTimeZone.NTZ), DoubleType(), StringType()], - timestamp_col="EVENT_TS", - store_type=OnlineStoreType.HYBRID_TABLE, - feature_granularity="1h", - aggregation_secondary_keys=["AD_ID"], - features=[Feature.sum("AMOUNT", "24h").alias("AMOUNT_SUM_24H")], - ) - fs = self._make_mock_feature_store() + def test_hybrid_online_rejected_for_tiled_fv(self) -> None: + """HYBRID_TABLE online is unsupported for aggregation (tiled) feature views. - fs._create_online_feature_table(fv, SqlIdentifier("DT_NAME"), version="v1") - - query = fs._session.sql.call_args_list[0][0][0] - self.assertIn('PRIMARY KEY ("USER_ID")', query) - self.assertNotIn("AD_ID", query) - self.assertNotIn("FROM SPECIFICATION", query) + Tiles and aggregation secondary keys require the spec-backed POSTGRES store; this is a + feature-view invariant, so construction is rejected. + """ + with self.assertRaisesRegex(ValueError, "not supported for aggregation"): + self._make_feature_view( + entity_keys=["USER_ID"], + columns=["USER_ID", "EVENT_TS", "AMOUNT", "AD_ID"], + column_types=[StringType(), TimestampType(TimestampTimeZone.NTZ), DoubleType(), StringType()], + timestamp_col="EVENT_TS", + store_type=OnlineStoreType.HYBRID_TABLE, + feature_granularity="1h", + aggregation_secondary_keys=["AD_ID"], + features=[Feature.sum("AMOUNT", "24h").alias("AMOUNT_SUM_24H")], + ) class PostgresOnlineLocalRowCoercionTest(absltest.TestCase): diff --git a/snowflake/ml/feature_store/realtime_registration.py b/snowflake/ml/feature_store/realtime_registration.py index 6cb77061..4ae3c85b 100644 --- a/snowflake/ml/feature_store/realtime_registration.py +++ b/snowflake/ml/feature_store/realtime_registration.py @@ -828,9 +828,9 @@ def append_realtime_listing_row( ] output_values.append(values) - # RTFVs have no BatchSource/StreamSource ref-list, so source_refs is always - # None. backup_source is also always None for RTFVs. - output_values_extra.append([None, None]) # source_refs, backup_source + # RTFVs have no initialization warehouse, BatchSource/StreamSource ref-list, + # or backup_source, so all verbose-only fields are None. + output_values_extra.append([None, None, None]) # initialization_warehouse, source_refs, backup_source def append_realtime_listing_rows( diff --git a/snowflake/ml/feature_store/spec/builder.py b/snowflake/ml/feature_store/spec/builder.py index cdd8f83f..84bb0ae4 100644 --- a/snowflake/ml/feature_store/spec/builder.py +++ b/snowflake/ml/feature_store/spec/builder.py @@ -259,6 +259,27 @@ def _resolve_tiled_fv_output_column( ) +def _as_secondary_key_array_column(col: FSColumn) -> FSColumn: + """Wrap a scalar aggregation output column as the ``ArrayType`` a secondary-key FV materializes. + + Args: + col: The scalar FSColumn resolved for the aggregation's logical output. + + Returns: + An ``ArrayType`` FSColumn carrying ``col``'s scalar type as + ``element_type`` and preserving precision/scale/length/timezone. + """ + return FSColumn( + name=col.name, + type="ArrayType", + element_type=col.type, + precision=col.precision, + scale=col.scale, + length=col.length, + timezone=col.timezone, + ) + + # Type alias for the polymorphic source input SourceInput = Union[StreamSource, RequestSource, "FeatureView", "FeatureViewSlice", BatchSource] @@ -779,7 +800,10 @@ def _columns_from_feature_view(fv: FeatureView, materialized_schema: Optional[St For tiled FVs the columns are derived from ``aggregation_specs`` because ``output_schema`` exposes ``_PARTIAL_*`` tile columns, not the FV's logical - aggregation outputs. + aggregation outputs. A secondary-key FV (``aggregation_secondary_keys``) + additionally wraps each value aggregation as ``ArrayType`` so the + column matches the upstream FV's stored ``OutputColumn`` for the Online + Service exact-shape check. For non-tiled, non-RTFV sources, when ``materialized_schema`` is provided, column shapes are read from it instead of the authoring-time @@ -803,11 +827,16 @@ def _columns_from_feature_view(fv: FeatureView, materialized_schema: Optional[St ] if fv.is_tiled and fv.aggregation_specs is not None: partials_by_name = {f.name: _make_tiled_fs_column(f.name, f.datatype) for f in fv.output_schema.fields} - return [ - _resolve_tiled_fv_output_column(spec, partials_by_name) - for spec in fv.aggregation_specs - if not spec.function.is_secondary_key_array() - ] + has_secondary_key = bool(fv.aggregation_secondary_keys) + columns: list[FSColumn] = [] + for spec in fv.aggregation_specs: + if spec.function.is_secondary_key_array(): + continue + resolved = _resolve_tiled_fv_output_column(spec, partials_by_name) + if has_secondary_key: + resolved = _as_secondary_key_array_column(resolved) + columns.append(resolved) + return columns feature_name_set = {fn.resolved() for fn in fv.feature_names} schema = materialized_schema if materialized_schema is not None else fv.output_schema return [_make_fs_column(f.name, f.datatype) for f in schema.fields if f.name in feature_name_set] diff --git a/snowflake/ml/feature_store/spec/builder_test.py b/snowflake/ml/feature_store/spec/builder_test.py index 9eb03ebe..78331b7c 100644 --- a/snowflake/ml/feature_store/spec/builder_test.py +++ b/snowflake/ml/feature_store/spec/builder_test.py @@ -1058,6 +1058,7 @@ def _make_tiled_fv( output_schema: StructType, name: str = "tiled_fv", version: str = "v1", + aggregation_secondary_keys: Optional[list[str]] = None, ) -> mock.MagicMock: """Build a mock tiled FeatureView with the given agg specs and tile schema.""" from snowflake.ml.feature_store.feature_view import FeatureView @@ -1070,6 +1071,7 @@ def _make_tiled_fv( fv.aggregation_specs = aggregation_specs fv.output_schema = output_schema fv.feature_names = [SqlIdentifier(s.output_column) for s in aggregation_specs] + fv.aggregation_secondary_keys = aggregation_secondary_keys return fv @staticmethod @@ -1214,12 +1216,59 @@ def test_secondary_key_array_filtered_out(self) -> None: StructField("_PARTIAL_COUNT_AMOUNT", LongType()), ] ) - fv = self._make_tiled_fv(aggregation_specs=specs, output_schema=schema) + fv = self._make_tiled_fv(aggregation_specs=specs, output_schema=schema, aggregation_secondary_keys=["MERCHANT"]) cols = FeatureViewSpecBuilder._columns_from_feature_view(fv) self.assertEqual([c.name for c in cols], ["AMOUNT_SUM_1D"]) + def test_secondary_key_fv_value_outputs_wrapped_as_arrays(self) -> None: + """A FeatureGroup source over a secondary-key FV must mirror the FV's stored ArrayType shape.""" + specs = [ + AggregationSpec(AggregationType.SUM, "AMOUNT", "1d", "AMOUNT_SUM_1D"), + AggregationSpec(AggregationType.COUNT, "AMOUNT", "1d", "AMOUNT_COUNT_1D"), + AggregationSpec( + function=AggregationType._SECONDARY_KEY_ARRAY, + source_column="MERCHANT", + window="1d", + output_column="_INTERNAL_MERCHANT_ARR", + ), + ] + schema = StructType( + [ + StructField("USER_ID", StringType()), + StructField("_PARTIAL_SUM_AMOUNT", DecimalType(10, 2)), + StructField("_PARTIAL_COUNT_AMOUNT", LongType()), + ] + ) + fv = self._make_tiled_fv(aggregation_specs=specs, output_schema=schema, aggregation_secondary_keys=["MERCHANT"]) + + cols = FeatureViewSpecBuilder._columns_from_feature_view(fv) + + self.assertEqual([c.name for c in cols], ["AMOUNT_SUM_1D", "AMOUNT_COUNT_1D"]) + sum_col, count_col = cols + # SUM inherits the source scalar's precision/scale as the array element. + self.assertEqual(sum_col.type, "ArrayType") + self.assertEqual(sum_col.element_type, "DecimalType") + self.assertEqual(sum_col.precision, 10) + self.assertEqual(sum_col.scale, 2) + # COUNT's predetermined DecimalType(18, 0) becomes the array element. + self.assertEqual(count_col.type, "ArrayType") + self.assertEqual(count_col.element_type, "DecimalType") + self.assertEqual(count_col.precision, 18) + self.assertEqual(count_col.scale, 0) + + def test_non_secondary_key_fv_outputs_stay_scalar(self) -> None: + """Without a secondary key, value aggregations stay scalar (no array wrapping).""" + spec = AggregationSpec(AggregationType.SUM, "AMOUNT", "1d", "AMOUNT_SUM_1D") + fv = self._make_tiled_fv(aggregation_specs=[spec], output_schema=self._tile_schema_for_amount()) + + cols = FeatureViewSpecBuilder._columns_from_feature_view(fv) + + self.assertEqual(len(cols), 1) + self.assertEqual(cols[0].type, "DecimalType") + self.assertIsNone(cols[0].element_type) + def test_slice_over_tiled_feature_view(self) -> None: """``_convert_feature_view_slice`` selects logical agg outputs in slice order.""" from snowflake.ml.feature_store.feature_view import FeatureViewSlice diff --git a/snowflake/ml/feature_store/streaming_registration.py b/snowflake/ml/feature_store/streaming_registration.py index 3c5fa21a..d2fc3ef4 100644 --- a/snowflake/ml/feature_store/streaming_registration.py +++ b/snowflake/ml/feature_store/streaming_registration.py @@ -471,8 +471,9 @@ def run_streaming_postamble( feature_view_name: Physical FV name (``$``). preamble: Result from ``run_streaming_preamble``. metadata_manager: Metadata manager. - default_warehouse: Feature Store default warehouse; used only when - ``feature_view.warehouse`` is unset (one of the two must be set). + default_warehouse: Feature Store default warehouse; used only when neither + ``feature_view.initialization_warehouse`` nor ``feature_view.warehouse`` + is set (one of the three must be set). get_fully_qualified_name_fn: Bound ``FeatureStore._get_fully_qualified_name``. telemetry_stmp: Telemetry statement parameters. on_resource_created: Optional callback invoked immediately after each @@ -504,7 +505,16 @@ def run_streaming_postamble( if stream_config is None: raise ValueError(f"FeatureView '{feature_view.name}' does not have a stream_config.") - task_warehouse = feature_view.warehouse if feature_view.warehouse is not None else default_warehouse + # The backfill is the streaming FV's one-time, full-scan initialization, so it + # runs on the initialization warehouse when set, mirroring the dynamic table's + # initial/reinit refresh. Falls back to the FV warehouse, then the FS default. + task_warehouse = ( + feature_view.initialization_warehouse + if feature_view.initialization_warehouse is not None + else feature_view.warehouse + if feature_view.warehouse is not None + else default_warehouse + ) if task_warehouse is None: raise ValueError( "No warehouse available for streaming backfill task graph. Either set " diff --git a/snowflake/ml/feature_store/streaming_registration_test.py b/snowflake/ml/feature_store/streaming_registration_test.py index 485d6372..3057a39d 100644 --- a/snowflake/ml/feature_store/streaming_registration_test.py +++ b/snowflake/ml/feature_store/streaming_registration_test.py @@ -1548,6 +1548,56 @@ def test_postamble_uses_default_warehouse_when_fv_has_none(self) -> None: if sql.startswith("CREATE OR REPLACE TASK"): self.assertIn("WAREHOUSE = DEFAULT_WH", sql) + def test_postamble_prefers_initialization_warehouse_for_backfill(self) -> None: + """The backfill task graph runs on the initialization warehouse when set, + since the backfill is the streaming FV's one-time, full-scan initialization.""" + backfill_df = _make_mock_backfill_df() + entity = _make_entity() + stream_config = StreamConfig( + stream_source="txn_events", + transformation_fn=_sample_transform, + backfill_df=backfill_df, + ) + fv = FeatureView( + name="test_fv", + entities=[entity], + stream_config=stream_config, + timestamp_col="EVENT_TIME", + warehouse="small_wh", + initialization_warehouse="large_wh", + ) + + session = self._make_session_with_udf_schema() + metadata_manager = self._make_metadata_manager() + + from snowflake.ml.feature_store.streaming_registration import ( + StreamingPreambleResult, + ) + + preamble = StreamingPreambleResult( + fq_udf_table="DB.SCH.UDF_TABLE", + fq_backfill_table="DB.SCH.UDF_TABLE$BACKFILL", + resolved_source_name="TXN_EVENTS", + ) + + run_streaming_postamble( + session=session, + feature_view=fv, + version=FeatureViewVersion("v1"), + feature_view_name=FeatureView._get_physical_name(SqlIdentifier("test_fv"), FeatureViewVersion("v1")), + preamble=preamble, + metadata_manager=metadata_manager, + default_warehouse=SqlIdentifier("default_wh"), + get_fully_qualified_name_fn=self._fq, + telemetry_stmp={}, + ) + + sql_calls = [str(c.args[0]) for c in session.sql.call_args_list] + for sql in sql_calls: + if sql.startswith("CREATE OR REPLACE TASK"): + self.assertIn("WAREHOUSE = LARGE_WH", sql) + self.assertNotIn("WAREHOUSE = SMALL_WH", sql) + def test_postamble_no_warehouse_raises(self) -> None: """Postamble raises if neither the FV nor the FS provides a warehouse.""" backfill_df = _make_mock_backfill_df() diff --git a/snowflake/ml/model/_client/model/BUILD.bazel b/snowflake/ml/model/_client/model/BUILD.bazel index 1a7e031e..303a31f1 100644 --- a/snowflake/ml/model/_client/model/BUILD.bazel +++ b/snowflake/ml/model/_client/model/BUILD.bazel @@ -5,7 +5,9 @@ package(default_visibility = ["//visibility:public"]) py_library( name = "batch_inference_specs", srcs = ["batch_inference_specs.py"], - deps = [], + deps = [ + "//snowflake/ml/model:inference_engine", + ], ) py_library( @@ -50,6 +52,7 @@ py_test( ], deps = [ ":batch_inference_specs", + "//snowflake/ml/model:inference_engine", ], ) diff --git a/snowflake/ml/model/_client/model/batch_inference_specs.py b/snowflake/ml/model/_client/model/batch_inference_specs.py index 63c1255d..57946594 100644 --- a/snowflake/ml/model/_client/model/batch_inference_specs.py +++ b/snowflake/ml/model/_client/model/batch_inference_specs.py @@ -4,6 +4,8 @@ from pydantic import BaseModel, model_validator from typing_extensions import TypedDict +from snowflake.ml.model import inference_engine as inference_engine_module + class SaveMode(str, Enum): """Save mode options for batch inference output. @@ -183,3 +185,60 @@ def _validate_job_name_exclusivity(self) -> "JobSpec": if self.job_name is not None and self.job_name_prefix is not None: raise ValueError("job_name and job_name_prefix are mutually exclusive. Please specify only one or neither.") return self + + +# ---------------------------------------------------------------------------- +# Spec types for EXECUTE INFERENCE JOB SERVICE. +# +# These types mirror the YAML body 1:1 (`input` / `output` / `resources` +# / `inference` / `image_build`). +# ---------------------------------------------------------------------------- + + +class Input(BaseModel): + """Input block of the batch inference YAML body.""" + + params: Optional[dict[str, Any]] = None + column_handling: Optional[dict[str, ColumnHandlingOptions]] = None + partition_column: Optional[str] = None + + +class Output(BaseModel): + """Output block of the batch inference YAML body. + + ``stage_location`` is required. ``base_stage_location`` is not part of the + YAML schema and is intentionally omitted. + """ + + stage_location: str + mode: SaveMode = SaveMode.ERROR + + +class Resources(BaseModel): + """Resources block of the batch inference YAML body.""" + + cpu_requests: Optional[str] = None + memory_requests: Optional[str] = None + gpu_requests: Optional[str] = None + + +class EngineOptions(BaseModel): + """``inference.engine_options`` sub-block of the YAML body.""" + + engine: Optional[inference_engine_module.InferenceEngine] = None + engine_args_override: Optional[list[str]] = None + + +class Inference(BaseModel): + """Inference block of the batch inference YAML body.""" + + num_workers: Optional[int] = None + max_batch_rows: Optional[int] = None + engine_options: Optional[EngineOptions] = None + + +class ImageBuild(BaseModel): + """Image-build block of the batch inference YAML body.""" + + image_repo: Optional[str] = None + force_rebuild: bool = False diff --git a/snowflake/ml/model/_client/model/batch_inference_specs_test.py b/snowflake/ml/model/_client/model/batch_inference_specs_test.py index 04e72404..abec6018 100644 --- a/snowflake/ml/model/_client/model/batch_inference_specs_test.py +++ b/snowflake/ml/model/_client/model/batch_inference_specs_test.py @@ -1,5 +1,6 @@ from absl.testing import absltest +from snowflake.ml.model import inference_engine as inference_engine_module from snowflake.ml.model._client.model import batch_inference_specs @@ -21,6 +22,38 @@ def test_output_spec_stage_exclusivity(self) -> None: with self.assertRaises(ValueError): batch_inference_specs.OutputSpec() + def test_inference_job_output_requires_stage_location(self) -> None: + with self.assertRaises(ValueError): + batch_inference_specs.Output() # type: ignore[call-arg] + + def test_inference_job_output_default_mode_is_error(self) -> None: + spec = batch_inference_specs.Output(stage_location="@db.schema.stage/out/") + self.assertEqual(spec.mode, batch_inference_specs.SaveMode.ERROR) + + def test_inference_job_input_defaults_are_none(self) -> None: + spec = batch_inference_specs.Input() + self.assertIsNone(spec.params) + self.assertIsNone(spec.column_handling) + self.assertIsNone(spec.partition_column) + + def test_inference_job_engine_options_accepts_enum(self) -> None: + opts = batch_inference_specs.EngineOptions( + engine=inference_engine_module.InferenceEngine.VLLM, + engine_args_override=["--max-num-seqs=128"], + ) + self.assertEqual(opts.engine, inference_engine_module.InferenceEngine.VLLM) + self.assertEqual(opts.engine_args_override, ["--max-num-seqs=128"]) + + def test_inference_job_resources_all_optional(self) -> None: + spec = batch_inference_specs.Resources() + self.assertIsNone(spec.cpu_requests) + self.assertIsNone(spec.memory_requests) + self.assertIsNone(spec.gpu_requests) + + def test_inference_job_image_build_default_force_rebuild_false(self) -> None: + spec = batch_inference_specs.ImageBuild() + self.assertFalse(spec.force_rebuild) + if __name__ == "__main__": absltest.main() diff --git a/snowflake/ml/model/_client/model/model_version_impl.py b/snowflake/ml/model/_client/model/model_version_impl.py index 366a3aee..b33c6844 100644 --- a/snowflake/ml/model/_client/model/model_version_impl.py +++ b/snowflake/ml/model/_client/model/model_version_impl.py @@ -892,6 +892,111 @@ def run_batch( inference_engine_args=inference_engine_args, ) + def _run_batch_v2( + self, + X: dataframe.DataFrame, + *, + compute_pool: str, + output_spec: batch_inference_specs.Output, + input_spec: Optional[batch_inference_specs.Input] = None, + resources_spec: Optional[batch_inference_specs.Resources] = None, + inference_spec: Optional[batch_inference_specs.Inference] = None, + image_build_spec: Optional[batch_inference_specs.ImageBuild] = None, + function_name: Optional[str] = None, + job_name: Optional[str] = None, + replicas: Optional[int] = None, + async_: bool = True, + ) -> job.MLJob[Any]: + """Run batch inference via ``EXECUTE INFERENCE JOB SERVICE`` (private). + + This is a private preview API. Prefer :meth:`run_batch` until this path + is generally available. + + Args: + X: Snowpark DataFrame with the input rows. + compute_pool: Compute pool used by the SPCS job and image build. + output_spec: Output block. + input_spec: Input block. + resources_spec: Resources block. + inference_spec: Inference block. + image_build_spec: Image build block. + function_name: Model function name. Resolved against the model's + function list when omitted. + job_name: Optional fully qualified job name. When omitted the + server generates a name. + replicas: Optional ``REPLICAS`` value. + async_: ``ASYNC`` clause value. Defaults to ``True``. + + Returns: + MLJob handle for the launched batch inference job. + + Raises: + ValueError: If ``input_spec.partition_column`` is supplied for a + HuggingFace pipeline model or a FUNCTION-type method, or if + the partition column collides with a partitioned model output. + """ + statement_params = telemetry.get_statement_params( + project=_TELEMETRY_PROJECT, + subproject=_TELEMETRY_SUBPROJECT, + ) + + effective_input_spec = input_spec if input_spec is not None else batch_inference_specs.Input() + partition_columns = ( + [effective_input_spec.partition_column] if effective_input_spec.partition_column is not None else None + ) + + if partition_columns is not None: + model_spec = self._get_model_spec(statement_params) + if model_spec.get("model_type") == huggingface.TransformersPipelineHandler.HANDLER_TYPE: + raise ValueError( + "partition_column is not supported for HuggingFace pipeline models in batch inference jobs. " + "Please remove the partition_column from input_spec." + ) + + gpu_requests = resources_spec.gpu_requests if resources_spec is not None else None + self._throw_error_if_gpu_is_not_supported(gpu_requests, statement_params) + + target_function_info = self._get_function_info(function_name=function_name) + + if ( + partition_columns is not None + and target_function_info["target_method_function_type"] + == model_manifest_schema.ModelMethodFunctionTypes.FUNCTION.value + ): + raise ValueError( + "partition_column is not supported for FUNCTION type methods in batch inference jobs. " + "Only TABLE_FUNCTION type methods support partitioning." + ) + + if partition_columns is not None and target_function_info["is_partitioned"]: + output_cols_upper = {spec.name.upper() for spec in target_function_info["signature"].outputs} + partition_cols_upper = {p.upper() for p in partition_columns} + collisions = sorted(partition_cols_upper & output_cols_upper) + if collisions: + raise ValueError( + f"Partitioned model output includes the partition column(s) {collisions}. " + f"Batch inference automatically appends the partition column to the output of " + f"partitioned models; please remove {collisions} from the model output " + f"and re-register the model." + ) + + return self._service_ops.execute_inference_job_service( + X=X, + model_name=self._model_name, + version_name=self._version_name, + compute_pool_name=sql_identifier.SqlIdentifier(compute_pool), + input_spec=input_spec, + output_spec=output_spec, + resources_spec=resources_spec, + inference_spec=inference_spec, + image_build_spec=image_build_spec, + function_name=target_function_info["target_method"], + job_name=job_name, + replicas=replicas, + async_=async_, + statement_params=statement_params, + ) + def _get_function_info(self, function_name: Optional[str]) -> model_manifest_schema.ModelFunctionInfo: functions: list[model_manifest_schema.ModelFunctionInfo] = self._functions diff --git a/snowflake/ml/model/_client/model/model_version_impl_test.py b/snowflake/ml/model/_client/model/model_version_impl_test.py index 281f0ce7..5f9675da 100644 --- a/snowflake/ml/model/_client/model/model_version_impl_test.py +++ b/snowflake/ml/model/_client/model/model_version_impl_test.py @@ -3192,6 +3192,98 @@ def test_create_service_forwards_feature_sources_per_function(self) -> None: _, kwargs = mock_create_service.call_args self.assertIs(kwargs["feature_sources_per_function"], sentinel_feature_sources) + def test_run_batch_v2_forwards_specs_and_resolved_function(self) -> None: + input_df = mock.MagicMock(spec=dataframe.DataFrame) + + output_spec = batch_inference_specs.Output(stage_location="@output_stage") + input_spec = batch_inference_specs.Input(params={"k": "v"}) + resources_spec = batch_inference_specs.Resources(cpu_requests="1") + inference_spec = batch_inference_specs.Inference(num_workers=2) + image_build_spec = batch_inference_specs.ImageBuild(image_repo="DB.SCHEMA.REPO") + mock_job = mock.MagicMock(spec=job.MLJob) + + with ( + mock.patch.object( + self.m_mv, + "_get_function_info", + return_value={ + "target_method": "predict", + "target_method_function_type": "FUNCTION", + "signature": _DUMMY_SIG["predict"], + "is_partitioned": False, + }, + ), + mock.patch.object( + self.m_mv._service_ops, "execute_inference_job_service", return_value=mock_job + ) as mock_execute, + ): + result = self.m_mv._run_batch_v2( + input_df, + compute_pool="POOL", + output_spec=output_spec, + input_spec=input_spec, + resources_spec=resources_spec, + inference_spec=inference_spec, + image_build_spec=image_build_spec, + function_name="predict", + job_name="JOB", + replicas=2, + async_=False, + ) + + mock_execute.assert_called_once_with( + X=input_df, + model_name=sql_identifier.SqlIdentifier("MODEL"), + version_name=sql_identifier.SqlIdentifier("v1", case_sensitive=True), + compute_pool_name=sql_identifier.SqlIdentifier("POOL"), + input_spec=input_spec, + output_spec=output_spec, + resources_spec=resources_spec, + inference_spec=inference_spec, + image_build_spec=image_build_spec, + function_name="predict", + job_name="JOB", + replicas=2, + async_=False, + statement_params=mock.ANY, + ) + self.assertEqual(result, mock_job) + + def test_run_batch_v2_minimal(self) -> None: + input_df = mock.MagicMock(spec=dataframe.DataFrame) + + output_spec = batch_inference_specs.Output(stage_location="@output_stage/") + mock_job = mock.MagicMock(spec=job.MLJob) + + with ( + mock.patch.object( + self.m_mv, + "_get_function_info", + return_value={ + "target_method": "predict", + "target_method_function_type": "FUNCTION", + "signature": _DUMMY_SIG["predict"], + "is_partitioned": False, + }, + ), + mock.patch.object( + self.m_mv._service_ops, "execute_inference_job_service", return_value=mock_job + ) as mock_execute, + ): + self.m_mv._run_batch_v2(input_df, compute_pool="POOL", output_spec=output_spec) + + kwargs = mock_execute.call_args.kwargs + self.assertIs(kwargs["X"], input_df) + self.assertIs(kwargs["output_spec"], output_spec) + self.assertTrue(kwargs["async_"]) + self.assertIsNone(kwargs["input_spec"]) + self.assertIsNone(kwargs["resources_spec"]) + self.assertIsNone(kwargs["inference_spec"]) + self.assertIsNone(kwargs["image_build_spec"]) + self.assertIsNone(kwargs["job_name"]) + self.assertIsNone(kwargs["replicas"]) + self.assertEqual(kwargs["function_name"], "predict") + if __name__ == "__main__": absltest.main() diff --git a/snowflake/ml/model/_client/ops/BUILD.bazel b/snowflake/ml/model/_client/ops/BUILD.bazel index ff617669..34d509c8 100644 --- a/snowflake/ml/model/_client/ops/BUILD.bazel +++ b/snowflake/ml/model/_client/ops/BUILD.bazel @@ -145,6 +145,7 @@ py_library( "//snowflake/ml/model/_client/model:batch_inference_serialization", "//snowflake/ml/model/_client/model:batch_inference_specs", "//snowflake/ml/model/_client/ops:deployment_step", + "//snowflake/ml/model/_client/service:inference_job_service_spec", "//snowflake/ml/model/_client/service:model_deployment_spec", "//snowflake/ml/model/_client/sql:service", "//snowflake/ml/model/_client/sql:stage", diff --git a/snowflake/ml/model/_client/ops/service_ops.py b/snowflake/ml/model/_client/ops/service_ops.py index 2f1c2e8e..bb68671d 100644 --- a/snowflake/ml/model/_client/ops/service_ops.py +++ b/snowflake/ml/model/_client/ops/service_ops.py @@ -5,6 +5,7 @@ import tempfile import threading import time +import uuid import warnings from collections.abc import Sequence from typing import Any, Optional, Union, cast @@ -20,10 +21,13 @@ batch_inference_specs, ) from snowflake.ml.model._client.ops import deployment_step, param_utils -from snowflake.ml.model._client.service import model_deployment_spec +from snowflake.ml.model._client.service import ( + inference_job_service_spec, + model_deployment_spec, +) from snowflake.ml.model._client.sql import service as service_sql, stage as stage_sql from snowflake.ml.model._signatures import core -from snowflake.snowpark import async_job, exceptions, row, session +from snowflake.snowpark import async_job, dataframe, exceptions, row, session from snowflake.snowpark._internal import utils as snowpark_utils module_logger = service_logger.get_logger(__name__, service_logger.LogColor.GREY) @@ -34,6 +38,13 @@ _UTF8_ENCODING = "utf-8" +# Reserved subdirectory under the output stage where input parquet is materialized. +# Must stay in sync with ``InferenceJobServiceSpecWrapper.RESERVED_INPUT_SUBDIR`` on +# the server side: the server treats files under this subdir as managed and +# excludes them from the output-mode preflight scan, and the SPCS container uses the +# same convention to decide whether to clean up the input stage after the job runs. +_BATCH_INFERENCE_RESERVED_INPUT_SUBDIR = "_snowflake_temporary" + @dataclasses.dataclass class ServiceLogInfo: @@ -1165,3 +1176,146 @@ def _create_temp_stage( statement_params=statement_params, ) return self._stage_client.fully_qualified_object_name(database_name, schema_name, stage_name) # stage path + + def execute_inference_job_service( + self, + *, + X: dataframe.DataFrame, + model_name: sql_identifier.SqlIdentifier, + version_name: sql_identifier.SqlIdentifier, + compute_pool_name: sql_identifier.SqlIdentifier, + input_spec: Optional[batch_inference_specs.Input], + output_spec: batch_inference_specs.Output, + resources_spec: Optional[batch_inference_specs.Resources], + inference_spec: Optional[batch_inference_specs.Inference], + image_build_spec: Optional[batch_inference_specs.ImageBuild], + function_name: Optional[str], + job_name: Optional[str], + replicas: Optional[int], + async_: bool, + statement_params: Optional[dict[str, Any]] = None, + ) -> job.MLJob[Any]: + """Materialize ``X``, build the YAML body, and run ``EXECUTE INFERENCE JOB SERVICE``. + + Args: + X: Input DataFrame. Materialized as parquet under + ``/_snowflake_temporary//`` + before the SQL command is issued. + model_name: Model identifier; combined with this operator's database + and schema to form the ``MODEL`` clause FQN. + version_name: Model version identifier for the ``VERSION`` clause. + compute_pool_name: Compute pool for ``IN COMPUTE POOL``. + input_spec: Optional input block of the YAML body. + output_spec: Required output block of the YAML body. ``stage_location`` + is normalized to end with ``/`` before being emitted. + resources_spec: Optional resources block of the YAML body. + inference_spec: Optional inference block of the YAML body. + image_build_spec: Optional image build block of the YAML body. + function_name: Optional model function name for ``FUNCTION``. + job_name: Optional fully qualified job name for ``NAME``. When + ``None`` the server generates a name and this method parses it + from the response. + replicas: Optional integer for ``REPLICAS``. + async_: ``ASYNC`` clause value. + statement_params: Optional statement params for telemetry. + + Returns: + MLJob for the launched batch inference job. + + Raises: + Exception: Any exception raised by the underlying SQL call. The + staged input is best-effort removed before re-raising so + rejected launches do not orphan files in the reserved subdir. + RuntimeError: If staging ``X`` fails, or if the server response does + not contain a parseable job name and ``job_name`` was not + provided. + """ + # All pure argument transformations happen before any I/O so a malformed + # input (e.g. an invalid job_name) fails with no side effects to clean up. + output_stage_location = output_spec.stage_location + if not output_stage_location.endswith("/"): + output_stage_location += "/" + # Stage the materialized input under the reserved subdirectory so the server-side + # preflight skips it and the SPCS container cleans it up after the job. The UUID + # avoids collisions between concurrent calls into the same output stage. + input_stage_location = f"{output_stage_location}{_BATCH_INFERENCE_RESERVED_INPUT_SUBDIR}/{uuid.uuid4().hex}/" + normalized_output_spec = output_spec.model_copy(update={"stage_location": output_stage_location}) + + spec_builder = inference_job_service_spec.InferenceJobServiceSpec() + if input_spec is not None: + spec_builder.add_input_spec(input_spec) + spec_builder.add_output_spec(normalized_output_spec) + if resources_spec is not None: + spec_builder.add_resources_spec(resources_spec) + if inference_spec is not None: + spec_builder.add_inference_spec(inference_spec) + if image_build_spec is not None: + spec_builder.add_image_build_spec(image_build_spec) + yaml_body = spec_builder.save() + + model_fqn = identifier.get_schema_level_object_identifier( + self._database_name.identifier(), self._schema_name.identifier(), model_name.identifier() + ) + + job_fqn: Optional[str] = None + job_database_name: Optional[sql_identifier.SqlIdentifier] = None + job_schema_name: Optional[sql_identifier.SqlIdentifier] = None + parsed_job_name: Optional[sql_identifier.SqlIdentifier] = None + if job_name is not None: + job_database_name, job_schema_name, parsed_job_name = sql_identifier.parse_fully_qualified_name(job_name) + job_database_name = job_database_name or self._database_name + job_schema_name = job_schema_name or self._schema_name + assert parsed_job_name is not None + job_fqn = identifier.get_schema_level_object_identifier( + job_database_name.identifier(), job_schema_name.identifier(), parsed_job_name.identifier() + ) + + # I/O starts here. + try: + X.write.copy_into_location( # type:ignore[call-overload] + location=input_stage_location, file_format_type="parquet", header=True, overwrite=True + ) + except Exception as e: + raise RuntimeError(f"Failed to process input data: {e}") + + try: + _, async_job_handle = self._service_client.execute_inference_job_service( + yaml_body=yaml_body, + compute_pool_name=compute_pool_name, + model_fqn=model_fqn, + version=version_name, + function_name=function_name, + job_fqn=job_fqn, + async_=async_, + replicas=replicas, + from_stage_path=input_stage_location, + statement_params=statement_params, + ) + result = async_job_handle.result() + except Exception: + # Server-side rejection leaves the staged input orphaned. Best-effort remove + # so files don't accumulate across failed launches. + try: + self._session.sql(f"REMOVE {input_stage_location}").collect(statement_params=statement_params) + except Exception as cleanup_err: + module_logger.warning(f"Failed to clean up staged input at {input_stage_location}: {cleanup_err}") + raise + + if job_name is not None: + assert parsed_job_name is not None + assert job_database_name is not None + assert job_schema_name is not None + fq_job_name = sql_identifier.get_fully_qualified_name(job_database_name, job_schema_name, parsed_job_name) + else: + response_msg = cast(str, cast(list[row.Row], result)[0][0]) + match = re.search(r"Batch inference job (\S+)", response_msg) + if match is None: + raise RuntimeError( + "batch inference job: failed to parse job name from server response. " f"Response: {response_msg}" + ) + fq_job_name = match.group(1) + + return job.MLJob( + id=fq_job_name, + session=self._session, + ) diff --git a/snowflake/ml/model/_client/ops/service_ops_test.py b/snowflake/ml/model/_client/ops/service_ops_test.py index d1379047..1d6b43b3 100644 --- a/snowflake/ml/model/_client/ops/service_ops_test.py +++ b/snowflake/ml/model/_client/ops/service_ops_test.py @@ -24,7 +24,7 @@ from snowflake.ml.model._signatures import core from snowflake.ml.test_utils import mock_data_frame, mock_session from snowflake.ml.test_utils.mock_progress import create_mock_progress_status -from snowflake.snowpark import Session, row +from snowflake.snowpark import Session, dataframe, row from snowflake.snowpark._internal import utils as snowpark_utils _DUMMY_SIG = { @@ -2758,6 +2758,196 @@ def test_check_if_service_exists_returns_false_when_no_match(self) -> None: ) self.assertFalse(exists) + def test_execute_inference_job_service_with_explicit_job_name(self) -> None: + m_async_job = self._create_mock_async_job() + m_async_job.result.return_value = [row.Row("Batch inference job IGNORED with model ...")] + input_df = mock.MagicMock(spec=dataframe.DataFrame) + fake_uuid = uuid.UUID("abcdef0123456789abcdef0123456789") + with ( + mock.patch.object( + self.m_ops._service_client, + "execute_inference_job_service", + return_value=("query_id", m_async_job), + ) as mock_execute, + mock.patch("snowflake.ml.model._client.ops.service_ops.uuid.uuid4", return_value=fake_uuid), + ): + result = self.m_ops.execute_inference_job_service( + X=input_df, + model_name=sql_identifier.SqlIdentifier("MODEL"), + version_name=sql_identifier.SqlIdentifier("V1"), + compute_pool_name=sql_identifier.SqlIdentifier("POOL"), + input_spec=batch_inference_specs.Input(params={"k": "v"}), + output_spec=batch_inference_specs.Output(stage_location="@DB.SCHEMA.STAGE/out"), + resources_spec=None, + inference_spec=None, + image_build_spec=None, + function_name="predict", + job_name="JOB", + replicas=2, + async_=False, + statement_params={"test": "1"}, + ) + expected_input_stage = f"@DB.SCHEMA.STAGE/out/_snowflake_temporary/{fake_uuid.hex}/" + input_df.write.copy_into_location.assert_called_once_with( + location=expected_input_stage, file_format_type="parquet", header=True, overwrite=True + ) + mock_execute.assert_called_once() + call_kwargs = mock_execute.call_args.kwargs + self.assertIn("input:", call_kwargs["yaml_body"]) + # output_spec.stage_location is normalized to add trailing / before being emitted. + self.assertIn("stage_location: '@DB.SCHEMA.STAGE/out/'", call_kwargs["yaml_body"]) + self.assertEqual(call_kwargs["model_fqn"], 'TEMP."test".MODEL') + self.assertEqual(call_kwargs["version"], sql_identifier.SqlIdentifier("V1")) + self.assertEqual(call_kwargs["function_name"], "predict") + self.assertEqual(call_kwargs["job_fqn"], 'TEMP."test".JOB') + self.assertEqual(call_kwargs["replicas"], 2) + self.assertFalse(call_kwargs["async_"]) + self.assertEqual(call_kwargs["from_stage_path"], expected_input_stage) + self.assertEqual(result.id, 'TEMP."test".JOB') + + def test_execute_inference_job_service_parses_server_generated_name(self) -> None: + m_async_job = self._create_mock_async_job() + m_async_job.result.return_value = [row.Row("Batch inference job DB.SCHEMA.SRV_GEN with model M ...")] + input_df = mock.MagicMock(spec=dataframe.DataFrame) + with mock.patch.object( + self.m_ops._service_client, + "execute_inference_job_service", + return_value=("query_id", m_async_job), + ): + result = self.m_ops.execute_inference_job_service( + X=input_df, + model_name=sql_identifier.SqlIdentifier("MODEL"), + version_name=sql_identifier.SqlIdentifier("V1"), + compute_pool_name=sql_identifier.SqlIdentifier("POOL"), + input_spec=None, + output_spec=batch_inference_specs.Output(stage_location="@DB.SCHEMA.STAGE/out/"), + resources_spec=None, + inference_spec=None, + image_build_spec=None, + function_name=None, + job_name=None, + replicas=None, + async_=True, + statement_params=None, + ) + self.assertEqual(result.id, "DB.SCHEMA.SRV_GEN") + + def test_execute_inference_job_service_unparsable_response_raises(self) -> None: + m_async_job = self._create_mock_async_job() + m_async_job.result.return_value = [row.Row("not the expected format")] + input_df = mock.MagicMock(spec=dataframe.DataFrame) + with mock.patch.object( + self.m_ops._service_client, + "execute_inference_job_service", + return_value=("query_id", m_async_job), + ): + with self.assertRaisesRegex(RuntimeError, "failed to parse job name"): + self.m_ops.execute_inference_job_service( + X=input_df, + model_name=sql_identifier.SqlIdentifier("MODEL"), + version_name=sql_identifier.SqlIdentifier("V1"), + compute_pool_name=sql_identifier.SqlIdentifier("POOL"), + input_spec=None, + output_spec=batch_inference_specs.Output(stage_location="@DB.SCHEMA.STAGE/out/"), + resources_spec=None, + inference_spec=None, + image_build_spec=None, + function_name=None, + job_name=None, + replicas=None, + async_=True, + statement_params=None, + ) + + def test_execute_inference_job_service_raises_when_copy_into_fails(self) -> None: + input_df = mock.MagicMock(spec=dataframe.DataFrame) + input_df.write.copy_into_location.side_effect = Exception("staging boom") + with mock.patch.object( + self.m_ops._service_client, + "execute_inference_job_service", + ) as mock_execute: + with self.assertRaisesRegex(RuntimeError, "Failed to process input data"): + self.m_ops.execute_inference_job_service( + X=input_df, + model_name=sql_identifier.SqlIdentifier("MODEL"), + version_name=sql_identifier.SqlIdentifier("V1"), + compute_pool_name=sql_identifier.SqlIdentifier("POOL"), + input_spec=None, + output_spec=batch_inference_specs.Output(stage_location="@DB.SCHEMA.STAGE/out/"), + resources_spec=None, + inference_spec=None, + image_build_spec=None, + function_name=None, + job_name=None, + replicas=None, + async_=True, + statement_params=None, + ) + mock_execute.assert_not_called() + + def test_execute_inference_job_service_cleans_up_staged_input_on_sql_failure(self) -> None: + """When the SQL call raises, the orphaned staged input must be REMOVED.""" + input_df = mock.MagicMock(spec=dataframe.DataFrame) + fake_uuid = uuid.UUID("abcdef0123456789abcdef0123456789") + expected_stage = f"@DB.SCHEMA.STAGE/out/_snowflake_temporary/{fake_uuid.hex}/" + with ( + mock.patch.object( + self.m_ops._service_client, + "execute_inference_job_service", + side_effect=RuntimeError("server rejected"), + ), + mock.patch("snowflake.ml.model._client.ops.service_ops.uuid.uuid4", return_value=fake_uuid), + mock.patch.object(self.m_ops._session, "sql") as mock_sql, + ): + with self.assertRaisesRegex(RuntimeError, "server rejected"): + self.m_ops.execute_inference_job_service( + X=input_df, + model_name=sql_identifier.SqlIdentifier("MODEL"), + version_name=sql_identifier.SqlIdentifier("V1"), + compute_pool_name=sql_identifier.SqlIdentifier("POOL"), + input_spec=None, + output_spec=batch_inference_specs.Output(stage_location="@DB.SCHEMA.STAGE/out/"), + resources_spec=None, + inference_spec=None, + image_build_spec=None, + function_name=None, + job_name=None, + replicas=None, + async_=True, + statement_params=None, + ) + mock_sql.assert_called_once_with(f"REMOVE {expected_stage}") + + def test_execute_inference_job_service_validates_job_name_before_staging(self) -> None: + """An invalid job_name must raise before COPY INTO so nothing orphans.""" + input_df = mock.MagicMock(spec=dataframe.DataFrame) + with ( + mock.patch( + "snowflake.ml.model._client.ops.service_ops.sql_identifier.parse_fully_qualified_name", + side_effect=ValueError("bad job_name"), + ), + mock.patch.object(self.m_ops._service_client, "execute_inference_job_service") as mock_execute, + ): + with self.assertRaisesRegex(ValueError, "bad job_name"): + self.m_ops.execute_inference_job_service( + X=input_df, + model_name=sql_identifier.SqlIdentifier("MODEL"), + version_name=sql_identifier.SqlIdentifier("V1"), + compute_pool_name=sql_identifier.SqlIdentifier("POOL"), + input_spec=None, + output_spec=batch_inference_specs.Output(stage_location="@DB.SCHEMA.STAGE/out/"), + resources_spec=None, + inference_spec=None, + image_build_spec=None, + function_name=None, + job_name="malformed", + replicas=None, + async_=True, + statement_params=None, + ) + input_df.write.copy_into_location.assert_not_called() + mock_execute.assert_not_called() + if __name__ == "__main__": absltest.main() diff --git a/snowflake/ml/model/_client/service/BUILD.bazel b/snowflake/ml/model/_client/service/BUILD.bazel index 4a3d7d84..32a7ee73 100644 --- a/snowflake/ml/model/_client/service/BUILD.bazel +++ b/snowflake/ml/model/_client/service/BUILD.bazel @@ -47,3 +47,25 @@ py_library( ":model_deployment_spec_schema", ], ) + +py_library( + name = "inference_job_service_spec", + srcs = ["inference_job_service_spec.py"], + deps = [ + "//snowflake/ml/model/_client/model:batch_inference_specs", + ], +) + +py_test( + name = "inference_job_service_spec_test", + srcs = ["inference_job_service_spec_test.py"], + tags = [ + "feature:model_serving", + "short_regress", + ], + deps = [ + ":inference_job_service_spec", + "//snowflake/ml/model:inference_engine", + "//snowflake/ml/model/_client/model:batch_inference_specs", + ], +) diff --git a/snowflake/ml/model/_client/service/inference_job_service_spec.py b/snowflake/ml/model/_client/service/inference_job_service_spec.py new file mode 100644 index 00000000..6ba198f3 --- /dev/null +++ b/snowflake/ml/model/_client/service/inference_job_service_spec.py @@ -0,0 +1,75 @@ +from typing import Any, Optional + +import yaml + +from snowflake.ml.model._client.model import batch_inference_specs + + +class InferenceJobServiceSpec: + """Builds the YAML body for ``EXECUTE INFERENCE JOB SERVICE``.""" + + def __init__(self) -> None: + self._input: Optional[dict[str, Any]] = None + self._output: Optional[dict[str, Any]] = None + self._resources: Optional[dict[str, Any]] = None + self._inference: Optional[dict[str, Any]] = None + self._image_build: Optional[dict[str, Any]] = None + + def clear(self) -> None: + self._input = None + self._output = None + self._resources = None + self._inference = None + self._image_build = None + + def add_input_spec(self, input_spec: batch_inference_specs.Input) -> "InferenceJobServiceSpec": + # ``params`` and ``column_handling`` are emitted as raw dicts; the + # server handles encoding for both. + self._input = input_spec.model_dump(mode="json", exclude_none=True) + return self + + def add_output_spec(self, output_spec: batch_inference_specs.Output) -> "InferenceJobServiceSpec": + self._output = output_spec.model_dump(mode="json", exclude_none=True) + return self + + def add_resources_spec(self, resources_spec: batch_inference_specs.Resources) -> "InferenceJobServiceSpec": + dumped = resources_spec.model_dump(mode="json", exclude_none=True) + self._resources = dumped if dumped else None + return self + + def add_inference_spec(self, inference_spec: batch_inference_specs.Inference) -> "InferenceJobServiceSpec": + dumped = inference_spec.model_dump(mode="json", exclude_none=True) + self._inference = dumped if dumped else None + return self + + def add_image_build_spec(self, image_build_spec: batch_inference_specs.ImageBuild) -> "InferenceJobServiceSpec": + self._image_build = image_build_spec.model_dump(mode="json", exclude_none=True) + return self + + def save(self) -> str: + """Return the YAML body as a string. + + Raises: + ValueError: If ``output`` has not been added. + + Returns: + YAML string of the batch inference spec body. + """ + if self._output is None: + raise ValueError("batch inference job: output spec is required. Call add_output_spec().") + body: dict[str, Any] = {"output": self._output} + if self._input is not None: + body["input"] = self._input + if self._resources is not None: + body["resources"] = self._resources + if self._inference is not None: + body["inference"] = self._inference + if self._image_build is not None: + body["image_build"] = self._image_build + # Order keys to match the design doc for readability: input, output, + # resources, inference, image_build. + ordered: dict[str, Any] = {} + for key in ("input", "output", "resources", "inference", "image_build"): + if key in body: + ordered[key] = body[key] + return yaml.safe_dump(ordered, sort_keys=False) diff --git a/snowflake/ml/model/_client/service/inference_job_service_spec_test.py b/snowflake/ml/model/_client/service/inference_job_service_spec_test.py new file mode 100644 index 00000000..6ebb00a0 --- /dev/null +++ b/snowflake/ml/model/_client/service/inference_job_service_spec_test.py @@ -0,0 +1,126 @@ +from typing import Any, cast + +import yaml +from absl.testing import absltest + +from snowflake.ml.model import inference_engine as inference_engine_module +from snowflake.ml.model._client.model import batch_inference_specs +from snowflake.ml.model._client.service import inference_job_service_spec + + +class InferenceJobServiceSpecTest(absltest.TestCase): + def _build( + self, + *, + with_input: bool = False, + with_resources: bool = False, + with_inference: bool = False, + with_image_build: bool = False, + ) -> dict[str, Any]: + builder = inference_job_service_spec.InferenceJobServiceSpec() + if with_input: + builder.add_input_spec( + batch_inference_specs.Input( + params={"temperature": 0.7}, + column_handling={ + "image_col": { + "input_format": batch_inference_specs.InputFormat.FULL_STAGE_PATH, + "convert_to": batch_inference_specs.FileEncoding.BASE64, + } + }, + partition_column="PART_COL", + ) + ) + builder.add_output_spec( + batch_inference_specs.Output( + stage_location="@DB.SCHEMA.STAGE/out/", + mode=batch_inference_specs.SaveMode.OVERWRITE, + ) + ) + if with_resources: + builder.add_resources_spec( + batch_inference_specs.Resources(cpu_requests="2", memory_requests="8GiB", gpu_requests=None) + ) + if with_inference: + builder.add_inference_spec( + batch_inference_specs.Inference( + num_workers=4, + max_batch_rows=2048, + engine_options=batch_inference_specs.EngineOptions( + engine=inference_engine_module.InferenceEngine.VLLM, + engine_args_override=["--max-num-seqs=128"], + ), + ) + ) + if with_image_build: + builder.add_image_build_spec( + batch_inference_specs.ImageBuild( + image_repo="DB.SCHEMA.REPO", + force_rebuild=True, + ) + ) + return cast(dict[str, Any], yaml.safe_load(builder.save())) + + def test_save_requires_output(self) -> None: + builder = inference_job_service_spec.InferenceJobServiceSpec() + with self.assertRaisesRegex(ValueError, "output spec is required"): + builder.save() + + def test_minimal_body_only_has_output(self) -> None: + body = self._build() + self.assertEqual(set(body.keys()), {"output"}) + self.assertEqual(body["output"], {"stage_location": "@DB.SCHEMA.STAGE/out/", "mode": "overwrite"}) + + def test_full_body_has_all_blocks_in_canonical_order(self) -> None: + builder = inference_job_service_spec.InferenceJobServiceSpec() + builder.add_image_build_spec(batch_inference_specs.ImageBuild(image_repo="DB.SCHEMA.REPO")) + builder.add_inference_spec(batch_inference_specs.Inference(num_workers=2)) + builder.add_resources_spec(batch_inference_specs.Resources(cpu_requests="1")) + builder.add_output_spec(batch_inference_specs.Output(stage_location="@stage/")) + builder.add_input_spec(batch_inference_specs.Input(params={"k": "v"})) + rendered = builder.save() + self.assertLess(rendered.index("input"), rendered.index("output")) + self.assertLess(rendered.index("output"), rendered.index("resources")) + self.assertLess(rendered.index("resources"), rendered.index("inference")) + self.assertLess(rendered.index("inference"), rendered.index("image_build")) + + def test_input_emits_raw_params_and_column_handling(self) -> None: + body = self._build(with_input=True) + self.assertEqual(body["input"]["params"], {"temperature": 0.7}) + self.assertEqual( + body["input"]["column_handling"], + {"image_col": {"input_format": "full_stage_path", "convert_to": "base64"}}, + ) + self.assertEqual(body["input"]["partition_column"], "PART_COL") + + def test_inference_engine_serializes_to_string_value(self) -> None: + body = self._build(with_inference=True) + self.assertEqual(body["inference"]["engine_options"]["engine"], "vllm") + self.assertEqual(body["inference"]["engine_options"]["engine_args_override"], ["--max-num-seqs=128"]) + + def test_resources_omits_none_fields(self) -> None: + body = self._build(with_resources=True) + self.assertEqual(body["resources"], {"cpu_requests": "2", "memory_requests": "8GiB"}) + + def test_image_build_emits_force_rebuild(self) -> None: + body = self._build(with_image_build=True) + self.assertEqual(body["image_build"], {"image_repo": "DB.SCHEMA.REPO", "force_rebuild": True}) + + def test_clear_resets_state(self) -> None: + builder = inference_job_service_spec.InferenceJobServiceSpec() + builder.add_output_spec(batch_inference_specs.Output(stage_location="@stage/")) + builder.add_input_spec(batch_inference_specs.Input(params={"a": 1})) + builder.clear() + with self.assertRaises(ValueError): + builder.save() + + def test_no_unexpected_top_level_keys(self) -> None: + body = self._build(with_input=True, with_resources=True, with_inference=True, with_image_build=True) + self.assertEqual( + set(body.keys()), + {"input", "output", "resources", "inference", "image_build"}, + ) + + +if __name__ == "__main__": + absltest.main() diff --git a/snowflake/ml/model/_client/sql/service.py b/snowflake/ml/model/_client/sql/service.py index 97176211..9ccc6121 100644 --- a/snowflake/ml/model/_client/sql/service.py +++ b/snowflake/ml/model/_client/sql/service.py @@ -116,6 +116,66 @@ def deploy_model( assert isinstance(async_job, snowpark.AsyncJob) return async_job.query_id, async_job + def execute_inference_job_service( + self, + *, + yaml_body: str, + compute_pool_name: sql_identifier.SqlIdentifier, + model_fqn: str, + version: Optional[sql_identifier.SqlIdentifier], + function_name: Optional[str], + job_fqn: Optional[str], + async_: bool, + replicas: Optional[int], + from_stage_path: str, + statement_params: Optional[dict[str, Any]] = None, + ) -> tuple[str, snowpark.AsyncJob]: + """Run an ``EXECUTE INFERENCE JOB SERVICE`` SQL command. + + Args: + yaml_body: Batch inference YAML body. Inlined as a single-quoted + SQL literal with embedded single quotes escaped. The server + does not accept bind parameters in the ``WITH SPECIFICATION`` + clause, and dollar-quoting would collide with any ``$$`` that + appears inside user-supplied content. + compute_pool_name: Compute pool identifier for ``IN COMPUTE POOL``. + model_fqn: Fully qualified model name for ``MODEL =``. + version: Optional model version for ``VERSION =``. + function_name: Optional model function for ``FUNCTION = '...'``. + job_fqn: Optional fully qualified job name for ``NAME =``. + async_: ``ASYNC = TRUE`` if true, ``ASYNC = FALSE`` if false. + replicas: Optional integer for ``REPLICAS =``. + from_stage_path: Stage path used as the ``FROM @`` source. + Must already start with ``@``. + statement_params: Optional statement params for telemetry. + + Returns: + Tuple of (query id, AsyncJob) for the submitted SQL. + """ + escaped_yaml = snowpark_utils.escape_single_quotes(yaml_body) # type: ignore[no-untyped-call] + clauses = [ + "EXECUTE INFERENCE JOB SERVICE", + f"IN COMPUTE POOL {compute_pool_name.identifier()}", + f"WITH SPECIFICATION '{escaped_yaml}'", + f"FROM {from_stage_path}", + f"MODEL = {model_fqn}", + ] + if version is not None: + clauses.append(f"VERSION = {version.identifier()}") + if function_name is not None: + escaped_fn = function_name.replace("'", "''") + clauses.append(f"FUNCTION = '{escaped_fn}'") + if job_fqn is not None: + clauses.append(f"NAME = {job_fqn}") + clauses.append(f"ASYNC = {'TRUE' if async_ else 'FALSE'}") + if replicas is not None: + clauses.append(f"REPLICAS = {replicas}") + sql_str = "\n".join(clauses) + logger.info(f"Executing inference job service: {sql_str}") + async_job = self._session.sql(sql_str).collect(block=False, statement_params=statement_params) + assert isinstance(async_job, snowpark.AsyncJob) + return async_job.query_id, async_job + def invoke_function_method( self, *, diff --git a/snowflake/ml/model/_client/sql/service_test.py b/snowflake/ml/model/_client/sql/service_test.py index 3030b652..6e2fe1ba 100644 --- a/snowflake/ml/model/_client/sql/service_test.py +++ b/snowflake/ml/model/_client/sql/service_test.py @@ -69,6 +69,170 @@ def test_deploy_model_inline_yaml(self) -> None: statement_params=m_statement_params, ) + def test_execute_inference_job_service_minimal(self) -> None: + m_statement_params = {"test": "1"} + m_async_job = mock.MagicMock(spec=snowpark.AsyncJob) + m_async_job.query_id = uuid.uuid4() + m_df = mock_data_frame.MockDataFrame( + collect_block=False, + collect_result=m_async_job, + collect_statement_params=m_statement_params, + ) + + expected = ( + "EXECUTE INFERENCE JOB SERVICE\n" + "IN COMPUTE POOL POOL\n" + "WITH SPECIFICATION 'minimal_body\n'\n" + "FROM @DB.SCHEMA.STAGE/out/_snowflake_temporary/abc123/\n" + "MODEL = DB.SCHEMA.MODEL\n" + "ASYNC = TRUE" + ) + self.m_session.add_mock_sql(query=expected, result=copy.deepcopy(m_df)) + c_session = cast(Session, self.m_session) + + service_sql.ServiceSQLClient( + c_session, + database_name=sql_identifier.SqlIdentifier("TEMP"), + schema_name=sql_identifier.SqlIdentifier("test", case_sensitive=True), + ).execute_inference_job_service( + yaml_body="minimal_body\n", + compute_pool_name=sql_identifier.SqlIdentifier("POOL"), + model_fqn="DB.SCHEMA.MODEL", + version=None, + function_name=None, + job_fqn=None, + async_=True, + replicas=None, + from_stage_path="@DB.SCHEMA.STAGE/out/_snowflake_temporary/abc123/", + statement_params=m_statement_params, + ) + + def test_execute_inference_job_service_all_clauses(self) -> None: + m_statement_params = {"test": "1"} + m_async_job = mock.MagicMock(spec=snowpark.AsyncJob) + m_async_job.query_id = uuid.uuid4() + m_df = mock_data_frame.MockDataFrame( + collect_block=False, + collect_result=m_async_job, + collect_statement_params=m_statement_params, + ) + + expected = ( + "EXECUTE INFERENCE JOB SERVICE\n" + "IN COMPUTE POOL POOL\n" + "WITH SPECIFICATION 'yaml_body\n'\n" + "FROM @stage/in/\n" + "MODEL = DB.SCHEMA.MODEL\n" + "VERSION = V1\n" + "FUNCTION = 'predict'\n" + "NAME = DB.SCHEMA.JOB\n" + "ASYNC = FALSE\n" + "REPLICAS = 2" + ) + self.m_session.add_mock_sql(query=expected, result=copy.deepcopy(m_df)) + c_session = cast(Session, self.m_session) + + service_sql.ServiceSQLClient( + c_session, + database_name=sql_identifier.SqlIdentifier("TEMP"), + schema_name=sql_identifier.SqlIdentifier("test", case_sensitive=True), + ).execute_inference_job_service( + yaml_body="yaml_body\n", + compute_pool_name=sql_identifier.SqlIdentifier("POOL"), + model_fqn="DB.SCHEMA.MODEL", + version=sql_identifier.SqlIdentifier("V1"), + function_name="predict", + job_fqn="DB.SCHEMA.JOB", + async_=False, + replicas=2, + from_stage_path="@stage/in/", + statement_params=m_statement_params, + ) + + def test_execute_inference_job_service_escapes_function_name(self) -> None: + m_statement_params = {"test": "1"} + m_async_job = mock.MagicMock(spec=snowpark.AsyncJob) + m_async_job.query_id = uuid.uuid4() + m_df = mock_data_frame.MockDataFrame( + collect_block=False, + collect_result=m_async_job, + collect_statement_params=m_statement_params, + ) + + expected = ( + "EXECUTE INFERENCE JOB SERVICE\n" + "IN COMPUTE POOL POOL\n" + "WITH SPECIFICATION 'yaml_body\n'\n" + "FROM @stage/\n" + "MODEL = M\n" + "FUNCTION = 'pred''ict'\n" + "ASYNC = TRUE" + ) + self.m_session.add_mock_sql(query=expected, result=copy.deepcopy(m_df)) + c_session = cast(Session, self.m_session) + + service_sql.ServiceSQLClient( + c_session, + database_name=sql_identifier.SqlIdentifier("TEMP"), + schema_name=sql_identifier.SqlIdentifier("test", case_sensitive=True), + ).execute_inference_job_service( + yaml_body="yaml_body\n", + compute_pool_name=sql_identifier.SqlIdentifier("POOL"), + model_fqn="M", + version=None, + function_name="pred'ict", + job_fqn=None, + async_=True, + replicas=None, + from_stage_path="@stage/", + statement_params=m_statement_params, + ) + + def test_execute_inference_job_service_escapes_yaml_body_quotes_and_dollar(self) -> None: + """YAML containing single quotes and ``$$`` must round-trip into a single-quoted SQL literal. + + Single quotes get backslash-escaped (snowpark_utils.escape_single_quotes); + ``$$`` passes through untouched because we use single-quoted, not dollar-quoted, + literals. + """ + m_statement_params = {"test": "1"} + m_async_job = mock.MagicMock(spec=snowpark.AsyncJob) + m_async_job.query_id = uuid.uuid4() + m_df = mock_data_frame.MockDataFrame( + collect_block=False, + collect_result=m_async_job, + collect_statement_params=m_statement_params, + ) + + yaml_body = "input:\n params:\n prompt: 'has $$ and ' quote'\n" + expected = ( + "EXECUTE INFERENCE JOB SERVICE\n" + "IN COMPUTE POOL POOL\n" + "WITH SPECIFICATION 'input:\n params:\n prompt: \\'has $$ and \\' quote\\'\n'\n" + "FROM @stage/\n" + "MODEL = M\n" + "ASYNC = TRUE" + ) + self.m_session.add_mock_sql(query=expected, result=copy.deepcopy(m_df)) + c_session = cast(Session, self.m_session) + + service_sql.ServiceSQLClient( + c_session, + database_name=sql_identifier.SqlIdentifier("TEMP"), + schema_name=sql_identifier.SqlIdentifier("test", case_sensitive=True), + ).execute_inference_job_service( + yaml_body=yaml_body, + compute_pool_name=sql_identifier.SqlIdentifier("POOL"), + model_fqn="M", + version=None, + function_name=None, + job_fqn=None, + async_=True, + replicas=None, + from_stage_path="@stage/", + statement_params=m_statement_params, + ) + def test_invoke_function_method(self) -> None: m_statement_params = {"test": "1"} m_df = mock_data_frame.MockDataFrame() diff --git a/snowflake/ml/model/_model_composer/BUILD.bazel b/snowflake/ml/model/_model_composer/BUILD.bazel index 6d233159..e86b54a5 100644 --- a/snowflake/ml/model/_model_composer/BUILD.bazel +++ b/snowflake/ml/model/_model_composer/BUILD.bazel @@ -2,6 +2,14 @@ load("//bazel:py_rules.bzl", "py_library", "py_test") package(default_visibility = ["//visibility:public"]) +py_library( + name = "huggingface_lazy_uploader", + srcs = ["huggingface_lazy_uploader.py"], + deps = [ + "//snowflake/ml/_internal:file_utils", + ], +) + py_library( name = "model_composer", srcs = ["model_composer.py"], @@ -11,11 +19,24 @@ py_library( "//snowflake/ml/_internal:file_utils", "//snowflake/ml/model:model_signature", "//snowflake/ml/model:type_hints", + "//snowflake/ml/model/_model_composer:huggingface_lazy_uploader", "//snowflake/ml/model/_model_composer/model_manifest", "//snowflake/ml/model/_packager:model_packager", ], ) +py_test( + name = "huggingface_lazy_uploader_test", + srcs = ["huggingface_lazy_uploader_test.py"], + tags = [ + "feature:model_registry", + "short_regress", + ], + deps = [ + ":huggingface_lazy_uploader", + ], +) + py_test( name = "model_composer_test", srcs = ["model_composer_test.py"], @@ -24,6 +45,7 @@ py_test( "short_regress", ], deps = [ + ":huggingface_lazy_uploader", ":model_composer", "//snowflake/ml/_internal:env_utils", "//snowflake/ml/_internal:file_utils", diff --git a/snowflake/ml/model/_model_composer/huggingface_lazy_uploader.py b/snowflake/ml/model/_model_composer/huggingface_lazy_uploader.py new file mode 100644 index 00000000..336431dd --- /dev/null +++ b/snowflake/ml/model/_model_composer/huggingface_lazy_uploader.py @@ -0,0 +1,244 @@ +import concurrent.futures +import contextlib +import logging +import os +import pathlib +import shutil +import tempfile +import threading +import time +from dataclasses import dataclass +from typing import Any, Generator, Optional, Union +from urllib import parse + +from snowflake import snowpark +from snowflake.ml._internal import file_utils + +logger = logging.getLogger(__name__) + +DEFAULT_MAX_WORKERS = 8 +DISK_SAFETY_MARGIN = 0.9 +DEFAULT_DISK_BUDGET_WAIT_TIMEOUT_SECONDS = 3600 +_LAZY_UPLOAD_TEMP_PREFIX = "snowml_hf_lazy_upload_" + + +def _lazy_upload_temp_root() -> pathlib.Path: + """Return the filesystem root where lazy-upload temp directories are created.""" + return pathlib.Path(tempfile.gettempdir()) + + +def _validate_repo_relative_path(path: str) -> None: + """Reject repo-relative paths that could escape the intended upload root.""" + if not path or path.startswith("/"): + raise ValueError( + "model upload: invalid HuggingFace repository file path. " + f"Expected a relative path within the repository, got {path!r}." + ) + parts = pathlib.PurePosixPath(path).parts + if ".." in parts: + raise ValueError( + "model upload: invalid HuggingFace repository file path. " + f"Path must not contain parent-directory segments: {path!r}." + ) + + +@dataclass(frozen=True) +class LazyHFUpload: + """Deferred HuggingFace repository upload metadata.""" + + download_kwargs: dict[str, Any] + files: list[str] + file_sizes: dict[str, int] + relative_stage_dir: pathlib.PurePosixPath + download_token: Optional[str] = None + + +class DiskBudget: + """Tracks available disk bytes for concurrent HuggingFace downloads.""" + + def __init__( + self, + available_bytes: int, + *, + acquire_timeout_seconds: float = DEFAULT_DISK_BUDGET_WAIT_TIMEOUT_SECONDS, + ) -> None: + self._available_bytes = available_bytes + self._condition = threading.Condition() + self._acquire_timeout_seconds = acquire_timeout_seconds + + @staticmethod + def from_files( + file_sizes: dict[str, int], + files: list[str], + ) -> "DiskBudget": + """Validate disk space and return a budget for concurrent downloads.""" + temp_root = _lazy_upload_temp_root() + usage = shutil.disk_usage(temp_root) + budget_bytes = int(usage.free * DISK_SAFETY_MARGIN) + + if not files: + return DiskBudget(budget_bytes) + + largest_file = max(files, key=lambda filename: file_sizes.get(filename, 0)) + largest_file_size = file_sizes.get(largest_file, 0) + if largest_file_size > budget_bytes: + raise ValueError( + "model upload: insufficient disk space to download HuggingFace model files. " + f"The largest file ({largest_file}) requires {_format_size(largest_file_size)} " + f"but only {_format_size(budget_bytes)} is available." + ) + + logger.info( + "HuggingFace lazy upload disk budget: %s available at %s (%s total free)", + _format_size(budget_bytes), + temp_root, + _format_size(usage.free), + ) + return DiskBudget(budget_bytes) + + def acquire(self, nbytes: int) -> None: + """Block until nbytes of disk budget are available, then reserve them.""" + with self._condition: + deadline = time.monotonic() + self._acquire_timeout_seconds + while self._available_bytes < nbytes: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError( + "model upload: timed out waiting for disk space to download HuggingFace model files. " + f"Required {_format_size(nbytes)} but only {_format_size(self._available_bytes)} " + "was available." + ) + self._condition.wait(timeout=remaining) + self._available_bytes -= nbytes + + def release(self, nbytes: int) -> None: + """Return nbytes to the disk budget.""" + with self._condition: + self._available_bytes += nbytes + self._condition.notify_all() + + @contextlib.contextmanager + def reserve(self, nbytes: int) -> Generator[None, None, None]: + """Reserve disk budget for the duration of the context.""" + self.acquire(nbytes) + try: + yield + finally: + self.release(nbytes) + + +def _format_size(num_bytes: int) -> str: + if num_bytes >= 1024**3: + return f"{num_bytes / (1024**3):.2f} GiB" + if num_bytes >= 1024**2: + return f"{num_bytes / (1024**2):.2f} MiB" + if num_bytes >= 1024: + return f"{num_bytes / 1024:.2f} KiB" + return f"{num_bytes} B" + + +def _upload_single_file( + *, + session: snowpark.Session, + stage_path: Union[pathlib.PurePosixPath, parse.ParseResult], + lazy: LazyHFUpload, + filename: str, + disk_budget: DiskBudget, + statement_params: Optional[dict[str, Any]], +) -> None: + """Download one HuggingFace file, upload it to stage, and free local disk space.""" + import huggingface_hub as hf_hub + + _validate_repo_relative_path(filename) + + file_size = lazy.file_sizes.get(filename, 0) + if file_size == 0: + logger.warning( + "HuggingFace file %s has no known size; disk budget will not be reserved for it.", + filename, + ) + + relative_path = lazy.relative_stage_dir / pathlib.PurePosixPath(filename) + tmp_dir = tempfile.mkdtemp(prefix=_LAZY_UPLOAD_TEMP_PREFIX, dir=str(_lazy_upload_temp_root())) + try: + stage_dir_path = file_utils._resolve_stage_dir_path(stage_path, relative_path) + with disk_budget.reserve(file_size): + local_path = hf_hub.hf_hub_download( + filename=filename, + local_dir=tmp_dir, + token=lazy.download_token, + **lazy.download_kwargs, + ) + file_utils.upload_file_to_stage( + session, + local_path, + stage_dir_path, + statement_params=statement_params, + ) + if os.path.isfile(local_path): + os.remove(local_path) + finally: + shutil.rmtree(tmp_dir, ignore_errors=True) + + +def stream_upload( + session: snowpark.Session, + stage_path: Union[pathlib.PurePosixPath, parse.ParseResult], + lazy: LazyHFUpload, + *, + max_workers: int = DEFAULT_MAX_WORKERS, + statement_params: Optional[dict[str, Any]] = None, +) -> None: + """Download HuggingFace repo files in parallel and upload each to stage. + + Args: + session: Snowpark Session. + stage_path: Base path in the stage. + lazy: Metadata describing which files to upload and where. + max_workers: Maximum number of concurrent download-and-upload workers. + statement_params: Statement Params. + + Raises: + Exception: If disk space is insufficient or any file download or stage upload fails. + """ + total_files = len(lazy.files) + if total_files == 0: + return + + for filename in lazy.files: + _validate_repo_relative_path(filename) + + disk_budget = DiskBudget.from_files(lazy.file_sizes, lazy.files) + logger.info( + "Starting parallel upload of %s HuggingFace files with %s workers", + total_files, + max_workers, + ) + + with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = { + executor.submit( + _upload_single_file, + session=session, + stage_path=stage_path, + lazy=lazy, + filename=filename, + disk_budget=disk_budget, + statement_params=statement_params, + ): filename + for filename in lazy.files + } + for future in concurrent.futures.as_completed(futures): + filename = futures[future] + try: + future.result() + except Exception: + for pending_future in futures: + pending_future.cancel() + raise + file_size = lazy.file_sizes.get(filename, 0) + logger.info( + "Uploaded HuggingFace file %s (size: %s)", + filename, + _format_size(file_size), + ) diff --git a/snowflake/ml/model/_model_composer/huggingface_lazy_uploader_test.py b/snowflake/ml/model/_model_composer/huggingface_lazy_uploader_test.py new file mode 100644 index 00000000..7be3e50c --- /dev/null +++ b/snowflake/ml/model/_model_composer/huggingface_lazy_uploader_test.py @@ -0,0 +1,211 @@ +import os +import pathlib +from unittest import mock + +from absl.testing import absltest + +from snowflake.ml.model._model_composer import huggingface_lazy_uploader + + +class HuggingFaceLazyUploaderTest(absltest.TestCase): + @mock.patch("huggingface_hub.utils.get_token", return_value=None) + @mock.patch("snowflake.ml._internal.file_utils.upload_file_to_stage") + @mock.patch("huggingface_hub.hf_hub_download") + def test_stream_upload_puts_each_file( + self, + mock_hf_hub_download: mock.Mock, + mock_upload_file_to_stage: mock.Mock, + mock_get_token: mock.Mock, + ) -> None: + """Each HuggingFace file should be downloaded, uploaded, and removed locally.""" + + def _fake_download(*, filename: str, local_dir: str, **kwargs: object) -> str: + local_path = os.path.join(local_dir, filename) + os.makedirs(os.path.dirname(local_path), exist_ok=True) + with open(local_path, "w", encoding="utf-8") as file: + file.write("data") + return local_path + + mock_hf_hub_download.side_effect = _fake_download + + lazy = huggingface_lazy_uploader.LazyHFUpload( + download_kwargs={"repo_id": "org/model", "revision": None}, + files=["config.json", "weights/model.safetensors"], + file_sizes={"config.json": 100, "weights/model.safetensors": 200}, + relative_stage_dir=pathlib.PurePosixPath("model", "models", "my_model", "model"), + ) + mock_session = mock.Mock() + + huggingface_lazy_uploader.stream_upload( + mock_session, + pathlib.PurePosixPath("@stage/path"), + lazy, + max_workers=2, + ) + + self.assertEqual(mock_hf_hub_download.call_count, 2) + self.assertEqual(mock_upload_file_to_stage.call_count, 2) + uploaded_local_paths = [call.args[1] for call in mock_upload_file_to_stage.call_args_list] + for local_path in uploaded_local_paths: + self.assertFalse(os.path.exists(local_path)) + + stage_dirs = {call.args[2] for call in mock_upload_file_to_stage.call_args_list} + self.assertEqual( + stage_dirs, + { + "@stage/path/model/models/my_model/model", + "@stage/path/model/models/my_model/model/weights", + }, + ) + + @mock.patch("snowflake.ml.model._model_composer.huggingface_lazy_uploader._lazy_upload_temp_root") + @mock.patch("shutil.disk_usage") + def test_create_disk_budget_raises_when_largest_file_exceeds_available_space( + self, + mock_disk_usage: mock.Mock, + mock_temp_root: mock.Mock, + ) -> None: + mock_temp_root.return_value = pathlib.Path("/tmp") + mock_disk_usage.return_value = mock.Mock(free=1000, total=1000, used=0) + + with self.assertRaises(ValueError) as error_context: + huggingface_lazy_uploader.DiskBudget.from_files( + {"config.json": 100, "weights/model.safetensors": 1800}, + ["config.json", "weights/model.safetensors"], + ) + self.assertEqual( + str(error_context.exception), + "model upload: insufficient disk space to download HuggingFace model files. " + "The largest file (weights/model.safetensors) requires 1.76 KiB but only 900 B is available.", + ) + mock_disk_usage.assert_called_once_with(pathlib.Path("/tmp")) + + def test_disk_budget_reserve_releases_on_exit(self) -> None: + budget = huggingface_lazy_uploader.DiskBudget(100) + with budget.reserve(80): + pass + with budget.reserve(100): + pass + + def test_disk_budget_reserve_releases_on_exception(self) -> None: + budget = huggingface_lazy_uploader.DiskBudget(100) + with self.assertRaises(RuntimeError): + with budget.reserve(80): + raise RuntimeError("upload failed") + with budget.reserve(100): + pass + + def test_disk_budget_acquire_raises_on_timeout(self) -> None: + budget = huggingface_lazy_uploader.DiskBudget(100, acquire_timeout_seconds=0.1) + with budget.reserve(80): + with self.assertRaises(TimeoutError) as error_context: + budget.acquire(50) + self.assertEqual( + str(error_context.exception), + "model upload: timed out waiting for disk space to download HuggingFace model files. " + "Required 50 B but only 20 B was available.", + ) + + @mock.patch("huggingface_hub.utils.get_token", return_value=None) + @mock.patch("snowflake.ml._internal.file_utils.upload_file_to_stage") + @mock.patch("huggingface_hub.hf_hub_download") + @mock.patch("shutil.disk_usage") + def test_stream_upload_throttles_concurrent_large_downloads( + self, + mock_disk_usage: mock.Mock, + mock_hf_hub_download: mock.Mock, + mock_upload_file_to_stage: mock.Mock, + mock_get_token: mock.Mock, + ) -> None: + """Workers should wait for disk budget before downloading large files.""" + mock_disk_usage.return_value = mock.Mock(free=300, total=300, used=0) + active_downloads = 0 + max_active_downloads = 0 + + def _fake_download(*, filename: str, local_dir: str, **kwargs: object) -> str: + nonlocal active_downloads, max_active_downloads + active_downloads += 1 + max_active_downloads = max(max_active_downloads, active_downloads) + local_path = os.path.join(local_dir, filename) + os.makedirs(os.path.dirname(local_path), exist_ok=True) + with open(local_path, "w", encoding="utf-8") as file: + file.write("data") + active_downloads -= 1 + return local_path + + mock_hf_hub_download.side_effect = _fake_download + + lazy = huggingface_lazy_uploader.LazyHFUpload( + download_kwargs={"repo_id": "org/model", "revision": None}, + files=["shard-0.bin", "shard-1.bin", "shard-2.bin"], + file_sizes={"shard-0.bin": 150, "shard-1.bin": 150, "shard-2.bin": 150}, + relative_stage_dir=pathlib.PurePosixPath("model"), + ) + + huggingface_lazy_uploader.stream_upload( + mock.Mock(), + pathlib.PurePosixPath("@stage/path"), + lazy, + max_workers=3, + ) + + self.assertEqual(mock_hf_hub_download.call_count, 3) + self.assertEqual(max_active_downloads, 1) + + def test_stream_upload_rejects_path_traversal(self) -> None: + lazy = huggingface_lazy_uploader.LazyHFUpload( + download_kwargs={"repo_id": "org/model", "revision": None}, + files=["../../functions/__call__.py"], + file_sizes={"../../functions/__call__.py": 100}, + relative_stage_dir=pathlib.PurePosixPath("model"), + ) + with self.assertRaises(ValueError) as error_context: + huggingface_lazy_uploader.stream_upload( + mock.Mock(), + pathlib.PurePosixPath("@stage/path"), + lazy, + ) + self.assertEqual( + str(error_context.exception), + "model upload: invalid HuggingFace repository file path. " + "Path must not contain parent-directory segments: '../../functions/__call__.py'.", + ) + + @mock.patch("huggingface_hub.utils.get_token", return_value=None) + @mock.patch("snowflake.ml._internal.file_utils.upload_file_to_stage") + @mock.patch("huggingface_hub.hf_hub_download") + def test_stream_upload_uses_download_token_from_lazy_metadata( + self, + mock_hf_hub_download: mock.Mock, + mock_upload_file_to_stage: mock.Mock, + mock_get_token: mock.Mock, + ) -> None: + local_path = "/tmp/config.json" + mock_hf_hub_download.return_value = local_path + + lazy = huggingface_lazy_uploader.LazyHFUpload( + download_kwargs={"repo_id": "org/model", "revision": None}, + files=["config.json"], + file_sizes={"config.json": 100}, + relative_stage_dir=pathlib.PurePosixPath("model"), + download_token="hf_test_token", + ) + + huggingface_lazy_uploader.stream_upload( + mock.Mock(), + pathlib.PurePosixPath("@stage/path"), + lazy, + ) + + mock_hf_hub_download.assert_called_once_with( + filename="config.json", + local_dir=mock.ANY, + token="hf_test_token", + repo_id="org/model", + revision=None, + ) + mock_upload_file_to_stage.assert_called_once() + + +if __name__ == "__main__": + absltest.main() diff --git a/snowflake/ml/model/_model_composer/model_composer.py b/snowflake/ml/model/_model_composer/model_composer.py index fe6549d2..f3d08612 100644 --- a/snowflake/ml/model/_model_composer/model_composer.py +++ b/snowflake/ml/model/_model_composer/model_composer.py @@ -10,6 +10,7 @@ from snowflake.ml._internal.lineage import lineage_utils from snowflake.ml.data import data_source from snowflake.ml.model import model_signature, type_hints as model_types +from snowflake.ml.model._model_composer import huggingface_lazy_uploader from snowflake.ml.model._model_composer.model_manifest import model_manifest from snowflake.ml.model._packager import model_packager from snowflake.ml.model._packager.model_meta import model_meta @@ -159,6 +160,13 @@ def save( stage_path=self.stage_path, statement_params=self._statement_params, ) + if model_metadata._lazy_hf_upload is not None: + huggingface_lazy_uploader.stream_upload( + self.session, + self.stage_path, + model_metadata._lazy_hf_upload, + statement_params=self._statement_params, + ) return model_metadata @staticmethod diff --git a/snowflake/ml/model/_model_composer/model_composer_test.py b/snowflake/ml/model/_model_composer/model_composer_test.py index 990c9a0d..58f6f792 100644 --- a/snowflake/ml/model/_model_composer/model_composer_test.py +++ b/snowflake/ml/model/_model_composer/model_composer_test.py @@ -12,7 +12,7 @@ from snowflake.ml._internal import env_utils, file_utils from snowflake.ml.model import type_hints as model_types -from snowflake.ml.model._model_composer import model_composer +from snowflake.ml.model._model_composer import huggingface_lazy_uploader, model_composer from snowflake.ml.model._packager import model_packager from snowflake.ml.modeling.linear_model import ( # type:ignore[attr-defined] LinearRegression, @@ -37,6 +37,7 @@ def test_save_interface(self, params: dict[str, Any]) -> None: mock_pk = mock.MagicMock() mock_pk.meta = mock.MagicMock() mock_pk.meta.signatures = mock.MagicMock() + mock_pk.meta._lazy_hf_upload = None if params["use_save_location"]: temp_dir = tempfile.mkdtemp() else: @@ -131,6 +132,45 @@ def test_save_interface(self, params: dict[str, Any]) -> None: mock_save.assert_called_once() mock_manifest_save.assert_called_once() + def test_save_calls_stream_upload_when_lazy_hf_upload_set(self) -> None: + m_session = mock_session.MockSession(conn=None, test_case=self) + c_session = cast(Session, m_session) + stage_path = '@"db"."schema"."stage"' + + lazy_upload = huggingface_lazy_uploader.LazyHFUpload( + download_kwargs={"repo_id": "org/model", "revision": None}, + files=["config.json"], + file_sizes={"config.json": 100}, + relative_stage_dir=pathlib.PurePosixPath("model", "models", "model1", "model"), + ) + + mock_pk = mock.MagicMock() + mock_pk.meta = mock.MagicMock() + mock_pk.meta.signatures = mock.MagicMock() + mock_pk.meta._lazy_hf_upload = lazy_upload + + m = model_composer.ModelComposer(session=c_session, stage_path=stage_path) + with open(os.path.join(m.packager_workspace_path, "model.yaml"), "w", encoding="utf-8") as f: + f.write("") + m.packager = mock_pk + + with mock.patch.object(m.packager, "save", return_value=mock_pk.meta), mock.patch.object( + m.manifest, "save" + ), mock.patch.object(file_utils, "upload_directory_to_stage"), mock.patch.object( + huggingface_lazy_uploader, "stream_upload" + ) as mock_stream_upload, mock.patch.object( + env_utils, + "get_matched_package_versions_in_information_schema", + return_value={env_utils.SNOWPARK_ML_PKG_NAME: []}, + ): + m.save(name="model1", model=LinearRegression()) + mock_stream_upload.assert_called_once_with( + c_session, + pathlib.PurePosixPath(stage_path), + lazy_upload, + statement_params=None, + ) + def test_load(self) -> None: m_options = model_types.PyTorchLoadOptions(use_gpu=False) with mock.patch.object(model_packager.ModelPackager, "load") as mock_load: diff --git a/snowflake/ml/model/_packager/model_handlers/BUILD.bazel b/snowflake/ml/model/_packager/model_handlers/BUILD.bazel index 4419e587..16600e78 100644 --- a/snowflake/ml/model/_packager/model_handlers/BUILD.bazel +++ b/snowflake/ml/model/_packager/model_handlers/BUILD.bazel @@ -542,6 +542,7 @@ py_library( "//snowflake/ml/model:model_signature", "//snowflake/ml/model:openai_signatures", "//snowflake/ml/model:type_hints", + "//snowflake/ml/model/_model_composer:huggingface_lazy_uploader", "//snowflake/ml/model/_packager/model_env", "//snowflake/ml/model/_packager/model_handlers_migrator:base_migrator", "//snowflake/ml/model/_packager/model_meta", @@ -574,6 +575,7 @@ py_library( "//snowflake/ml/model:custom_model", "//snowflake/ml/model:model_signature", "//snowflake/ml/model:type_hints", + "//snowflake/ml/model/_model_composer:huggingface_lazy_uploader", "//snowflake/ml/model/_packager/model_env", "//snowflake/ml/model/_packager/model_handlers_migrator:base_migrator", "//snowflake/ml/model/_packager/model_meta", diff --git a/snowflake/ml/model/_packager/model_handlers/huggingface/_handler.py b/snowflake/ml/model/_packager/model_handlers/huggingface/_handler.py index f706db16..1fec7b74 100644 --- a/snowflake/ml/model/_packager/model_handlers/huggingface/_handler.py +++ b/snowflake/ml/model/_packager/model_handlers/huggingface/_handler.py @@ -1,5 +1,6 @@ import logging import os +import pathlib import shutil import warnings from typing import TYPE_CHECKING, Any, Callable, Optional, Union, cast, final @@ -11,6 +12,7 @@ from snowflake.ml._internal import type_utils from snowflake.ml.model import custom_model, model_signature, type_hints as model_types +from snowflake.ml.model._model_composer import huggingface_lazy_uploader from snowflake.ml.model._packager.model_env import model_env from snowflake.ml.model._packager.model_handlers import _base, _utils as handlers_utils from snowflake.ml.model._packager.model_handlers.huggingface import _utils as _hf_utils @@ -241,7 +243,26 @@ def save_model( "wb", ) as f: cloudpickle.dump(model, f) - if model.repo_snapshot_dir: + lazy_repo_files = getattr(model, "_lazy_repo_files", None) + if lazy_repo_files is not None: + lazy_download_kwargs = getattr(model, "_lazy_download_kwargs", None) + if lazy_download_kwargs is None: + raise ValueError("HuggingFace model metadata is incomplete; cannot determine download parameters.") + lazy_file_sizes = getattr(model, "_lazy_file_sizes", None) or {} + model_meta._lazy_hf_upload = huggingface_lazy_uploader.LazyHFUpload( + download_kwargs=lazy_download_kwargs, + files=lazy_repo_files, + file_sizes=lazy_file_sizes, + relative_stage_dir=pathlib.PurePosixPath( + "model", + "models", + name, + cls.MODEL_BLOB_FILE_OR_DIR, + ), + download_token=_hf_utils.download_token_for_lazy_upload(model), + ) + is_repo_downloaded = True + elif model.repo_snapshot_dir: logger.info("model's repo_snapshot_dir is available, copying snapshot") shutil.copytree( model.repo_snapshot_dir, @@ -392,7 +413,7 @@ def _create_pipeline_from_dir( m = transformers.pipeline( model_blob_options["task"], model=model_blob_file_or_dir_path, - trust_remote_code=True, + trust_remote_code=False, torch_dtype="auto", **additional_pipeline_params, **device_config, diff --git a/snowflake/ml/model/_packager/model_handlers/huggingface/_utils.py b/snowflake/ml/model/_packager/model_handlers/huggingface/_utils.py index 00d18d59..fe0138aa 100644 --- a/snowflake/ml/model/_packager/model_handlers/huggingface/_utils.py +++ b/snowflake/ml/model/_packager/model_handlers/huggingface/_utils.py @@ -1,11 +1,36 @@ import logging -from typing import Any +from typing import Any, Optional import numpy as np import pandas as pd logger = logging.getLogger(__name__) +_LAZY_UPLOAD_SECRET_ERROR = ( + "model upload: HuggingFace lazy upload cannot resolve auth from a Snowflake secret during local logging. " + "Set the HF_TOKEN environment variable or pass a HuggingFace token when constructing the model." +) + + +def download_token_for_lazy_upload(model: object) -> Optional[str]: + """Return the HuggingFace token to use for lazy upload from a wrapper model. + + Args: + model: A HuggingFace wrapper with optional ``token_or_secret`` and ``secret_identifier`` attributes. + + Returns: + The HuggingFace auth token when the wrapper was constructed with a plain token, otherwise None. + + Raises: + ValueError: If the wrapper was constructed with a Snowflake secret reference. + """ + if getattr(model, "secret_identifier", None) is not None: + raise ValueError(_LAZY_UPLOAD_SECRET_ERROR) + token_or_secret = getattr(model, "token_or_secret", None) + if isinstance(token_or_secret, str): + return token_or_secret + return None + def _resolve_chat_params(row: pd.Series, kwargs: dict[str, Any]) -> dict[str, Any]: """Resolve chat completion params from kwargs (ParamSpec) or DataFrame row columns.""" diff --git a/snowflake/ml/model/_packager/model_handlers/sentence_transformers.py b/snowflake/ml/model/_packager/model_handlers/sentence_transformers.py index 2d1e4453..792df6cd 100644 --- a/snowflake/ml/model/_packager/model_handlers/sentence_transformers.py +++ b/snowflake/ml/model/_packager/model_handlers/sentence_transformers.py @@ -2,19 +2,23 @@ import json import logging import os +import pathlib import shutil from importlib import metadata as importlib_metadata from typing import TYPE_CHECKING, Any, Callable, Optional, Sequence, Union, cast, final import cloudpickle +import numpy as np import pandas as pd from packaging import version from typing_extensions import TypeGuard, Unpack from snowflake.ml._internal import type_utils from snowflake.ml.model import custom_model, model_signature, type_hints as model_types +from snowflake.ml.model._model_composer import huggingface_lazy_uploader from snowflake.ml.model._packager.model_env import model_env from snowflake.ml.model._packager.model_handlers import _base, _utils as handlers_utils +from snowflake.ml.model._packager.model_handlers.huggingface import _utils as _hf_utils from snowflake.ml.model._packager.model_handlers_migrator import base_migrator from snowflake.ml.model._packager.model_meta import ( model_blob_meta, @@ -129,6 +133,47 @@ def _capture_pretrained_model_name(model: "sentence_transformers.SentenceTransfo return None +def _is_null_sentence(value: Any) -> bool: + """Return True when a sentence input should produce a null embedding output.""" + if value is None: + return True + if isinstance(value, (list, dict, np.ndarray)): + return False + return bool(pd.isna(value)) + + +def _encode_sentences_with_nulls( + sentences: Sequence[Any], + method_to_call: Callable[..., Any], + encode_kwargs: dict[str, Any], +) -> list[Any]: + """Encode non-null sentences and return None for null sentence inputs.""" + if not sentences: + return [] + + null_indices: list[int] = [] + non_null_indices: list[int] = [] + non_null_sentences: list[Any] = [] + for index, sentence in enumerate(sentences): + if _is_null_sentence(sentence): + null_indices.append(index) + else: + non_null_indices.append(index) + non_null_sentences.append(sentence) + + if not null_indices: + return cast(list[Any], method_to_call(list(sentences), **encode_kwargs).tolist()) + + if len(null_indices) == len(sentences): + return [None] * len(sentences) + + outputs: list[Any] = [None] * len(sentences) + encoded = method_to_call(non_null_sentences, **encode_kwargs).tolist() + for index, embedding in zip(non_null_indices, encoded): + outputs[index] = embedding + return outputs + + def _auto_infer_signature( target_method: str, embedding_dim: int, @@ -507,6 +552,7 @@ def save_model( if is_wrapper: assert isinstance(model, snowml_huggingface.SentenceTransformer) cls._save_wrapper_model( + name=name, model=model, model_meta=model_meta, model_blob_path=model_blob_path, @@ -515,7 +561,9 @@ def save_model( target_methods=target_methods, ) blob_options["model"] = model.model - if model.compute_pool_for_log is None and model.repo_snapshot_dir is not None: + if model.compute_pool_for_log is None and ( + getattr(model, "_lazy_repo_files", None) is not None or model.repo_snapshot_dir is not None + ): blob_options["is_repo_downloaded"] = True if batch_size is not None: blob_options["batch_size"] = batch_size @@ -574,6 +622,7 @@ def save_model( @classmethod def _save_wrapper_model( cls, + name: str, model: snowml_huggingface.SentenceTransformer, model_meta: model_meta_api.ModelMetadata, model_blob_path: str, @@ -582,15 +631,19 @@ def _save_wrapper_model( is_sub_model: Optional[bool] = False, target_methods: Optional[Sequence[str]] = None, ) -> None: - """Save a SentenceTransformer wrapper using pickle + copytree. + """Save a SentenceTransformer wrapper using pickle and optional snapshot copy. Args: + name: Model blob name used for lazy-upload stage paths. model: The wrapper instance. model_meta: Model metadata to update with signatures. model_blob_path: Base path for model blobs. batch_size: Batch size for inference. is_sub_model: Whether this is a sub-model. target_methods: Optional list of target methods for signature inference. + + Raises: + ValueError: If lazy-upload metadata on the wrapper is incomplete. """ model_blob_file_or_dir = os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR) os.makedirs(model_blob_file_or_dir, exist_ok=True) @@ -599,13 +652,32 @@ def _save_wrapper_model( with open(pickle_file, "wb") as f: cloudpickle.dump(model, f) - if model.compute_pool_for_log is None and model.repo_snapshot_dir: - logger.info("Wrapper repo_snapshot_dir is available, copying snapshot") - shutil.copytree( - model.repo_snapshot_dir, - model_blob_file_or_dir, - dirs_exist_ok=True, - ) + if model.compute_pool_for_log is None: + lazy_repo_files = getattr(model, "_lazy_repo_files", None) + if lazy_repo_files is not None: + lazy_download_kwargs = getattr(model, "_lazy_download_kwargs", None) + if lazy_download_kwargs is None: + raise ValueError("HuggingFace model metadata is incomplete; cannot determine download parameters.") + lazy_file_sizes = getattr(model, "_lazy_file_sizes", None) or {} + model_meta._lazy_hf_upload = huggingface_lazy_uploader.LazyHFUpload( + download_kwargs=lazy_download_kwargs, + files=lazy_repo_files, + file_sizes=lazy_file_sizes, + relative_stage_dir=pathlib.PurePosixPath( + "model", + "models", + name, + cls.MODEL_BLOB_FILE_OR_DIR, + ), + download_token=_hf_utils.download_token_for_lazy_upload(model), + ) + elif model.repo_snapshot_dir: + logger.info("Wrapper repo_snapshot_dir is available, copying snapshot") + shutil.copytree( + model.repo_snapshot_dir, + model_blob_file_or_dir, + dirs_exist_ok=True, + ) if not is_sub_model and model.compute_pool_for_log is None: cls._set_signatures_from_config( @@ -740,7 +812,7 @@ def get_prediction( encode_kwargs: dict[str, Any] = {} if batch_size is not None: encode_kwargs["batch_size"] = batch_size - return pd.DataFrame({0: method_to_call(X_list, **encode_kwargs).tolist()}) + return pd.DataFrame({0: _encode_sentences_with_nulls(X_list, method_to_call, encode_kwargs)}) # Case 1: User provided explicit signatures if model_meta.signatures: @@ -957,7 +1029,9 @@ def fn( if truncate_dim is not None: encode_kwargs["truncate_dim"] = truncate_dim - return pd.DataFrame({output_name: method_to_call(X_list, **encode_kwargs).tolist()}) + return pd.DataFrame( + {output_name: _encode_sentences_with_nulls(X_list, method_to_call, encode_kwargs)} + ) return fn diff --git a/snowflake/ml/model/_packager/model_handlers_test/BUILD.bazel b/snowflake/ml/model/_packager/model_handlers_test/BUILD.bazel index 64619e64..a702f125 100644 --- a/snowflake/ml/model/_packager/model_handlers_test/BUILD.bazel +++ b/snowflake/ml/model/_packager/model_handlers_test/BUILD.bazel @@ -239,6 +239,18 @@ py_test( ], ) +py_test( + name = "huggingface_utils_test", + srcs = ["huggingface_utils_test.py"], + tags = [ + "feature:model_registry", + "short_regress", + ], + deps = [ + "//snowflake/ml/model/_packager/model_handlers:huggingface_utils", + ], +) + py_test( name = "sentence_transformers_test", timeout = "eternal", diff --git a/snowflake/ml/model/_packager/model_handlers_test/huggingface_pipeline_test.py b/snowflake/ml/model/_packager/model_handlers_test/huggingface_pipeline_test.py index 464ab279..2b57444f 100644 --- a/snowflake/ml/model/_packager/model_handlers_test/huggingface_pipeline_test.py +++ b/snowflake/ml/model/_packager/model_handlers_test/huggingface_pipeline_test.py @@ -404,6 +404,51 @@ def test_save_wrapper_captures_model(self) -> None: self.assertEqual(blob_options["task"], "fill-mask") self.assertEqual(blob_options["model"], expected_model_name) + @mock.patch("huggingface_hub.hf_hub_download") + @mock.patch("huggingface_hub.HfApi") + def test_save_model_lazy_upload_skips_copytree( + self, + mock_hf_api: mock.Mock, + mock_hf_hub_download: mock.Mock, + ) -> None: + """Lazy HuggingFace wrapper logging should defer weight copies and attach upload metadata.""" + mock_hf_api.return_value.model_info.return_value.siblings = [ + mock.Mock(rfilename="config.json", size=570), + mock.Mock(rfilename="model.safetensors", size=1000), + ] + + wrapper_model = hf_base.TransformersPipeline( + task="text-generation", + model="facebook/opt-125m", + compute_pool_for_log=None, + lazy_upload=True, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + packager = model_packager.ModelPackager(os.path.join(tmpdir, "model")) + with mock.patch("shutil.copytree") as mock_copytree: + packager.save( + name="model", + model=wrapper_model, + metadata={"author": "test", "version": "1"}, + options=model_types.HuggingFaceSaveOptions(), + ) + mock_copytree.assert_not_called() + + assert packager.meta is not None + lazy_hf_upload = getattr(packager.meta, "_lazy_hf_upload", None) + self.assertIsNotNone(lazy_hf_upload) + assert lazy_hf_upload is not None + self.assertEqual(lazy_hf_upload.files, ["config.json", "model.safetensors"]) + self.assertEqual( + lazy_hf_upload.file_sizes, + {"config.json": 570, "model.safetensors": 1000}, + ) + self.assertEqual( + lazy_hf_upload.download_kwargs, + {"repo_id": "facebook/opt-125m", "revision": None}, + ) + @parameterized.parameters( # type: ignore[misc] {"transformers_version": "4.41.2", "tokenizers_version": "0.19.1"}, {"transformers_version": "5.3.0", "tokenizers_version": "0.22.2"}, diff --git a/snowflake/ml/model/_packager/model_handlers_test/huggingface_utils_test.py b/snowflake/ml/model/_packager/model_handlers_test/huggingface_utils_test.py new file mode 100644 index 00000000..4e879f72 --- /dev/null +++ b/snowflake/ml/model/_packager/model_handlers_test/huggingface_utils_test.py @@ -0,0 +1,41 @@ +from unittest import mock + +from absl.testing import absltest + +from snowflake.ml.model._packager.model_handlers.huggingface import _utils as hf_utils + + +class HuggingFaceUtilsTest(absltest.TestCase): + def test_download_token_for_lazy_upload_returns_plain_token(self) -> None: + model = mock.Mock() + model.secret_identifier = None + model.token_or_secret = "hf_test_token" + + self.assertEqual( + hf_utils.download_token_for_lazy_upload(model), + "hf_test_token", + ) + + def test_download_token_for_lazy_upload_returns_none_without_token(self) -> None: + model = mock.Mock() + model.secret_identifier = None + model.token_or_secret = None + + self.assertIsNone(hf_utils.download_token_for_lazy_upload(model)) + + def test_download_token_for_lazy_upload_raises_for_secret(self) -> None: + model = mock.Mock() + model.secret_identifier = "db.schema.secret" + model.token_or_secret = "db.schema.secret" + + with self.assertRaises(ValueError) as error_context: + hf_utils.download_token_for_lazy_upload(model) + self.assertEqual( + str(error_context.exception), + "model upload: HuggingFace lazy upload cannot resolve auth from a Snowflake secret during local logging. " + "Set the HF_TOKEN environment variable or pass a HuggingFace token when constructing the model.", + ) + + +if __name__ == "__main__": + absltest.main() diff --git a/snowflake/ml/model/_packager/model_handlers_test/sentence_transformers_test.py b/snowflake/ml/model/_packager/model_handlers_test/sentence_transformers_test.py index aefe2281..59b7cb14 100644 --- a/snowflake/ml/model/_packager/model_handlers_test/sentence_transformers_test.py +++ b/snowflake/ml/model/_packager/model_handlers_test/sentence_transformers_test.py @@ -5,7 +5,7 @@ import tempfile import warnings from importlib import metadata as importlib_metadata -from typing import Any, cast +from typing import Any, Mapping, Optional, cast from unittest import mock import numpy as np @@ -23,8 +23,10 @@ SentenceTransformerHandler, _auto_infer_signature, _capture_model_truncate_dim, + _encode_sentences_with_nulls, _get_available_default_methods, _get_embedding_dim_from_config, + _is_null_sentence, _supports_encode_truncate_dim_param, _supports_init_truncate_dim, _validate_sentence_transformers_signatures, @@ -247,6 +249,83 @@ def test_validate_sentence_transformers_signatures_valid(self) -> None: } _validate_sentence_transformers_signatures(valid_all_methods) + def test_is_null_sentence(self) -> None: + self.assertTrue(_is_null_sentence(None)) + self.assertTrue(_is_null_sentence(pd.NA)) + self.assertTrue(_is_null_sentence(float("nan"))) + self.assertFalse(_is_null_sentence("hello")) + self.assertFalse(_is_null_sentence("")) + self.assertFalse(_is_null_sentence(["hello"])) + + def test_encode_sentences_with_nulls(self) -> None: + def fake_encode(sentences: list[str], **kwargs: Any) -> npt.NDArray[Any]: + del kwargs + return np.array([[float(index), float(index + 1)] for index, _ in enumerate(sentences)]) + + outputs = _encode_sentences_with_nulls( + ["a", None, pd.NA, "b"], + fake_encode, + {}, + ) + self.assertEqual(outputs[0], [0.0, 1.0]) + self.assertIsNone(outputs[1]) + self.assertIsNone(outputs[2]) + self.assertEqual(outputs[3], [1.0, 2.0]) + + all_valid_outputs = _encode_sentences_with_nulls(["a", "b"], fake_encode, {}) + self.assertEqual(all_valid_outputs, [[0.0, 1.0], [1.0, 2.0]]) + + all_null_outputs = _encode_sentences_with_nulls([None, pd.NA], fake_encode, {}) + self.assertEqual(all_null_outputs, [None, None]) + + def test_null_sentence_custom_model_returns_none_embedding(self) -> None: + model = self._st_model + sentences = pd.DataFrame( + { + "sentence": [ + "Why don't scientists trust atoms? Because they make up everything.", + None, + "Parallel lines have so much in common. It's a shame they'll never meet.", + ] + } + ) + expected_embeddings = model.encode( + [sentences.iloc[0, 0], sentences.iloc[2, 0]], + ).tolist() + expected = pd.DataFrame( + { + "output": [ + expected_embeddings[0], + None, + expected_embeddings[1], + ] + } + ) + + sig = {"encode": model_signature.infer_signature(sentences.iloc[[0, 2]], expected.iloc[[0, 2]])} + + with tempfile.TemporaryDirectory() as tmpdir: + model_packager.ModelPackager(os.path.join(tmpdir, "model")).save( + name="model", + model=model, + signatures=sig, + metadata={"author": "test", "version": "1"}, + options=model_types.SentenceTransformersSaveOptions(), + ) + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + pk = model_packager.ModelPackager(os.path.join(tmpdir, "model")) + pk.load(as_custom_model=True) + assert pk.model + predict_method = getattr(pk.model, "encode", None) + assert callable(predict_method) + embeddings_load = predict_method(sentences) + embeddings_load.columns = expected.columns + self.assertIsNone(embeddings_load.iloc[1, 0]) + assert_frame_equal(embeddings_load.iloc[[0, 2]], expected.iloc[[0, 2]]) + assert_frame_equal(embeddings_load.iloc[[1]], expected.iloc[[1]]) + def test_validate_sentence_transformers_signatures_empty(self) -> None: """Test that empty signatures raise ValueError.""" with self.assertRaises(ValueError) as ctx: @@ -1572,6 +1651,240 @@ def test_get_embedding_dim_from_config_missing_pooling_config(self) -> None: dim = _get_embedding_dim_from_config(snapshot_dir) self.assertIsNone(dim) + def _mock_lazy_st_repo_download( + self, + *, + local_dir: str, + filename: str, + modules: list[dict[str, object]], + pooling_config: Optional[Mapping[str, Any]] = None, + dense_config: Optional[Mapping[str, Any]] = None, + ) -> str: + if filename == "modules.json": + with open(os.path.join(local_dir, "modules.json"), "w") as f: + json.dump(modules, f) + return os.path.join(local_dir, "modules.json") + if filename == "1_Pooling/config.json" and pooling_config is not None: + pooling_dir = os.path.join(local_dir, "1_Pooling") + os.makedirs(pooling_dir, exist_ok=True) + config_path = os.path.join(pooling_dir, "config.json") + with open(config_path, "w") as f: + json.dump(pooling_config, f) + return config_path + if filename == "2_Dense/config.json" and dense_config is not None: + dense_dir = os.path.join(local_dir, "2_Dense") + os.makedirs(dense_dir, exist_ok=True) + config_path = os.path.join(dense_dir, "config.json") + with open(config_path, "w") as f: + json.dump(dense_config, f) + return config_path + raise AssertionError(f"Unexpected download: {filename}") + + @mock.patch("huggingface_hub.hf_hub_download") + @mock.patch("huggingface_hub.HfApi") + def test_save_wrapper_lazy_upload_skips_copytree_and_sets_lazy_hf_upload( + self, + mock_hf_api: mock.Mock, + mock_hf_hub_download: mock.Mock, + ) -> None: + """Lazy ST wrapper logging should defer weight copies and attach upload metadata.""" + from snowflake.ml.model.models import huggingface as snowml_huggingface + + modules = [ + {"idx": 0, "path": "", "type": "sentence_transformers.models.Transformer"}, + {"idx": 1, "path": "1_Pooling", "type": "sentence_transformers.models.Pooling"}, + ] + pooling_config = {"word_embedding_dimension": 384, "pooling_mode_mean_tokens": True} + + def fake_download( + *, + repo_id: str, + filename: str, + revision: object, + token: object, + local_dir: str, + **kwargs: object, + ) -> str: + del repo_id, revision, token, kwargs + return self._mock_lazy_st_repo_download( + local_dir=local_dir, + filename=filename, + modules=modules, + pooling_config=pooling_config, + ) + + mock_hf_hub_download.side_effect = fake_download + mock_hf_api.return_value.model_info.return_value.siblings = [ + mock.Mock(rfilename="modules.json", size=120), + mock.Mock(rfilename="1_Pooling/config.json", size=200), + mock.Mock(rfilename="model.safetensors", size=1000), + ] + + wrapper = snowml_huggingface.SentenceTransformer( + model="sentence-transformers/all-MiniLM-L6-v2", + compute_pool_for_log=None, + lazy_upload=True, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + packager = model_packager.ModelPackager(os.path.join(tmpdir, "model")) + with mock.patch("shutil.copytree") as mock_copytree: + packager.save( + name="model", + model=wrapper, + metadata={"author": "test", "version": "1"}, + options=model_types.SentenceTransformersSaveOptions(), + ) + mock_copytree.assert_not_called() + + assert packager.meta is not None + lazy_hf_upload = getattr(packager.meta, "_lazy_hf_upload", None) + self.assertIsNotNone(lazy_hf_upload) + assert lazy_hf_upload is not None + self.assertEqual( + lazy_hf_upload.files, + ["modules.json", "1_Pooling/config.json", "model.safetensors"], + ) + self.assertEqual( + lazy_hf_upload.file_sizes, + {"modules.json": 120, "1_Pooling/config.json": 200, "model.safetensors": 1000}, + ) + blob_options = cast( + model_meta_schema.SentenceTransformersModelBlobOptions, + packager.meta.models["model"].options, + ) + self.assertTrue(blob_options.get("is_repo_downloaded", False)) + + @mock.patch("huggingface_hub.hf_hub_download") + @mock.patch("huggingface_hub.HfApi") + def test_save_wrapper_lazy_upload_auto_infers_signature( + self, + mock_hf_api: mock.Mock, + mock_hf_hub_download: mock.Mock, + ) -> None: + """Lazy ST wrapper save should auto-infer signatures after config download.""" + from snowflake.ml.model.models import huggingface as snowml_huggingface + + modules = [ + {"idx": 0, "path": "", "type": "sentence_transformers.models.Transformer"}, + {"idx": 1, "path": "1_Pooling", "type": "sentence_transformers.models.Pooling"}, + ] + pooling_config = {"word_embedding_dimension": 384, "pooling_mode_mean_tokens": True} + + def fake_download( + *, + repo_id: str, + filename: str, + revision: object, + token: object, + local_dir: str, + **kwargs: object, + ) -> str: + del repo_id, revision, token, kwargs + return self._mock_lazy_st_repo_download( + local_dir=local_dir, + filename=filename, + modules=modules, + pooling_config=pooling_config, + ) + + mock_hf_hub_download.side_effect = fake_download + mock_hf_api.return_value.model_info.return_value.siblings = [ + mock.Mock(rfilename="modules.json", size=120), + mock.Mock(rfilename="1_Pooling/config.json", size=200), + mock.Mock(rfilename="model.safetensors", size=1000), + ] + + wrapper = snowml_huggingface.SentenceTransformer( + model="sentence-transformers/all-MiniLM-L6-v2", + compute_pool_for_log=None, + lazy_upload=True, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + model_packager.ModelPackager(os.path.join(tmpdir, "model")).save( + name="model", + model=wrapper, + metadata={"author": "test", "version": "1"}, + options=model_types.SentenceTransformersSaveOptions(), + ) + + pk = model_packager.ModelPackager(os.path.join(tmpdir, "model")) + pk.load(meta_only=True) + assert pk.meta is not None + + for method_name in _DEFAULT_WRAPPER_TARGET_METHODS: + self.assertIn(method_name, pk.meta.signatures) + sig = pk.meta.signatures[method_name] + assert isinstance(sig.outputs[0], model_signature.FeatureSpec) + self.assertEqual(sig.outputs[0]._shape, (384,)) + + @mock.patch("huggingface_hub.hf_hub_download") + @mock.patch("huggingface_hub.HfApi") + def test_save_wrapper_lazy_upload_dense_module_dim( + self, + mock_hf_api: mock.Mock, + mock_hf_hub_download: mock.Mock, + ) -> None: + """Lazy ST wrapper save should prefer Dense out_features when present.""" + from snowflake.ml.model.models import huggingface as snowml_huggingface + + modules = [ + {"idx": 0, "path": "", "type": "sentence_transformers.models.Transformer"}, + {"idx": 1, "path": "1_Pooling", "type": "sentence_transformers.models.Pooling"}, + {"idx": 2, "path": "2_Dense", "type": "sentence_transformers.models.Dense"}, + ] + pooling_config = {"word_embedding_dimension": 768, "pooling_mode_mean_tokens": True} + dense_config = {"out_features": 256} + + def fake_download( + *, + repo_id: str, + filename: str, + revision: object, + token: object, + local_dir: str, + **kwargs: object, + ) -> str: + del repo_id, revision, token, kwargs + return self._mock_lazy_st_repo_download( + local_dir=local_dir, + filename=filename, + modules=modules, + pooling_config=pooling_config, + dense_config=dense_config, + ) + + mock_hf_hub_download.side_effect = fake_download + mock_hf_api.return_value.model_info.return_value.siblings = [ + mock.Mock(rfilename="modules.json", size=120), + mock.Mock(rfilename="1_Pooling/config.json", size=200), + mock.Mock(rfilename="2_Dense/config.json", size=150), + mock.Mock(rfilename="model.safetensors", size=1000), + ] + + wrapper = snowml_huggingface.SentenceTransformer( + model="sentence-transformers/all-MiniLM-L6-v2", + compute_pool_for_log=None, + lazy_upload=True, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + model_packager.ModelPackager(os.path.join(tmpdir, "model")).save( + name="model", + model=wrapper, + metadata={"author": "test", "version": "1"}, + options=model_types.SentenceTransformersSaveOptions(target_methods=["encode"]), + ) + + pk = model_packager.ModelPackager(os.path.join(tmpdir, "model")) + pk.load(meta_only=True) + assert pk.meta is not None + + sig = pk.meta.signatures["encode"] + assert isinstance(sig.outputs[0], model_signature.FeatureSpec) + self.assertEqual(sig.outputs[0]._shape, (256,)) + def test_can_handle_wrapper(self) -> None: from snowflake.ml.model.models import huggingface as snowml_huggingface @@ -1602,6 +1915,7 @@ def test_save_and_load_wrapper_with_snapshot(self) -> None: wrapper = snowml_huggingface.SentenceTransformer( model=MODEL_NAMES[0], compute_pool_for_log=None, + lazy_upload=False, ) wrapper.repo_snapshot_dir = self._wrapper_snapshot_dir @@ -1636,6 +1950,7 @@ def test_save_wrapper_sets_is_repo_downloaded(self) -> None: wrapper = snowml_huggingface.SentenceTransformer( model=MODEL_NAMES[0], compute_pool_for_log=None, + lazy_upload=False, ) wrapper.repo_snapshot_dir = self._wrapper_snapshot_dir @@ -1700,6 +2015,7 @@ def test_save_wrapper_auto_infers_signature_from_config(self) -> None: wrapper = snowml_huggingface.SentenceTransformer( model=MODEL_NAMES[0], compute_pool_for_log=None, + lazy_upload=False, ) wrapper.repo_snapshot_dir = self._wrapper_snapshot_dir @@ -1732,6 +2048,7 @@ def test_save_wrapper_target_methods_subset(self) -> None: wrapper = snowml_huggingface.SentenceTransformer( model=MODEL_NAMES[0], compute_pool_for_log=None, + lazy_upload=False, ) wrapper.repo_snapshot_dir = self._wrapper_snapshot_dir @@ -1759,6 +2076,7 @@ def test_save_wrapper_target_methods_invalid(self) -> None: wrapper = snowml_huggingface.SentenceTransformer( model=MODEL_NAMES[0], compute_pool_for_log=None, + lazy_upload=False, ) wrapper.repo_snapshot_dir = self._wrapper_snapshot_dir @@ -1784,6 +2102,7 @@ def test_save_wrapper_dependency_unpinned(self) -> None: wrapper = snowml_huggingface.SentenceTransformer( model=MODEL_NAMES[0], compute_pool_for_log=None, + lazy_upload=False, ) wrapper.repo_snapshot_dir = self._wrapper_snapshot_dir @@ -1816,6 +2135,7 @@ def test_save_wrapper_multi_method_inference(self) -> None: wrapper = snowml_huggingface.SentenceTransformer( model=MODEL_NAMES[0], compute_pool_for_log=None, + lazy_upload=False, ) wrapper.repo_snapshot_dir = self._wrapper_snapshot_dir diff --git a/snowflake/ml/model/_packager/model_meta/BUILD.bazel b/snowflake/ml/model/_packager/model_meta/BUILD.bazel index 9650811b..badab7e7 100644 --- a/snowflake/ml/model/_packager/model_meta/BUILD.bazel +++ b/snowflake/ml/model/_packager/model_meta/BUILD.bazel @@ -30,6 +30,7 @@ py_library( "//snowflake/ml/model:code_path", "//snowflake/ml/model:model_signature", "//snowflake/ml/model:type_hints", + "//snowflake/ml/model/_model_composer:huggingface_lazy_uploader", "//snowflake/ml/model/_packager/model_env", "//snowflake/ml/model/_packager/model_meta_migrator:migrator_plans", "//snowflake/ml/model/_packager/model_runtime", diff --git a/snowflake/ml/model/_packager/model_meta/model_meta.py b/snowflake/ml/model/_packager/model_meta/model_meta.py index 1fbf86e2..65a24b7f 100644 --- a/snowflake/ml/model/_packager/model_meta/model_meta.py +++ b/snowflake/ml/model/_packager/model_meta/model_meta.py @@ -16,6 +16,7 @@ from snowflake.ml import version as snowml_version from snowflake.ml._internal import env_utils, file_utils from snowflake.ml.model import model_signature, type_hints as model_types +from snowflake.ml.model._model_composer import huggingface_lazy_uploader from snowflake.ml.model._packager.model_env import model_env from snowflake.ml.model._packager.model_meta import model_blob_meta, model_meta_schema from snowflake.ml.model._packager.model_meta_migrator import migrator_plans @@ -348,6 +349,8 @@ def __init__( self.sample_input_file_paths: dict[str, str] = sample_input_file_paths or {} self._packaged_env_dict = packaged_env_dict + self._lazy_hf_upload: Optional[huggingface_lazy_uploader.LazyHFUpload] = None + @property def min_snowpark_ml_version(self) -> str: return self._min_snowpark_ml_version.base_version diff --git a/snowflake/ml/model/models/huggingface.py b/snowflake/ml/model/models/huggingface.py index 71f95aa6..628b51af 100644 --- a/snowflake/ml/model/models/huggingface.py +++ b/snowflake/ml/model/models/huggingface.py @@ -1,7 +1,10 @@ import enum +import fnmatch import json import logging import os +import pathlib +import tempfile import warnings from typing import Any, Optional, Union @@ -29,6 +32,7 @@ class _LoggingMode(enum.Enum): class TransformersPipeline: _requires_task: bool = True + _CHAT_TEMPLATE_METADATA_FILES = ("tokenizer_config.json", "chat_template.jinja") def __init__( self, @@ -43,6 +47,7 @@ def __init__( # repo snapshot download args allow_patterns: Optional[Union[list[str], str]] = None, ignore_patterns: Optional[Union[list[str], str]] = None, + lazy_upload: bool = True, **kwargs: Any, ) -> None: """ @@ -77,6 +82,9 @@ def __init__( downloaded from the HuggingFace repository. allow_patterns: If provided, only files matching at least one pattern are downloaded. ignore_patterns: If provided, files matching any of the patterns are not downloaded. + lazy_upload: When ``compute_pool_for_log`` is None, list HuggingFace repository files at construction time + and stream each file to Snowflake during ``log_model`` instead of downloading the full snapshot locally. + Defaults to True. Set to False to download the entire repository before logging. kwargs: Additional keyword arguments passed along to the specific pipeline init (see the documentation for the corresponding pipeline class for possible values). @@ -96,6 +104,10 @@ def __init__( self._validate_device_map(kwargs, model_kwargs) + self._lazy_repo_files: Optional[list[str]] = None + self._lazy_file_sizes: Optional[dict[str, int]] = None + self._lazy_download_kwargs: Optional[dict[str, Any]] = None + repo_snapshot_dir = self._download_snapshot_if_needed( logging_mode=logging_mode, uses_secret=uses_secret, @@ -104,6 +116,8 @@ def __init__( token_or_secret=token_or_secret, allow_patterns=allow_patterns, ignore_patterns=ignore_patterns, + lazy_upload=lazy_upload, + task=task, ) self.has_chat_template = self._has_chat_template( @@ -260,7 +274,29 @@ def _validate_device_map(kwargs: dict[str, Any], model_kwargs: dict[str, Any]) - ) @staticmethod + def _filter_repo_files( + repo_files: list[str], + *, + allow_patterns: Optional[Union[list[str], str]] = None, + ignore_patterns: Optional[Union[list[str], str]] = None, + ) -> list[str]: + """Filter repo file paths using the same glob semantics as ``snapshot_download``.""" + filtered_files = repo_files + if allow_patterns is not None: + allow = [allow_patterns] if isinstance(allow_patterns, str) else allow_patterns + filtered_files = [ + filename for filename in filtered_files if any(fnmatch.fnmatch(filename, p) for p in allow) + ] + if ignore_patterns is not None: + ignore = [ignore_patterns] if isinstance(ignore_patterns, str) else ignore_patterns + filtered_files = [ + filename for filename in filtered_files if not any(fnmatch.fnmatch(filename, p) for p in ignore) + ] + return filtered_files + def _download_snapshot_if_needed( + self, + *, logging_mode: _LoggingMode, uses_secret: bool, model: Optional[str], @@ -268,6 +304,8 @@ def _download_snapshot_if_needed( token_or_secret: Optional[str], allow_patterns: Optional[Union[list[str], str]], ignore_patterns: Optional[Union[list[str], str]], + lazy_upload: bool, + task: Optional[str], ) -> Optional[str]: """Download the model snapshot if in SNAPSHOT_DOWNLOAD mode. @@ -279,6 +317,8 @@ def _download_snapshot_if_needed( token_or_secret: The auth token or secret. allow_patterns: Patterns of files to include. ignore_patterns: Patterns of files to exclude. + lazy_upload: Whether to list repo files and defer weight downloads until log time. + task: The pipeline task. Returns: The path to the downloaded snapshot directory, or None if not downloaded. @@ -295,17 +335,161 @@ def _download_snapshot_if_needed( try: import huggingface_hub as hf_hub + except ImportError: + logger.info("huggingface_hub package is not installed, skipping snapshot download") + return None - return hf_hub.snapshot_download( - repo_id=model, + if lazy_upload: + return self._prepare_lazy_repo_upload( + hf_hub=hf_hub, + model=model, revision=revision, - token=token_or_secret, + token_or_secret=token_or_secret, allow_patterns=allow_patterns, ignore_patterns=ignore_patterns, + task=task, ) - except ImportError: - logger.info("huggingface_hub package is not installed, skipping snapshot download") - return None + + return hf_hub.snapshot_download( + repo_id=model, + revision=revision, + token=token_or_secret, + allow_patterns=allow_patterns, + ignore_patterns=ignore_patterns, + ) + + def _lazy_metadata_seed_files( + self, + *, + task: Optional[str], + filtered_files: list[str], + ) -> list[str]: + """Return repo-relative paths to download before follow-up metadata resolution.""" + if task in self._CHAT_TEMPLATE_TASKS: + return [filename for filename in self._CHAT_TEMPLATE_METADATA_FILES if filename in filtered_files] + return [] + + def _lazy_metadata_followup_files( + self, + *, + metadata_dir: str, + filtered_files: list[str], + ) -> list[str]: + """Return additional repo-relative metadata paths after seed files are downloaded.""" + return [] + + @staticmethod + def _download_lazy_metadata_file( + *, + hf_hub: Any, + model: str, + filename: str, + revision: Optional[str], + token_or_secret: Optional[str], + metadata_dir: str, + ) -> None: + """Download one HuggingFace metadata file into the lazy-upload staging directory.""" + try: + hf_hub.hf_hub_download( + repo_id=model, + filename=filename, + revision=revision, + token=token_or_secret, + local_dir=metadata_dir, + ) + except hf_hub.errors.EntryNotFoundError: + logger.error( + "HuggingFace metadata file %s was listed but not found during download.", + filename, + ) + raise + + def _download_lazy_metadata_files( + self, + *, + hf_hub: Any, + model: str, + revision: Optional[str], + token_or_secret: Optional[str], + metadata_dir: str, + filtered_files: list[str], + task: Optional[str], + ) -> None: + """Download metadata files needed for local inference before log_model.""" + seed_files = self._lazy_metadata_seed_files(task=task, filtered_files=filtered_files) + for filename in seed_files: + self._download_lazy_metadata_file( + hf_hub=hf_hub, + model=model, + filename=filename, + revision=revision, + token_or_secret=token_or_secret, + metadata_dir=metadata_dir, + ) + + followup_files = self._lazy_metadata_followup_files( + metadata_dir=metadata_dir, + filtered_files=filtered_files, + ) + for filename in followup_files: + self._download_lazy_metadata_file( + hf_hub=hf_hub, + model=model, + filename=filename, + revision=revision, + token_or_secret=token_or_secret, + metadata_dir=metadata_dir, + ) + + def _prepare_lazy_repo_upload( + self, + *, + hf_hub: Any, + model: str, + revision: Optional[str], + token_or_secret: Optional[str], + allow_patterns: Optional[Union[list[str], str]], + ignore_patterns: Optional[Union[list[str], str]], + task: Optional[str], + ) -> str: + """List repo files and download only metadata needed before log_model.""" + api = hf_hub.HfApi() + repo_info = api.model_info( + repo_id=model, + revision=revision, + token=token_or_secret, + files_metadata=True, + ) + repo_files = [sibling.rfilename for sibling in repo_info.siblings] + size_by_file = {sibling.rfilename: sibling.size for sibling in repo_info.siblings if sibling.size is not None} + filtered_files = self._filter_repo_files( + repo_files, + allow_patterns=allow_patterns, + ignore_patterns=ignore_patterns, + ) + filtered_file_sizes = {filename: size_by_file.get(filename, 0) for filename in filtered_files} + for filename in filtered_files: + if filename not in size_by_file: + logger.warning("HuggingFace file %s has no known size; disk checks may be incomplete.", filename) + + metadata_dir = tempfile.mkdtemp(prefix="snowml_hf_lazy_") + self._download_lazy_metadata_files( + hf_hub=hf_hub, + model=model, + revision=revision, + token_or_secret=token_or_secret, + metadata_dir=metadata_dir, + filtered_files=filtered_files, + task=task, + ) + + self._lazy_repo_files = filtered_files + self._lazy_file_sizes = filtered_file_sizes + self._lazy_download_kwargs = { + "repo_id": model, + "revision": revision, + } + return metadata_dir _CHAT_TEMPLATE_TASKS = frozenset( { @@ -393,8 +577,10 @@ class SentenceTransformer(TransformersPipeline): logged via ``SYSTEM$IMPORT_MODEL`` in a SPCS job. **Local logging**: when ``compute_pool_for_log`` is ``None``, artifacts are - downloaded locally via ``huggingface_hub.snapshot_download`` before registry - upload. Signatures are inferred from snapshot config files (no local + streamed to Snowflake during ``log_model`` by default (``lazy_upload=True``), + or downloaded locally via ``huggingface_hub.snapshot_download`` when + ``lazy_upload=False``. Signatures are inferred from metadata or snapshot + config files (no local ``sentence_transformers`` import). The model environment records ``sentence-transformers`` without a version pin so serve-time resolves the latest compatible release. By default, these inference methods are registered: @@ -434,6 +620,68 @@ class SentenceTransformer(TransformersPipeline): """ _requires_task = False + _SIGNATURE_METADATA_SEED_FILES = ("modules.json",) + + @staticmethod + def _module_config_relative_paths(modules: list[dict[str, Any]]) -> list[str]: + """Return HuggingFace repo-relative paths to each module's config.json.""" + config_paths: list[str] = [] + for module_entry in modules: + module_path = module_entry.get("path", "") + if module_path: + config_paths.append(f"{module_path}/config.json") + else: + config_paths.append("config.json") + return config_paths + + def _lazy_metadata_seed_files( + self, + *, + task: Optional[str], + filtered_files: list[str], + ) -> list[str]: + del task + return [filename for filename in self._SIGNATURE_METADATA_SEED_FILES if filename in filtered_files] + + def _lazy_metadata_followup_files( + self, + *, + metadata_dir: str, + filtered_files: list[str], + ) -> list[str]: + """Download module config files needed for embedding-dimension inference.""" + modules_json_path = os.path.join(metadata_dir, "modules.json") + if not os.path.isfile(modules_json_path): + return [] + + try: + with open(modules_json_path, encoding="utf-8") as f: + modules = json.load(f) + except (json.JSONDecodeError, OSError): + return [] + + if not isinstance(modules, list): + return [] + + followup_files: list[str] = [] + for config_rel_path in self._module_config_relative_paths(modules): + if config_rel_path not in filtered_files: + continue + module_path = os.path.dirname(config_rel_path) + if module_path and (".." in pathlib.PurePosixPath(module_path).parts or module_path.startswith("/")): + logger.warning( + "Skipping HuggingFace module config with an invalid path: %s", + config_rel_path, + ) + continue + if module_path: + local_config_path = os.path.join(metadata_dir, module_path, "config.json") + else: + local_config_path = os.path.join(metadata_dir, "config.json") + if os.path.isfile(local_config_path): + continue + followup_files.append(config_rel_path) + return followup_files def __init__( self, @@ -445,6 +693,7 @@ def __init__( compute_pool_for_log: Optional[str] = DEFAULT_CPU_COMPUTE_POOL, allow_patterns: Optional[Union[list[str], str]] = None, ignore_patterns: Optional[Union[list[str], str]] = None, + lazy_upload: bool = True, ) -> None: """Initialize a SentenceTransformer wrapper. @@ -461,6 +710,10 @@ def __init__( passed to ``log_model``. allow_patterns: File patterns to include when downloading. Defaults to None. ignore_patterns: File patterns to exclude when downloading. Defaults to None. + lazy_upload: When ``compute_pool_for_log`` is None, list HuggingFace repository + files at construction time and stream each file to Snowflake during + ``log_model`` instead of downloading the full snapshot locally. + Defaults to True. Set to False to download the entire repository before logging. """ super().__init__( task=None, @@ -472,4 +725,5 @@ def __init__( compute_pool_for_log=compute_pool_for_log, allow_patterns=allow_patterns, ignore_patterns=ignore_patterns, + lazy_upload=lazy_upload, ) diff --git a/snowflake/ml/model/models/huggingface_test.py b/snowflake/ml/model/models/huggingface_test.py index ff4a5b6f..056b1493 100644 --- a/snowflake/ml/model/models/huggingface_test.py +++ b/snowflake/ml/model/models/huggingface_test.py @@ -100,6 +100,115 @@ def test_wrapper(self) -> None: device=0, ) + def test_lazy_upload_lists_repo_files_without_snapshot_download(self) -> None: + """Lazy upload lists repo files and avoids downloading the full snapshot.""" + repo_files = ["config.json", "model.safetensors", "tokenizer_config.json"] + repo_file_sizes = { + "config.json": 570, + "model.safetensors": 1000, + "tokenizer_config.json": 48, + } + + with absltest.mock.patch("huggingface_hub.snapshot_download") as mock_snapshot_download, absltest.mock.patch( + "huggingface_hub.hf_hub_download" + ) as mock_hf_hub_download, absltest.mock.patch("huggingface_hub.HfApi") as mock_hf_api: + mock_hf_api.return_value.model_info.return_value.siblings = [ + absltest.mock.Mock(rfilename=filename, size=repo_file_sizes[filename]) for filename in repo_files + ] + + pipeline = huggingface.TransformersPipeline( + task="text-generation", + model="facebook/opt-125m", + compute_pool_for_log=None, + lazy_upload=True, + ) + + mock_snapshot_download.assert_not_called() + mock_hf_api.return_value.model_info.assert_called_once_with( + repo_id="facebook/opt-125m", + revision=None, + token=None, + files_metadata=True, + ) + self.assertEqual(pipeline._lazy_repo_files, repo_files) + self.assertEqual(pipeline._lazy_file_sizes, repo_file_sizes) + self.assertEqual( + pipeline._lazy_download_kwargs, + {"repo_id": "facebook/opt-125m", "revision": None}, + ) + downloaded_filenames = {call.kwargs["filename"] for call in mock_hf_hub_download.call_args_list} + self.assertEqual(downloaded_filenames, {"tokenizer_config.json"}) + + def test_lazy_upload_skips_chat_template_metadata_for_non_chat_tasks(self) -> None: + """Lazy upload should not download chat-template metadata for unrelated tasks.""" + repo_files = ["config.json", "model.safetensors", "tokenizer_config.json"] + + with absltest.mock.patch("huggingface_hub.snapshot_download") as mock_snapshot_download, absltest.mock.patch( + "huggingface_hub.hf_hub_download" + ) as mock_hf_hub_download, absltest.mock.patch("huggingface_hub.HfApi") as mock_hf_api: + mock_hf_api.return_value.model_info.return_value.siblings = [ + absltest.mock.Mock(rfilename=filename, size=100) for filename in repo_files + ] + + huggingface.TransformersPipeline( + task="text-classification", + model="distilbert-base-uncased-finetuned-sst-2-english", + compute_pool_for_log=None, + lazy_upload=True, + ) + + mock_snapshot_download.assert_not_called() + mock_hf_hub_download.assert_not_called() + + def test_lazy_upload_detects_chat_template_from_metadata_download(self) -> None: + """Lazy upload still detects chat templates from metadata files.""" + with tempfile.TemporaryDirectory() as repo_dir: + config_path = os.path.join(repo_dir, "tokenizer_config.json") + with open(config_path, "w") as f: + json.dump({"chat_template": "{% for msg in messages %}{{ msg }}{% endfor %}"}, f) + + with absltest.mock.patch( + "huggingface_hub.snapshot_download" + ) as mock_snapshot_download, absltest.mock.patch( + "huggingface_hub.hf_hub_download", + return_value=config_path, + ), absltest.mock.patch( + "huggingface_hub.HfApi" + ) as mock_hf_api, absltest.mock.patch( + "snowflake.ml.model.models.huggingface.tempfile.mkdtemp", + return_value=repo_dir, + ): + mock_hf_api.return_value.model_info.return_value.siblings = [ + absltest.mock.Mock(rfilename="config.json", size=100), + absltest.mock.Mock(rfilename="tokenizer_config.json", size=200), + ] + + pipeline = huggingface.TransformersPipeline( + task="text-generation", + model="some-model/with-chat-template", + compute_pool_for_log=None, + lazy_upload=True, + ) + + mock_snapshot_download.assert_not_called() + self.assertTrue(pipeline.has_chat_template) + + def test_lazy_upload_false_uses_snapshot_download(self) -> None: + """Opting out of lazy upload restores eager snapshot download.""" + with absltest.mock.patch( + "huggingface_hub.snapshot_download", return_value="/some/path" + ) as mock_snapshot_download: + pipeline = huggingface.TransformersPipeline( + task="text-generation", + model="facebook/opt-125m", + compute_pool_for_log=None, + lazy_upload=False, + ) + + mock_snapshot_download.assert_called_once() + self.assertIsNone(pipeline._lazy_repo_files) + self.assertEqual(pipeline.repo_snapshot_dir, "/some/path") + def test_remote_logging_skips_snapshot_download(self) -> None: """Test that when compute_pool_for_log is set, we don't download from huggingface_hub.""" with absltest.mock.patch("huggingface_hub.snapshot_download") as mock_snapshot_download: @@ -195,6 +304,7 @@ def test_detect_chat_template_set_for_text_generation_snapshot_download(self) -> task="text-generation", model="some-model/with-chat-template", compute_pool_for_log=None, + lazy_upload=False, ) self.assertTrue(pipeline.has_chat_template) @@ -205,6 +315,7 @@ def test_detect_chat_template_false_for_non_text_generation_task(self) -> None: task="fill-mask", model="some-model/fill-mask", compute_pool_for_log=None, + lazy_upload=False, ) self.assertFalse(pipeline.has_chat_template) @@ -229,6 +340,7 @@ def test_detect_chat_template_for_image_text_to_text_task(self) -> None: task="image-text-to-text", model="some-model/vision-llm", compute_pool_for_log=None, + lazy_upload=False, ) self.assertTrue(pipeline.has_chat_template) @@ -244,6 +356,7 @@ def test_detect_chat_template_for_video_text_to_text_task(self) -> None: task="video-text-to-text", model="some-model/video-llm", compute_pool_for_log=None, + lazy_upload=False, ) self.assertTrue(pipeline.has_chat_template) @@ -259,6 +372,7 @@ def test_detect_chat_template_for_audio_text_to_text_task(self) -> None: task="audio-text-to-text", model="some-model/audio-llm", compute_pool_for_log=None, + lazy_upload=False, ) self.assertTrue(pipeline.has_chat_template) @@ -274,6 +388,7 @@ def test_detect_chat_template_false_for_multimodal_without_template(self) -> Non task="image-text-to-text", model="some-model/vision-no-template", compute_pool_for_log=None, + lazy_upload=False, ) self.assertFalse(pipeline.has_chat_template) diff --git a/snowflake/ml/model/models/sentence_transformers_test.py b/snowflake/ml/model/models/sentence_transformers_test.py index c80a4f15..1e087b43 100644 --- a/snowflake/ml/model/models/sentence_transformers_test.py +++ b/snowflake/ml/model/models/sentence_transformers_test.py @@ -1,3 +1,4 @@ +import json import os import tempfile @@ -46,10 +47,92 @@ def test_local_mode_downloads_snapshot(self) -> None: model = huggingface.SentenceTransformer( model="sentence-transformers/all-MiniLM-L6-v2", compute_pool_for_log=None, + lazy_upload=False, ) mock_snapshot.assert_called_once() self.assertEqual(model.repo_snapshot_dir, fake_snapshot_dir) + def test_lazy_upload_lists_repo_files_without_snapshot_download(self) -> None: + """Lazy upload lists repo files and avoids downloading the full snapshot.""" + repo_files = ["modules.json", "1_Pooling/config.json", "model.safetensors"] + repo_file_sizes = { + "modules.json": 120, + "1_Pooling/config.json": 200, + "model.safetensors": 1000, + } + modules = [ + {"idx": 0, "path": "", "type": "sentence_transformers.models.Transformer"}, + {"idx": 1, "path": "1_Pooling", "type": "sentence_transformers.models.Pooling"}, + ] + + def fake_download( + *, + repo_id: str, + filename: str, + revision: object, + token: object, + local_dir: str, + **kwargs: object, + ) -> str: + del repo_id, revision, token, kwargs + if filename == "modules.json": + with open(os.path.join(local_dir, "modules.json"), "w") as f: + json.dump(modules, f) + return os.path.join(local_dir, "modules.json") + if filename == "1_Pooling/config.json": + pooling_dir = os.path.join(local_dir, "1_Pooling") + os.makedirs(pooling_dir, exist_ok=True) + config_path = os.path.join(pooling_dir, "config.json") + with open(config_path, "w") as f: + json.dump({"word_embedding_dimension": 384}, f) + return config_path + raise AssertionError(f"Unexpected download: {filename}") + + with absltest.mock.patch("huggingface_hub.snapshot_download") as mock_snapshot_download, absltest.mock.patch( + "huggingface_hub.hf_hub_download", + side_effect=fake_download, + ) as mock_hf_hub_download, absltest.mock.patch("huggingface_hub.HfApi") as mock_hf_api: + mock_hf_api.return_value.model_info.return_value.siblings = [ + absltest.mock.Mock(rfilename=filename, size=repo_file_sizes[filename]) for filename in repo_files + ] + + model = huggingface.SentenceTransformer( + model="sentence-transformers/all-MiniLM-L6-v2", + compute_pool_for_log=None, + lazy_upload=True, + ) + + mock_snapshot_download.assert_not_called() + mock_hf_api.return_value.model_info.assert_called_once_with( + repo_id="sentence-transformers/all-MiniLM-L6-v2", + revision=None, + token=None, + files_metadata=True, + ) + self.assertEqual(model._lazy_repo_files, repo_files) + self.assertEqual(model._lazy_file_sizes, repo_file_sizes) + self.assertEqual( + model._lazy_download_kwargs, + {"repo_id": "sentence-transformers/all-MiniLM-L6-v2", "revision": None}, + ) + downloaded_filenames = {call.kwargs["filename"] for call in mock_hf_hub_download.call_args_list} + self.assertEqual(downloaded_filenames, {"modules.json", "1_Pooling/config.json"}) + + def test_lazy_upload_false_uses_snapshot_download(self) -> None: + """Opting out of lazy upload restores eager snapshot download.""" + with absltest.mock.patch( + "huggingface_hub.snapshot_download", return_value="/some/path" + ) as mock_snapshot_download: + model = huggingface.SentenceTransformer( + model="sentence-transformers/all-MiniLM-L6-v2", + compute_pool_for_log=None, + lazy_upload=False, + ) + + mock_snapshot_download.assert_called_once() + self.assertIsNone(model._lazy_repo_files) + self.assertEqual(model.repo_snapshot_dir, "/some/path") + def test_requires_task_is_false(self) -> None: self.assertFalse(huggingface.SentenceTransformer._requires_task) diff --git a/snowflake/ml/modeling/_internal/snowpark_implementations/snowpark_handlers.py b/snowflake/ml/modeling/_internal/snowpark_implementations/snowpark_handlers.py index 9dac22b5..c55bc1f7 100644 --- a/snowflake/ml/modeling/_internal/snowpark_implementations/snowpark_handlers.py +++ b/snowflake/ml/modeling/_internal/snowpark_implementations/snowpark_handlers.py @@ -154,6 +154,7 @@ def load_estimator() -> object: imports=imports, # type: ignore[arg-type] ) def vec_batch_infer(input_df: pd.DataFrame) -> T.PandasSeries[dict]: # type: ignore[type-arg] + import joblib import numpy as np # noqa: F401 import pandas as pd @@ -164,7 +165,16 @@ def vec_batch_infer(input_df: pd.DataFrame) -> T.PandasSeries[dict]: # type: ig if hasattr(estimator, "n_jobs"): # Vectorized UDF cannot handle joblib multiprocessing right now, deactivate the n_jobs estimator.n_jobs = 1 - inference_res = getattr(estimator, inference_method)(input_df, *args, **kwargs) + + # Birch exposes no n_jobs param, so the guard above does not reach it. Its internal pairwise + # distance computation hits scikit-learn #33877 (_parallel_pairwise slices Y but not + # Y_norm_squared when effective_n_jobs > 1) on the pinned scikit-learn (<1.9). Pin a single-job + # backend for it. Remove this carve-out once scikit-learn is bumped past 1.9.0. (SNOW-952252) + if type(estimator).__name__ == "Birch": + with joblib.parallel_backend("threading", n_jobs=1): + inference_res = getattr(estimator, inference_method)(input_df, *args, **kwargs) + else: + inference_res = getattr(estimator, inference_method)(input_df, *args, **kwargs) transformed_numpy_array, _ = handle_inference_result( inference_res=inference_res, diff --git a/snowflake/ml/modeling/_internal/snowpark_implementations/snowpark_trainer.py b/snowflake/ml/modeling/_internal/snowpark_implementations/snowpark_trainer.py index 1d17feab..73d395ca 100644 --- a/snowflake/ml/modeling/_internal/snowpark_implementations/snowpark_trainer.py +++ b/snowflake/ml/modeling/_internal/snowpark_implementations/snowpark_trainer.py @@ -193,11 +193,16 @@ def fit_and_return_estimator() -> str: args["sample_weight"] = df[sample_weight_col].squeeze() X = np.ascontiguousarray(args["X"].to_numpy(), dtype=np.float64) - if ArgKmin.is_usable_for(X, X, "euclidean"): - estimator.fit(**args) - else: + # Birch's internal clustering routes through scikit-learn's _parallel_pairwise, which on + # the pinned scikit-learn (<1.9) slices Y but not Y_norm_squared when effective_n_jobs > 1, + # raising a dimension mismatch (scikit-learn #33877). It takes that path regardless of + # ArgKmin usability on the training data, so force a single-job backend for it. Remove this + # carve-out once scikit-learn is bumped past 1.9.0. (SNOW-952252) + if type(estimator).__name__ == "Birch" or not ArgKmin.is_usable_for(X, X, "euclidean"): with joblib.parallel_backend("threading", n_jobs=1): estimator.fit(**args) + else: + estimator.fit(**args) local_result_file_name = temp_file_utils.get_temp_file_path() @@ -313,11 +318,13 @@ def fit_predict_wrapper_function( estimator = cp.load(local_transform_file_obj) X = np.ascontiguousarray(df[input_cols].to_numpy(), dtype=np.float64) - if ArgKmin.is_usable_for(X, X, "euclidean"): - fit_predict_result = estimator.fit_predict(X=df[input_cols]) - else: + # Force a single-job backend for Birch to avoid scikit-learn #33877 (see fit_wrapper_function + # for the full explanation). (SNOW-952252) + if type(estimator).__name__ == "Birch" or not ArgKmin.is_usable_for(X, X, "euclidean"): with joblib.parallel_backend("threading", n_jobs=1): fit_predict_result = estimator.fit_predict(X=df[input_cols]) + else: + fit_predict_result = estimator.fit_predict(X=df[input_cols]) local_result_file_name = temp_file_utils.get_temp_file_path() @@ -451,11 +458,13 @@ def fit_transform_wrapper_function( args["sample_weight"] = df[sample_weight_col].squeeze() X = np.ascontiguousarray(args["X"].to_numpy(), dtype=np.float64) - if ArgKmin.is_usable_for(X, X, "euclidean"): - fit_transform_result = estimator.fit_transform(**args) - else: + # Force a single-job backend for Birch to avoid scikit-learn #33877 (see fit_wrapper_function + # for the full explanation). (SNOW-952252) + if type(estimator).__name__ == "Birch" or not ArgKmin.is_usable_for(X, X, "euclidean"): with joblib.parallel_backend("threading", n_jobs=1): fit_transform_result = estimator.fit_transform(**args) + else: + fit_transform_result = estimator.fit_transform(**args) local_result_file_name = temp_file_utils.get_temp_file_path() diff --git a/snowflake/ml/version.py b/snowflake/ml/version.py index 8dce423e..b72595d2 100644 --- a/snowflake/ml/version.py +++ b/snowflake/ml/version.py @@ -1,2 +1,2 @@ # This is parsed by regex in conda recipe meta file. Make sure not to break it. -VERSION = "1.44.0" +VERSION = "1.45.0" diff --git a/tests/integ/snowflake/ml/feature_store/feature_store_feature_group_bundled.py b/tests/integ/snowflake/ml/feature_store/feature_store_feature_group_bundled.py index 12e2b7c3..cef6fb7a 100644 --- a/tests/integ/snowflake/ml/feature_store/feature_store_feature_group_bundled.py +++ b/tests/integ/snowflake/ml/feature_store/feature_store_feature_group_bundled.py @@ -31,6 +31,7 @@ from __future__ import annotations +import json import logging import time import uuid @@ -333,6 +334,74 @@ def _register_postgres_tiled_bfv(self, *, suffix: str) -> tuple[str, str]: return fv_name, seeded_user_id + def _register_postgres_secondary_key_bfv(self, *, suffix: str) -> tuple[str, str]: + """Create + register a tiled batch FV on Postgres with an aggregation secondary key. + + Args: + suffix: Short label embedded in the FV / source-table names. + + Returns: + Tuple of ``(fv_name, seeded_user_id)``. ``seeded_user_id`` matches a + seeded row so downstream reads return non-empty data. + """ + s = uuid.uuid4().hex[:8] + fv_name = f"FG_INTEG_SK_BFV_{suffix}_{s}" + seeded_user_id = f"U_SK_{suffix}_{s}" + src_table = f"{self.test_db}.{self.fs._config.schema.identifier()}.FG_SK_SRC_{suffix}_{s}" + + self._session.sql( + f""" + CREATE OR REPLACE TABLE {src_table} ( + USER_ID VARCHAR, + EVENT_TIME TIMESTAMP_NTZ, + AD_ID VARCHAR, + AMOUNT FLOAT + ) + """ + ).collect() + yesterday = "DATEADD('day', -1, DATE_TRUNC('day', CURRENT_TIMESTAMP()::TIMESTAMP_NTZ))" + self._session.sql( + f""" + INSERT INTO {src_table} + SELECT column1, column2, column3, column4 + FROM VALUES + ({seeded_user_id!r}, DATEADD('hour', 1, {yesterday}), 'ad_a', 10.0), + ({seeded_user_id!r}, DATEADD('hour', 2, {yesterday}), 'ad_a', 20.0), + ({seeded_user_id!r}, DATEADD('hour', 3, {yesterday}), 'ad_b', 70.0) + """ + ).collect() + feature_df = self._session.table(src_table) + + # Secondary keys reject list aggregations; SUM/COUNT are supported. + features = [ + Feature.sum("AMOUNT", "3d").alias(f"AMOUNT_SUM_3D_{suffix}"), + Feature.count("AMOUNT", "3d").alias(f"TXN_COUNT_3D_{suffix}"), + ] + fv = FeatureView( + name=fv_name, + entities=[self.user_entity], + feature_df=feature_df, + timestamp_col="EVENT_TIME", + refresh_mode="FULL", + refresh_freq="1 minute", + feature_granularity="1d", + features=features, + aggregation_secondary_keys=["AD_ID"], + online_config=OnlineConfig(enable=True, target_lag="10s", store_type=OnlineStoreType.POSTGRES), + ) + registered_fv = self.fs.register_feature_view(fv, "v1") + self.assertTrue(registered_fv.is_tiled) + self.assertEqual(registered_fv.aggregation_secondary_keys, ["AD_ID"]) + + deadline = time.time() + 180.0 + while time.time() < deadline: + count = self._session.table(registered_fv.fully_qualified_name()).count() + if count > 0: + break + time.sleep(5) + + return fv_name, seeded_user_id + def _wait_until_fg_read_returns_rows(self, fg_live: FeatureGroup, key: str, timeout: float = 600.0) -> None: """Poll ``read_feature_group`` until at least one row is returned.""" deadline = time.time() + timeout @@ -1044,6 +1113,69 @@ def test_feature_group_with_tiled_bfv_online_read(self) -> None: finally: self.fs.delete_feature_group(fg_name, fg_version) + def test_feature_group_with_secondary_key_bfv_online_read(self) -> None: + """A FG over a secondary-key tiled BFV registers and reads back array-shaped agg outputs.""" + sfx = "SK" + fv_sk_name, seeded_user_id = self._register_postgres_secondary_key_bfv(suffix=sfx) + fv_sk = self.fs.get_feature_view(fv_sk_name, "v1") + + fg_name = f"FG_INTEG_SK_BFV_{uuid.uuid4().hex[:8].upper()}" + fg_version = "v1" + fg = FeatureGroup(name=fg_name, features=[fv_sk], auto_prefix=False) + + try: + # This call regressed before the secondary-key array-wrapping fix. + registered = self.fs.register_feature_group(fg, fg_version) + normalized_outputs = {_normalize_column_name(c) for c in registered.output_columns} + for agg in (f"AMOUNT_SUM_3D_{sfx}", f"TXN_COUNT_3D_{sfx}"): + self.assertTrue( + any(agg in c for c in normalized_outputs), + f"expected agg '{agg}' in FG output_columns; got {sorted(normalized_outputs)}", + ) + + fg_live = self.fs.get_feature_group(fg_name, fg_version) + + sum_col_token = f"AMOUNT_SUM_3D_{sfx}" + deadline = time.time() + 600.0 + pdf: Optional[pd.DataFrame] = None + value = None + while time.time() < deadline: + pdf = self._read_feature_group_with_retry(fg_live, keys=[[seeded_user_id]]) + if len(pdf) > 0: + sum_col = next((c for c in pdf.columns if sum_col_token in _normalize_column_name(c)), None) + candidate = pdf.iloc[0][sum_col] if sum_col is not None else None + # ``x != x`` detects float NaN without importing math; arrays + # are object-dtype so this only trips on a scalar-null cell. + is_nan = isinstance(candidate, float) and candidate != candidate + if candidate is not None and not is_nan: + value = candidate + break + time.sleep(5) + + self.assertIsNotNone(pdf, "read_feature_group never returned a DataFrame for the secondary-key FG.") + assert pdf is not None + # Output columns are structural, present regardless of data lag. + pdf_cols = {_normalize_column_name(c) for c in pdf.columns} + for agg in (f"AMOUNT_SUM_3D_{sfx}", f"TXN_COUNT_3D_{sfx}"): + self.assertTrue( + any(agg in c for c in pdf_cols), + f"expected agg '{agg}' in read columns; got {sorted(pdf_cols)}", + ) + + self.assertIsNotNone( + value, + f"secondary-key value column '{sum_col_token}' did not materialize a non-null value " + f"within 600s; cannot verify the array-shape regression guard.", + ) + # The secondary-key value columns come back as arrays (one element per + # AD_ID bucket), never bare scalars. ``list(...)`` accepts + # list/tuple/ndarray and raises on a scalar, so this fails loudly if + # the array wrapping regresses. + coerced = json.loads(value) if isinstance(value, str) else value + self.assertIsInstance(list(coerced), list) + finally: + self.fs.delete_feature_group(fg_name, fg_version) + def test_generate_training_set_from_feature_group_with_tiled_sfv(self) -> None: """``generate_training_set(feature_group=...)`` works when one upstream FV is tiled.""" sfx = "TS2" diff --git a/tests/integ/snowflake/ml/feature_store/feature_store_online_management_test.py b/tests/integ/snowflake/ml/feature_store/feature_store_online_management_test.py index 0a123a12..afde1f1b 100644 --- a/tests/integ/snowflake/ml/feature_store/feature_store_online_management_test.py +++ b/tests/integ/snowflake/ml/feature_store/feature_store_online_management_test.py @@ -3,9 +3,10 @@ import common_utils from absl.testing import absltest, parameterized -from fs_integ_test_base import FeatureStoreIntegTestBase +from fs_integ_test_base import FeatureStoreIntegTestBase from snowflake.ml.feature_store import entity, feature_store, feature_view +from snowflake.ml.feature_store.feature import Feature class FeatureStoreOnlineTest(FeatureStoreIntegTestBase, parameterized.TestCase): @@ -199,6 +200,58 @@ def test_update_feature_view_enable_online(self) -> None: self.assertIsNotNone(updated_fv.online_config) self.assertEqual(updated_fv.online_config.target_lag, "15 seconds") + def test_update_enable_hybrid_online_on_tiled_fv_rejected(self) -> None: + """Enabling HYBRID_TABLE online on a tiled FV is unsupported and rejected with a clear error. + + The rejection fires fast during update_feature_view input validation, before any planning + or resource work. + """ + fv_name = "test_enable_hybrid_online_tiled_fv" + + # Offline-only tiled feature view. + fv = feature_view.FeatureView( + name=fv_name, + entities=[self.user_entity], + feature_df=self.sample_data.select("user_id", "purchase_amount", "purchase_time"), + timestamp_col="purchase_time", + refresh_freq="15m", + feature_granularity="1h", + features=[Feature.sum("purchase_amount", "2h").alias("amount_sum_2h")], + desc="Test enable online on tiled FV", + online_config=feature_view.OnlineConfig(enable=False), + ) + + registered_fv = self.fs.register_feature_view(fv, "v1") + self.assertFalse(registered_fv.online) + self.assertTrue(registered_fv.is_tiled) + + # Enabling HYBRID_TABLE (the default) online on a tiled FV must be rejected clearly. + with self.assertRaisesRegex(Exception, "not supported for aggregation"): + self.fs.update_feature_view( + name=fv_name, + version="v1", + online_config=feature_view.OnlineConfig(enable=True), + ) + + def test_register_tiled_fv_with_hybrid_online_rejected(self) -> None: + """A tiled FV with HYBRID_TABLE online is unsupported and rejected with a clear error. + + The rejection is a feature-view invariant, so it fires at construction (before register). + """ + with self.assertRaisesRegex(Exception, "not supported for aggregation"): + fv = feature_view.FeatureView( + name="test_register_hybrid_online_tiled_fv", + entities=[self.user_entity], + feature_df=self.sample_data.select("user_id", "purchase_amount", "purchase_time"), + timestamp_col="purchase_time", + refresh_freq="15m", + feature_granularity="1h", + features=[Feature.sum("purchase_amount", "2h").alias("amount_sum_2h")], + desc="Test register tiled FV with HYBRID online", + online_config=feature_view.OnlineConfig(enable=True), # default HYBRID_TABLE + ) + self.fs.register_feature_view(fv, "v1") + def test_update_feature_view_disable_online(self) -> None: """Test disabling online storage for existing feature view.""" fv_name = "test_update_disable_online_fv" diff --git a/tests/integ/snowflake/ml/feature_store/feature_store_rollup_test.py b/tests/integ/snowflake/ml/feature_store/feature_store_rollup_test.py index a799811c..7cc51d54 100644 --- a/tests/integ/snowflake/ml/feature_store/feature_store_rollup_test.py +++ b/tests/integ/snowflake/ml/feature_store/feature_store_rollup_test.py @@ -11,6 +11,7 @@ import snowflake.ml.version as snowml_version from fs_integ_test_base import FeatureStoreIntegTestBase +from snowflake.ml._internal.utils.sql_identifier import SqlIdentifier from snowflake.ml.feature_store import Feature, RollupConfig from snowflake.ml.feature_store.entity import Entity from snowflake.ml.feature_store.feature_store import CreationMode, FeatureStore @@ -168,6 +169,54 @@ def test_rollup_fv_registration(self) -> None: self.assertEqual(registered_subscriber.name, "SUBSCRIBER_EVENTS") self.assertEqual(registered_subscriber.version, "v1") + def test_rollup_fv_initialization_warehouse(self) -> None: + """A rollup FV's dynamic table carries INITIALIZATION_WAREHOUSE when set.""" + fs = self._create_feature_store() + + visitor_entity = self._create_visitor_entity() + subscriber_entity = self._create_subscriber_entity() + fs.register_entity(visitor_entity) + fs.register_entity(subscriber_entity) + + visitor_fv = FeatureView( + name="visitor_events", + entities=[visitor_entity], + feature_df=self._get_events_df(), + timestamp_col="event_ts", + refresh_freq="1h", + feature_granularity="1h", + features=[Feature.count("visitor_id", "24h").alias("event_count")], + ) + registered_visitor = fs.register_feature_view(visitor_fv, "v1") + + subscriber_fv = FeatureView( + name="subscriber_events", + entities=[subscriber_entity], + rollup_config=RollupConfig( + source=registered_visitor, + mapping_df=self._get_mapping_df(), + ), + initialization_warehouse=self._alt_warehouse_name, + ) + registered_subscriber = fs.register_feature_view(subscriber_fv, "v1") + self.assertEqual( + registered_subscriber.initialization_warehouse, + SqlIdentifier(self._alt_warehouse_name), + ) + + self.assertEqual( + fs.get_feature_view("subscriber_events", "v1").initialization_warehouse, + SqlIdentifier(self._alt_warehouse_name), + ) + fv_name = FeatureView._get_physical_name(SqlIdentifier("SUBSCRIBER_EVENTS"), "v1") + dt_row = self._session.sql( + f"SHOW DYNAMIC TABLES LIKE '{fv_name.resolved()}' IN SCHEMA {fs._config.full_schema_path}" + ).collect()[0] + self.assertEqual( + SqlIdentifier(dt_row["initialization_warehouse"], case_sensitive=True), + SqlIdentifier(self._alt_warehouse_name), + ) + def test_rollup_fv_is_tiled(self) -> None: """Test that rollup FV has is_tiled=True.""" fs = self._create_feature_store() diff --git a/tests/integ/snowflake/ml/feature_store/feature_store_streaming_fv_bundled.py b/tests/integ/snowflake/ml/feature_store/feature_store_streaming_fv_bundled.py index 3a4623ea..317066ad 100644 --- a/tests/integ/snowflake/ml/feature_store/feature_store_streaming_fv_bundled.py +++ b/tests/integ/snowflake/ml/feature_store/feature_store_streaming_fv_bundled.py @@ -103,6 +103,52 @@ def test_register_streaming_fv_basic(self) -> None: udf_count = self._session.table(fq_udf).count() self.assertEqual(udf_count, 3, "udf_transformed table should have 3 backfill rows") + def test_streaming_fv_initialization_warehouse(self) -> None: + """A streaming FV's offline dynamic table carries INITIALIZATION_WAREHOUSE when set. + + The backfill task graph's use of the initialization warehouse is covered + deterministically by the streaming_registration unit tests; the graph + self-drops on completion, so it is not asserted here. + """ + s = uuid.uuid4().hex[:8] + stream = f"TXN_{s}" + fv_name = f"STREAM_FV_{s}" + fs = self._create_feature_store() + self._make_stream_source(fs, stream) + backfill_table = self._create_backfill_table(fs, s) + backfill_df = self._session.table(backfill_table) + stream_config = StreamConfig( + stream_source=stream, + transformation_fn=identity_transform, + backfill_df=backfill_df, + ) + fv = FeatureView( + name=fv_name, + entities=[self.user_entity], + stream_config=stream_config, + timestamp_col="EVENT_TIME", + refresh_freq="1 minute", + warehouse=self._test_warehouse_name, + initialization_warehouse=self._alt_warehouse_name, + ) + + registered_fv = fs.register_feature_view(fv, "v1") + self.assertTrue(registered_fv.is_streaming) + self.assertEqual(registered_fv.initialization_warehouse, SqlIdentifier(self._alt_warehouse_name)) + + self.assertEqual( + fs.get_feature_view(fv_name, "v1").initialization_warehouse, + SqlIdentifier(self._alt_warehouse_name), + ) + physical_name = FeatureView._get_physical_name(registered_fv.name, registered_fv.version) + dt_row = self._session.sql( + f"SHOW DYNAMIC TABLES LIKE '{physical_name.resolved()}' IN SCHEMA {fs._config.full_schema_path}" + ).collect()[0] + self.assertEqual( + SqlIdentifier(dt_row["initialization_warehouse"], case_sensitive=True), + SqlIdentifier(self._alt_warehouse_name), + ) + def _create_backfill_table_with_schema(self, fs, suffix: str, schema_sql: str, insert_values_sql: str) -> str: """Create a backfill table with caller-supplied column DDL and INSERT values.""" table_name = f"{self.test_db}.{fs._config.schema.identifier()}.BACKFILL_{suffix}" diff --git a/tests/integ/snowflake/ml/feature_store/feature_store_streaming_fv_integ_base.py b/tests/integ/snowflake/ml/feature_store/feature_store_streaming_fv_integ_base.py index a5c465f9..1568ca62 100644 --- a/tests/integ/snowflake/ml/feature_store/feature_store_streaming_fv_integ_base.py +++ b/tests/integ/snowflake/ml/feature_store/feature_store_streaming_fv_integ_base.py @@ -22,6 +22,7 @@ from fs_integ_test_base import FeatureStoreIntegTestBase, cleanup_spec_oft_e2e_databases from snowflake.connector.errors import DatabaseError +from snowflake.ml._internal.exceptions.exceptions import SnowflakeMLException from snowflake.ml._internal.utils.sql_identifier import SqlIdentifier from snowflake.ml.feature_store.entity import Entity from snowflake.ml.feature_store.feature_store import CreationMode, FeatureStore @@ -201,7 +202,7 @@ def _has_query_endpoint(st: Any) -> bool: on_recreate() try: fs.create_online_service(producer_role, consumer_role) - except DatabaseError as e: + except (DatabaseError, SnowflakeMLException) as e: logger.warning("create_online_service (recreate) raised: %s; will poll for RUNNING.", e) break time.sleep(poll_interval_s) diff --git a/tests/integ/snowflake/ml/feature_store/feature_store_test.py b/tests/integ/snowflake/ml/feature_store/feature_store_test.py index 453a67e5..8a8480c7 100644 --- a/tests/integ/snowflake/ml/feature_store/feature_store_test.py +++ b/tests/integ/snowflake/ml/feature_store/feature_store_test.py @@ -571,6 +571,108 @@ def test_register_feature_view_with_warehouse(self) -> None: self._alt_warehouse_name, ) + def test_register_feature_view_with_initialization_warehouse(self) -> None: + fs = self._create_feature_store() + + e = Entity("foo", ["id"]) + fs.register_entity(e) + + sql = f"SELECT id, name, title, ts FROM {self._mock_table}" + + # init-WH set: round-trips on the object and the list view, and the + # backing dynamic table reports it via the SHOW initialization_warehouse column. + d1 = FeatureView( + name="fv1", + entities=[e], + feature_df=self._session.sql(sql), + timestamp_col="ts", + refresh_freq="1d", + warehouse=self._test_warehouse_name, + initialization_warehouse=self._alt_warehouse_name, + ) + r1 = fs.register_feature_view(feature_view=d1, version="1.0") + self.assertEqual(r1.warehouse, self._test_warehouse_name) + self.assertEqual(r1.initialization_warehouse, SqlIdentifier(self._alt_warehouse_name)) + + # Full read-back round-trip recovers the value from the DT definition. + self.assertEqual( + fs.get_feature_view("fv1", "1.0").initialization_warehouse, + SqlIdentifier(self._alt_warehouse_name), + ) + self.assertEqual( + fs.list_feature_views(verbose=True) + .select("initialization_warehouse") + .filter(col("version") == "1.0") + .collect()[0]["INITIALIZATION_WAREHOUSE"], + self._alt_warehouse_name, + ) + fv_name = FeatureView._get_physical_name(SqlIdentifier("FV1"), "1.0") + dt_row = self._session.sql( + f"SHOW DYNAMIC TABLES LIKE '{fv_name.resolved()}' IN SCHEMA {fs._config.full_schema_path}" + ).collect()[0] + self.assertEqual( + SqlIdentifier(dt_row["initialization_warehouse"], case_sensitive=True), + SqlIdentifier(self._alt_warehouse_name), + ) + + # init-WH unset: the clause is omitted, behaving like single-warehouse FVs. + d2 = FeatureView( + name="fv2", + entities=[e], + feature_df=self._session.sql(sql), + timestamp_col="ts", + refresh_freq="1d", + warehouse=self._test_warehouse_name, + ) + r2 = fs.register_feature_view(feature_view=d2, version="1.0") + self.assertIsNone(r2.initialization_warehouse) + self.assertIsNone(fs.get_feature_view("fv2", "1.0").initialization_warehouse) + fv2_name = FeatureView._get_physical_name(SqlIdentifier("FV2"), "1.0") + dt2_row = self._session.sql( + f"SHOW DYNAMIC TABLES LIKE '{fv2_name.resolved()}' IN SCHEMA {fs._config.full_schema_path}" + ).collect()[0] + self.assertFalse(dt2_row["initialization_warehouse"]) + + def test_update_initialization_warehouse(self) -> None: + fs = self._create_feature_store() + + e = Entity("FOO", ["id"]) + fs.register_entity(e) + + sql = f"SELECT id, name, title FROM {self._mock_table}" + fv = FeatureView( + name="fv1", + entities=[e], + feature_df=self._session.sql(sql), + refresh_freq="1 minute", + warehouse=self._test_warehouse_name, + ) + fv = fs.register_feature_view(feature_view=fv, version="v1") + self.assertIsNone(fv.initialization_warehouse) + + # SET via update. + fs.update_feature_view("fv1", "v1", initialization_warehouse=self._alt_warehouse_name) + self.assertEqual( + fs.get_feature_view("fv1", "v1").initialization_warehouse, + SqlIdentifier(self._alt_warehouse_name), + ) + + # Omitting the argument leaves it unchanged. + fs.update_feature_view("fv1", "v1", refresh_freq="2 minute") + self.assertEqual( + fs.get_feature_view("fv1", "v1").initialization_warehouse, + SqlIdentifier(self._alt_warehouse_name), + ) + + # UNSET via explicit None. + fs.update_feature_view("fv1", "v1", initialization_warehouse=None) + self.assertIsNone(fs.get_feature_view("fv1", "v1").initialization_warehouse) + fv_name = FeatureView._get_physical_name(SqlIdentifier("FV1"), "v1") + dt_text = self._session.sql( + f"SHOW DYNAMIC TABLES LIKE '{fv_name.resolved()}' IN SCHEMA {fs._config.full_schema_path}" + ).collect()[0]["text"] + self.assertNotIn("INITIALIZATION_WAREHOUSE", dt_text.upper()) + def test_register_feature_view_with_unregistered_entity(self) -> None: fs = self._create_feature_store() @@ -2747,15 +2849,20 @@ def test_update_static_feature_view(self) -> None: ) fv = fs.register_feature_view(feature_view=fv, version="v1") with self.assertRaisesRegex( - RuntimeError, "Static feature view '.*' does not support refresh_freq and warehouse." + RuntimeError, "Static feature view '.*' does not support refresh_freq, warehouse, " ): fs.update_feature_view("fv1", "v1", refresh_freq="1 minute") with self.assertRaisesRegex( - RuntimeError, "Static feature view '.*' does not support refresh_freq and warehouse." + RuntimeError, "Static feature view '.*' does not support refresh_freq, warehouse, " ): fs.update_feature_view("fv1", "v1", warehouse=self._session.get_current_warehouse()) + with self.assertRaisesRegex( + RuntimeError, "Static feature view '.*' does not support refresh_freq, warehouse, " + ): + fs.update_feature_view("fv1", "v1", initialization_warehouse=self._session.get_current_warehouse()) + updated_fv = fs.update_feature_view(fv, desc="") self.assertEqual(updated_fv.desc, "") # Static feature views should have no storage_config diff --git a/tests/integ/snowflake/ml/jobs/jobs_integ_test.py b/tests/integ/snowflake/ml/jobs/jobs_integ_test.py index 74071ebc..bc576928 100644 --- a/tests/integ/snowflake/ml/jobs/jobs_integ_test.py +++ b/tests/integ/snowflake/ml/jobs/jobs_integ_test.py @@ -684,7 +684,8 @@ def job_sproc(session: snowpark.Session) -> str: return job.get_logs() result = job_sproc(self.session) - self.assertEqual("Hello from remote function!", result) + # Use containment since get_logs() can include framework footer markers around the user output. + self.assertIn("Hello from remote function!", result) def test_job_decorator_negative_result(self) -> None: @jobs.remote(self.compute_pool, stage_name="payload_stage", session=self.session) diff --git a/tests/integ/snowflake/ml/model/_client/model/model_version_impl_integ_test.py b/tests/integ/snowflake/ml/model/_client/model/model_version_impl_integ_test.py index 6c9c999f..a3961e98 100644 --- a/tests/integ/snowflake/ml/model/_client/model/model_version_impl_integ_test.py +++ b/tests/integ/snowflake/ml/model/_client/model/model_version_impl_integ_test.py @@ -17,6 +17,51 @@ MODEL_NAME = "TEST_MODEL" VERSION_NAME = "V1" +_MODEL_EXPORT_FILES = [ + "model", + "model/model.yaml", + "model/runtimes", + "model/models", + "model/env", + "model/runtimes/cpu", + "model/runtimes/cpu/env", + "model/runtimes/cpu/env/requirements.txt", + "model/models/explain_artifacts", + "model/models/TEST_MODEL", + "model/models/explain_artifacts/TEST_MODEL_background_data.pqt", + "model/models/TEST_MODEL/model.pkl", + "model/env/requirements.txt", + "model/sample_input_data.json", +] + +_FULL_EXPORT_FILES = [ + "runtimes", + "model", + "MANIFEST.yml", + "functions", + "runtimes/python_runtime", + "runtimes/python_runtime/env", + "runtimes/python_runtime/env/requirements.txt", + "model/model.yaml", + "model/runtimes", + "model/models", + "model/env", + "model/runtimes/cpu", + "model/runtimes/cpu/env", + "model/runtimes/cpu/env/requirements.txt", + "model/models/explain_artifacts", + "model/models/TEST_MODEL", + "model/models/explain_artifacts/TEST_MODEL_background_data.pqt", + "model/models/TEST_MODEL/model.pkl", + "model/env/requirements.txt", + "functions/decision_function.py", + "functions/predict_log_proba.py", + "functions/predict_proba.py", + "functions/predict.py", + "functions/explain.py", + "model/sample_input_data.json", +] + class TestModelVersionImplInteg(parameterized.TestCase): @classmethod @@ -76,70 +121,29 @@ def test_metrics(self) -> None: with self.assertRaises(KeyError): self._mv.get_metric("b") + def _assert_export_file_list( + self, + tmpdir: str, + export_mode: ExportMode, + actual_file_list: list[str], + ) -> None: + base_files = _FULL_EXPORT_FILES if export_mode == ExportMode.FULL else _MODEL_EXPORT_FILES + expected_file_list = [os.path.join(tmpdir, rel_path) for rel_path in base_files] + expected_file_list.extend(path for path in actual_file_list if path.endswith("conda.yml")) + self.assertSameElements(expected_file_list, actual_file_list) + def test_export(self) -> None: with tempfile.TemporaryDirectory() as tmpdir: self._mv.export(tmpdir) - expected_file_list = [ - "model", - "model/model.yaml", - "model/runtimes", - "model/models", - "model/env", - "model/runtimes/cpu", - "model/runtimes/cpu/env", - "model/runtimes/cpu/env/requirements.txt", - "model/runtimes/cpu/env/conda.yml", - "model/models/explain_artifacts", - "model/models/TEST_MODEL", - "model/models/explain_artifacts/TEST_MODEL_background_data.pqt", - "model/models/TEST_MODEL/model.pkl", - "model/env/requirements.txt", - "model/env/conda.yml", - "model/sample_input_data.json", - ] - expected_file_list = [os.path.join(tmpdir, expected_file) for expected_file in expected_file_list] actual_file_list = list(glob.iglob(os.path.join(tmpdir, "**", "*"), recursive=True)) - # remove "snowflake-ml-python.zip" from the actual file list actual_file_list = [file for file in actual_file_list if not file.endswith("snowflake-ml-python.zip")] - self.assertSameElements(actual_file_list, expected_file_list) + self._assert_export_file_list(tmpdir, ExportMode.MODEL, actual_file_list) with tempfile.TemporaryDirectory() as tmpdir: self._mv.export(tmpdir, export_mode=ExportMode.FULL) - expected_file_list = [ - "runtimes", - "model", - "MANIFEST.yml", - "functions", - "runtimes/python_runtime", - "runtimes/python_runtime/env", - "runtimes/python_runtime/env/requirements.txt", - "runtimes/python_runtime/env/conda.yml", - "model/model.yaml", - "model/runtimes", - "model/models", - "model/env", - "model/runtimes/cpu", - "model/runtimes/cpu/env", - "model/runtimes/cpu/env/requirements.txt", - "model/runtimes/cpu/env/conda.yml", - "model/models/explain_artifacts", - "model/models/TEST_MODEL", - "model/models/explain_artifacts/TEST_MODEL_background_data.pqt", - "model/models/TEST_MODEL/model.pkl", - "model/env/requirements.txt", - "model/env/conda.yml", - "functions/decision_function.py", - "functions/predict_log_proba.py", - "functions/predict_proba.py", - "functions/predict.py", - "functions/explain.py", - "model/sample_input_data.json", - ] - expected_file_list = [os.path.join(tmpdir, expected_file) for expected_file in expected_file_list] actual_file_list = list(glob.iglob(os.path.join(tmpdir, "**", "*"), recursive=True)) - # remove "snowflake-ml-python.zip" from the actual file list actual_file_list = [file for file in actual_file_list if not file.endswith("snowflake-ml-python.zip")] - self.assertSameElements(actual_file_list, expected_file_list) + self._assert_export_file_list(tmpdir, ExportMode.FULL, actual_file_list) def test_load(self) -> None: loaded_model = self._mv.load(force=True) diff --git a/tests/integ/snowflake/ml/registry/jobs/registry_batch_inference_cse_stage_test.py b/tests/integ/snowflake/ml/registry/jobs/registry_batch_inference_cse_stage_test.py index a56f1b13..4ea0711b 100644 --- a/tests/integ/snowflake/ml/registry/jobs/registry_batch_inference_cse_stage_test.py +++ b/tests/integ/snowflake/ml/registry/jobs/registry_batch_inference_cse_stage_test.py @@ -16,6 +16,7 @@ from tests.integ.snowflake.ml.registry.jobs import registry_batch_inference_test_base +@absltest.skip("SNOW-3691662") class TestRegistryBatchInferenceCSEStageInteg(registry_batch_inference_test_base.RegistryBatchInferenceTestBase): """Test batch inference with CSE (SNOWFLAKE_FULL) encrypted stage.""" diff --git a/tests/integ/snowflake/ml/registry/jobs/registry_custom_multi_modality_batch_inference_test.py b/tests/integ/snowflake/ml/registry/jobs/registry_custom_multi_modality_batch_inference_test.py index b0768731..bd95d0a3 100644 --- a/tests/integ/snowflake/ml/registry/jobs/registry_custom_multi_modality_batch_inference_test.py +++ b/tests/integ/snowflake/ml/registry/jobs/registry_custom_multi_modality_batch_inference_test.py @@ -61,6 +61,7 @@ def signature(cls) -> core.ModelSignature: ) +@absltest.skip("SNOW-3691662") class TestRegistryCustomMultiModalityBatchInferenceInteg( registry_batch_inference_test_base.RegistryBatchInferenceTestBase ): diff --git a/tests/integ/snowflake/ml/registry/jobs/registry_huggingface_pipeline_batch_inference_test.py b/tests/integ/snowflake/ml/registry/jobs/registry_huggingface_pipeline_batch_inference_test.py index d7d4b1e5..bd16feb8 100644 --- a/tests/integ/snowflake/ml/registry/jobs/registry_huggingface_pipeline_batch_inference_test.py +++ b/tests/integ/snowflake/ml/registry/jobs/registry_huggingface_pipeline_batch_inference_test.py @@ -33,6 +33,8 @@ def test_text_generation( self, pip_requirements: Optional[list[str]], ) -> None: + if pip_requirements is None: + self.skipTest("SNOW-3691662") import transformers model = transformers.pipeline( diff --git a/tests/integ/snowflake/ml/registry/jobs/registry_keras_batch_inference_test.py b/tests/integ/snowflake/ml/registry/jobs/registry_keras_batch_inference_test.py index 50cab1d8..c94b2409 100644 --- a/tests/integ/snowflake/ml/registry/jobs/registry_keras_batch_inference_test.py +++ b/tests/integ/snowflake/ml/registry/jobs/registry_keras_batch_inference_test.py @@ -24,6 +24,7 @@ def _prepare_keras_functional_model() -> tuple[keras.Model, npt.ArrayLike, npt.A class TestKerasBatchInferenceInteg(registry_batch_inference_test_base.RegistryBatchInferenceTestBase): + @absltest.skip("SNOW-3691662") @parameterized.parameters( # type: ignore[misc] {"gpu_requests": None, "cpu_requests": None, "memory_requests": None}, # uncomment this after https://snowflakecomputing.atlassian.net/browse/SNOW-2369772 is fixed diff --git a/tests/integ/snowflake/ml/registry/jobs/registry_multi_modality_hugging_face_pipeline_batch_inference_test.py b/tests/integ/snowflake/ml/registry/jobs/registry_multi_modality_hugging_face_pipeline_batch_inference_test.py index 2509ca1b..77a7ccae 100644 --- a/tests/integ/snowflake/ml/registry/jobs/registry_multi_modality_hugging_face_pipeline_batch_inference_test.py +++ b/tests/integ/snowflake/ml/registry/jobs/registry_multi_modality_hugging_face_pipeline_batch_inference_test.py @@ -13,6 +13,7 @@ from tests.integ.snowflake.ml.registry.jobs import registry_batch_inference_test_base +@absltest.skip("SNOW-3691662") class TestRegistryMultiModalityHuggingFacePipelineBatchInferenceInteg( registry_batch_inference_test_base.RegistryBatchInferenceTestBase ): diff --git a/tests/integ/snowflake/ml/registry/jobs/registry_new_hf_tasks_batch_inference_test.py b/tests/integ/snowflake/ml/registry/jobs/registry_new_hf_tasks_batch_inference_test.py index ac40dcda..1352439f 100644 --- a/tests/integ/snowflake/ml/registry/jobs/registry_new_hf_tasks_batch_inference_test.py +++ b/tests/integ/snowflake/ml/registry/jobs/registry_new_hf_tasks_batch_inference_test.py @@ -159,6 +159,7 @@ def check_vqa(res: pd.DataFrame) -> None: prediction_assert_fn=check_vqa, ) + @absltest.skip("SNOW-3691662") def test_image_feature_extraction(self) -> None: from transformers import pipeline @@ -285,6 +286,7 @@ def check_image_to_text(res: pd.DataFrame) -> None: prediction_assert_fn=check_image_to_text, ) + @absltest.skip("SNOW-3691662") def test_object_detection(self) -> None: from transformers import pipeline @@ -414,6 +416,7 @@ def check_zero_shot_image_classification(res: pd.DataFrame) -> None: prediction_assert_fn=check_zero_shot_image_classification, ) + @absltest.skip("SNOW-3691662") def test_zero_shot_object_detection(self) -> None: from transformers import pipeline diff --git a/tests/integ/snowflake/ml/registry/jobs/registry_pip_only_batch_inference_test.py b/tests/integ/snowflake/ml/registry/jobs/registry_pip_only_batch_inference_test.py index 02371b4d..f347ccda 100644 --- a/tests/integ/snowflake/ml/registry/jobs/registry_pip_only_batch_inference_test.py +++ b/tests/integ/snowflake/ml/registry/jobs/registry_pip_only_batch_inference_test.py @@ -1,19 +1,15 @@ """Integration tests for pip-only model packaging with batch inference.""" -import sys -from typing import Any, Callable +from typing import Any import pandas as pd -from absl.testing import absltest, parameterized +from absl.testing import absltest from snowflake.ml.model import custom_model from snowflake.ml.model.batch import JobSpec, OutputSpec from tests.integ.snowflake.ml.registry import pip_only_packaging_integ_util from tests.integ.snowflake.ml.registry.jobs import registry_batch_inference_test_base -# Python versions supported by the base image's cached standalone tarballs (dockerfile_template_pip path). -PIP_ONLY_PYTHON_VERSIONS = ("3.10", "3.11", "3.12", "3.13", "3.14") - class PipOnlyModel(custom_model.CustomModel): """A simple custom model with only pip dependencies.""" @@ -67,6 +63,7 @@ def check_env(self, input: pd.DataFrame) -> pd.DataFrame: ) +@absltest.skip("SNOW-3691688") class TestRegistryPipOnlyBatchInferenceInteg( pip_only_packaging_integ_util.PipOnlyPackagingIntegMixin, registry_batch_inference_test_base.RegistryBatchInferenceTestBase, @@ -124,19 +121,6 @@ def _assert_pip_only_env(self, res: pd.DataFrame) -> None: f"python_executable={res['python_executable'].iloc[0]}", ) - def _assert_pip_only_predict_result(self, py_ver: str, expected: pd.DataFrame) -> Callable[[pd.DataFrame], None]: - """Assert predict result matches expected and runtime python_version matches requested py_ver.""" - - def fn(res: pd.DataFrame) -> None: - self.assertEqual(res["python_version"].iloc[0], py_ver) - pd.testing.assert_series_equal( - res["output"].sort_values().reset_index(drop=True), - expected["output"].sort_values().reset_index(drop=True), - check_dtype=False, - ) - - return fn - def _run_pip_only_env_check_batch_inference( self, model: custom_model.CustomModel, @@ -206,52 +190,6 @@ def test_pip_only_batch_inference(self) -> None: options={"enable_explainability": False}, ) - @parameterized.parameters(*PIP_ONLY_PYTHON_VERSIONS) # type: ignore[misc] - def test_pip_only_batch_inference_python_versions(self, py_ver: str) -> None: - """Batch inference with a pip-only model for each supported Python version (3.10, 3.11, 3.12, 3.13, 3.14). - - Verifies: - 1. Model runs with the correct Python version - 2. Environment uses a pip-only path (venv) - """ - current_ver = f"{sys.version_info.major}.{sys.version_info.minor}" - if py_ver != current_ver: - self.skipTest( - f"Skipping Python {py_ver} test: model is pickled with {current_ver} and " - f"cloudpickle cannot deserialize across Python versions." - ) - - model = PipOnlyModel(custom_model.ModelContext()) - input_pandas_df = pd.DataFrame({"value": [1.0, 2.0, 3.0]}) - - # Generate expected predictions for assertion - model_output = model.predict(input_pandas_df) - - sp_df, input_df, _, job_name, output_stage_location = self._prepare_pip_only_test(model, input_pandas_df) - - self._test_registry_batch_inference( - model=model, - sample_input_data=sp_df, - X=input_df, - output_spec=OutputSpec(stage_location=output_stage_location), - job_spec=JobSpec( - job_name=job_name, - num_workers=1, - replicas=1, - function_name="predict", - ), - pip_requirements=["requests>=2.28.0"], - options={"enable_explainability": False}, - python_version=py_ver, - prediction_assert_fn=self._assert_pip_only_predict_result(py_ver, model_output), - conda_dependencies=[], - ) - self._run_pip_only_env_check_batch_inference( - model=model, - pip_requirements=["requests>=2.28.0"], - options={"enable_explainability": False}, - ) - if __name__ == "__main__": absltest.main() diff --git a/tests/integ/snowflake/ml/registry/jobs/registry_pip_only_gpu_batch_inference_test.py b/tests/integ/snowflake/ml/registry/jobs/registry_pip_only_gpu_batch_inference_test.py index 600b70ca..b49ac43c 100644 --- a/tests/integ/snowflake/ml/registry/jobs/registry_pip_only_gpu_batch_inference_test.py +++ b/tests/integ/snowflake/ml/registry/jobs/registry_pip_only_gpu_batch_inference_test.py @@ -39,6 +39,7 @@ def predict(self, input: pd.DataFrame) -> pd.DataFrame: return pd.DataFrame({"output": y.detach().cpu().numpy().ravel()}) +@absltest.skip("SNOW-3691688") class TestRegistryPipOnlyGpuBatchInferenceInteg( pip_only_packaging_integ_util.PipOnlyPackagingIntegMixin, registry_batch_inference_test_base.RegistryBatchInferenceTestBase, diff --git a/tests/integ/snowflake/ml/registry/jobs/registry_pytorch_batch_inference_test.py b/tests/integ/snowflake/ml/registry/jobs/registry_pytorch_batch_inference_test.py index 402553a7..59e6325f 100644 --- a/tests/integ/snowflake/ml/registry/jobs/registry_pytorch_batch_inference_test.py +++ b/tests/integ/snowflake/ml/registry/jobs/registry_pytorch_batch_inference_test.py @@ -9,6 +9,7 @@ from tests.integ.snowflake.ml.test_utils import model_factory +@absltest.skip("SNOW-3691662") class TestPyTorchBatchInferenceInteg(registry_batch_inference_test_base.RegistryBatchInferenceTestBase): @parameterized.parameters( # type: ignore[misc] {"gpu_requests": None, "cpu_requests": None, "memory_requests": None}, diff --git a/tests/integ/snowflake/ml/registry/jobs/registry_sentence_transformers_batch_inference_test.py b/tests/integ/snowflake/ml/registry/jobs/registry_sentence_transformers_batch_inference_test.py index 22b01126..c56ba246 100644 --- a/tests/integ/snowflake/ml/registry/jobs/registry_sentence_transformers_batch_inference_test.py +++ b/tests/integ/snowflake/ml/registry/jobs/registry_sentence_transformers_batch_inference_test.py @@ -8,6 +8,9 @@ from packaging import version as pkg_version from snowflake.ml.model._packager.model_env import model_env +from snowflake.ml.model._packager.model_handlers.sentence_transformers import ( + _encode_sentences_with_nulls, +) from snowflake.ml.model.batch import JobSpec, OutputSpec from snowflake.ml.model.models import huggingface as snowml_huggingface from tests.integ.snowflake.ml.registry.jobs import registry_batch_inference_test_base @@ -17,6 +20,7 @@ HF_HOME = "HF_HOME" +@absltest.skip("SNOW-3691662") class TestRegistrySentenceTransformerBatchInferenceInteg( registry_batch_inference_test_base.RegistryBatchInferenceTestBase ): @@ -101,6 +105,47 @@ def test_sentence_transformers( expected_predictions=expected_predictions, ) + def test_sentence_transformers_null_input(self) -> None: + import sentence_transformers + + sentences = pd.DataFrame( + { + "sentence": [ + "Why don't scientists trust atoms? Because they make up everything.", + None, + "Parallel lines have so much in common. It's a shame they'll never meet.", + ] + } + ) + job_name, output_stage_location, _ = self._prepare_job_name_and_stage_for_batch_inference() + + model = sentence_transformers.SentenceTransformer(random.choice(MODEL_NAMES)) + model_output_normalized = _encode_sentences_with_nulls( + sentences["sentence"].tolist(), + model.encode, + {}, + ) + model_output_df = pd.DataFrame({"output": model_output_normalized}) + + input_df, expected_predictions = self._prepare_batch_inference_data(sentences, model_output_df) + + self._test_registry_batch_inference( + model=model, + sample_input_data=sentences, + options={"cuda_version": model_env.DEFAULT_CUDA_VERSION}, + pip_requirements=["sentence-transformers"], + X=input_df, + output_spec=OutputSpec(stage_location=output_stage_location), + job_spec=JobSpec( + job_name=job_name, + num_workers=1, + replicas=1, + gpu_requests=None, + function_name="encode", + ), + expected_predictions=expected_predictions, + ) + @parameterized.parameters( # type: ignore[misc] {"pip_requirements": ["sentence-transformers"], "gpu_requests": None}, {"pip_requirements": None, "gpu_requests": "1"}, diff --git a/tests/integ/snowflake/ml/registry/jobs/registry_tensorflow_batch_inference_test.py b/tests/integ/snowflake/ml/registry/jobs/registry_tensorflow_batch_inference_test.py index 70c1476d..ff38ec58 100644 --- a/tests/integ/snowflake/ml/registry/jobs/registry_tensorflow_batch_inference_test.py +++ b/tests/integ/snowflake/ml/registry/jobs/registry_tensorflow_batch_inference_test.py @@ -8,6 +8,7 @@ class TestTensorFlowBatchInferenceInteg(registry_batch_inference_test_base.RegistryBatchInferenceTestBase): + @absltest.skip("SNOW-3691662") @parameterized.parameters( # type: ignore[misc] {"gpu_requests": None, "cpu_requests": None, "memory_requests": None}, {"gpu_requests": "1", "cpu_requests": None, "memory_requests": None}, diff --git a/tests/integ/snowflake/ml/registry/model/BUILD.bazel b/tests/integ/snowflake/ml/registry/model/BUILD.bazel index 0f03e6a3..d78f16a7 100644 --- a/tests/integ/snowflake/ml/registry/model/BUILD.bazel +++ b/tests/integ/snowflake/ml/registry/model/BUILD.bazel @@ -325,10 +325,10 @@ py_test( py_test( name = "registry_in_sproc_test", + timeout = "long", srcs = ["registry_in_sproc_test.py"], tags = [ "feature:model_registry", - "short_regress", ], deps = [ ":registry_model_test_base", diff --git a/tests/integ/snowflake/ml/registry/model/registry_huggingface_pipeline_model_test.py b/tests/integ/snowflake/ml/registry/model/registry_huggingface_pipeline_model_test.py index 5c38d8cf..1ea1581d 100644 --- a/tests/integ/snowflake/ml/registry/model/registry_huggingface_pipeline_model_test.py +++ b/tests/integ/snowflake/ml/registry/model/registry_huggingface_pipeline_model_test.py @@ -151,7 +151,7 @@ def check_res_with_params(res: pd.DataFrame) -> None: pd.testing.assert_index_equal(res.columns, pd.Index(["score", "start", "end", "answer"])) self.assertEqual(res["score"].dtype.type, np.float64) for answer in res["answer"]: - self.assertLessEqual(len(answer), 5) + self.assertIsNotNone(answer) self._test_registry_model( model=model, @@ -216,7 +216,7 @@ def check_res_with_params(res: pd.DataFrame) -> None: self.assertIsInstance(row, list) self.assertIn("answer", row[0]) for entry in row: - self.assertLessEqual(len(entry["answer"]), 5) + self.assertIsNotNone(entry["answer"]) self._test_registry_model( model=model, @@ -353,7 +353,7 @@ def check_res_with_params(res: pd.DataFrame) -> None: params_assert_fns={ "": ( x_df, - {"sequential": True, "padding": "max_length", "truncation": "only_first"}, + {"sequential": True, "padding": "max_length", "truncation": "drop_rows_to_fit"}, check_res_with_params, ), }, diff --git a/tests/integ/snowflake/ml/registry/services/registry_pip_only_model_deployment_test.py b/tests/integ/snowflake/ml/registry/services/registry_pip_only_model_deployment_test.py index befb482c..deff8a42 100644 --- a/tests/integ/snowflake/ml/registry/services/registry_pip_only_model_deployment_test.py +++ b/tests/integ/snowflake/ml/registry/services/registry_pip_only_model_deployment_test.py @@ -1,9 +1,7 @@ """Integration tests for pip-only model packaging and deployment.""" -from typing import Callable - import pandas as pd -from absl.testing import absltest, parameterized +from absl.testing import absltest from snowflake.ml.model import custom_model from snowflake.ml.model._packager.model_env import model_env @@ -12,9 +10,6 @@ registry_model_deployment_test_base, ) -# Python versions supported by the base image's cached standalone tarballs (dockerfile_template_pip path). -PIP_ONLY_PYTHON_VERSIONS = ("3.10", "3.11", "3.12", "3.13", "3.14") - class PipOnlyModel(custom_model.CustomModel): """A simple custom model with only pip dependencies.""" @@ -114,15 +109,6 @@ class TestRegistryPipOnlyModelDeploymentInteg( Uses the same Kaniko builder override as ``RegistryModelDeploymentTestBase`` (``BUILDER_IMAGE_PATH``). """ - def _assert_pip_only_predict_result(self, py_ver: str, expected: pd.DataFrame) -> Callable[[pd.DataFrame], None]: - """Assert predict result matches expected and runtime python_version matches requested py_ver.""" - - def fn(res: pd.DataFrame) -> None: - self.assertEqual(res["python_version"].iloc[0], py_ver) - pd.testing.assert_frame_equal(res, expected, check_dtype=False) - - return fn - def _assert_pip_only_env(self, res: pd.DataFrame) -> None: """Assert that the deployed service is using a pip-only path (venv).""" self.assertTrue( @@ -205,39 +191,6 @@ def test_pip_only_gpu_model(self) -> None: conda_dependencies=[], ) - @parameterized.parameters(*PIP_ONLY_PYTHON_VERSIONS) # type: ignore[misc] - def test_pip_only_model_python_versions(self, py_ver: str) -> None: - """E2E test: deploy a pip-only model with each supported Python version (3.10, 3.11, 3.12, 3.13, 3.14). - - Verifies: - 1. Model runs with the correct Python version - 2. Environment uses a pip-only path (venv) - """ - if not self._has_image_override(): - self.skipTest("Skipping pip-only model deployment test: image override not enabled.") - - test_input = pd.DataFrame({"value": [1.0, 2.0, 3.0]}) - model = PipOnlyModel(custom_model.ModelContext()) - self._test_registry_model_deployment( - model=model, - sample_input_data=test_input, - prediction_assert_fns={ - "predict": ( - test_input, - self._assert_pip_only_predict_result(py_ver, model.predict(test_input)), - ), - "check_env": ( - pd.DataFrame({"value": [1.0]}), - self._assert_pip_only_env, - ), - }, - pip_requirements=["requests>=2.28.0"], - options={"enable_explainability": False}, - python_version=py_ver, - service_name=f"service_pip_only_python_versions_{self._run_id}_{py_ver.replace('.', '')}", - conda_dependencies=[], - ) - def test_pip_only_pytorch_gpu_model(self) -> None: """E2E test: deploy a pip-only PyTorch model on GPU and verify CUDA is available.