snowflake-ml-python 1.6.1__py3-none-any.whl → 1.6.3__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- snowflake/cortex/__init__.py +4 -0
- snowflake/cortex/_classify_text.py +2 -2
- snowflake/cortex/_embed_text_1024.py +37 -0
- snowflake/cortex/_embed_text_768.py +37 -0
- snowflake/cortex/_extract_answer.py +2 -2
- snowflake/cortex/_sentiment.py +2 -2
- snowflake/cortex/_summarize.py +2 -2
- snowflake/cortex/_translate.py +2 -2
- snowflake/cortex/_util.py +4 -4
- snowflake/ml/_internal/env_utils.py +5 -5
- snowflake/ml/_internal/exceptions/error_codes.py +2 -0
- snowflake/ml/_internal/telemetry.py +142 -20
- snowflake/ml/_internal/utils/db_utils.py +50 -0
- snowflake/ml/_internal/utils/identifier.py +48 -11
- snowflake/ml/_internal/utils/service_logger.py +63 -0
- snowflake/ml/_internal/utils/snowflake_env.py +23 -13
- snowflake/ml/_internal/utils/sql_identifier.py +26 -2
- snowflake/ml/_internal/utils/table_manager.py +19 -1
- snowflake/ml/data/_internal/arrow_ingestor.py +1 -11
- snowflake/ml/data/data_connector.py +33 -7
- snowflake/ml/data/ingestor_utils.py +20 -10
- snowflake/ml/data/torch_utils.py +68 -0
- snowflake/ml/dataset/dataset.py +1 -3
- snowflake/ml/feature_store/access_manager.py +3 -3
- snowflake/ml/feature_store/feature_store.py +60 -19
- snowflake/ml/feature_store/feature_view.py +84 -30
- snowflake/ml/fileset/embedded_stage_fs.py +1 -1
- snowflake/ml/fileset/fileset.py +1 -1
- snowflake/ml/fileset/sfcfs.py +9 -3
- snowflake/ml/fileset/stage_fs.py +2 -1
- snowflake/ml/lineage/lineage_node.py +7 -2
- snowflake/ml/model/__init__.py +1 -2
- snowflake/ml/model/_client/model/model_version_impl.py +96 -12
- snowflake/ml/model/_client/ops/model_ops.py +124 -6
- snowflake/ml/model/_client/ops/service_ops.py +309 -9
- snowflake/ml/model/_client/service/model_deployment_spec.py +8 -5
- snowflake/ml/model/_client/service/model_deployment_spec_schema.py +2 -2
- snowflake/ml/model/_client/sql/_base.py +5 -0
- snowflake/ml/model/_client/sql/model.py +1 -0
- snowflake/ml/model/_client/sql/model_version.py +9 -5
- snowflake/ml/model/_client/sql/service.py +121 -20
- snowflake/ml/model/_model_composer/model_composer.py +11 -39
- snowflake/ml/model/_model_composer/model_manifest/model_manifest.py +31 -11
- snowflake/ml/model/_packager/model_env/model_env.py +4 -38
- snowflake/ml/model/_packager/model_handlers/_utils.py +134 -28
- snowflake/ml/model/_packager/model_handlers/catboost.py +31 -30
- snowflake/ml/model/_packager/model_handlers/huggingface_pipeline.py +26 -18
- snowflake/ml/model/_packager/model_handlers/lightgbm.py +31 -58
- snowflake/ml/model/_packager/model_handlers/mlflow.py +3 -5
- snowflake/ml/model/_packager/model_handlers/model_objective_utils.py +169 -0
- snowflake/ml/model/_packager/model_handlers/sentence_transformers.py +15 -8
- snowflake/ml/model/_packager/model_handlers/sklearn.py +56 -60
- snowflake/ml/model/_packager/model_handlers/snowmlmodel.py +141 -9
- snowflake/ml/model/_packager/model_handlers/torchscript.py +2 -2
- snowflake/ml/model/_packager/model_handlers/xgboost.py +63 -48
- snowflake/ml/model/_packager/model_meta/model_meta.py +16 -42
- snowflake/ml/model/_packager/model_meta/model_meta_schema.py +1 -14
- snowflake/ml/model/_packager/model_packager.py +14 -8
- snowflake/ml/model/_packager/model_runtime/model_runtime.py +11 -0
- snowflake/ml/model/_signatures/pytorch_handler.py +1 -1
- snowflake/ml/model/_signatures/snowpark_handler.py +3 -2
- snowflake/ml/model/_signatures/utils.py +9 -0
- snowflake/ml/model/type_hints.py +12 -145
- snowflake/ml/modeling/_internal/constants.py +1 -0
- snowflake/ml/modeling/_internal/local_implementations/pandas_handlers.py +5 -5
- snowflake/ml/modeling/_internal/local_implementations/pandas_trainer.py +9 -6
- snowflake/ml/modeling/_internal/model_specifications.py +2 -0
- snowflake/ml/modeling/_internal/model_trainer.py +1 -0
- snowflake/ml/modeling/_internal/snowpark_implementations/distributed_hpo_trainer.py +2 -4
- snowflake/ml/modeling/_internal/snowpark_implementations/snowpark_handlers.py +5 -5
- snowflake/ml/modeling/_internal/snowpark_implementations/snowpark_trainer.py +130 -166
- snowflake/ml/modeling/_internal/snowpark_implementations/xgboost_external_memory_trainer.py +0 -1
- snowflake/ml/modeling/calibration/calibrated_classifier_cv.py +61 -21
- snowflake/ml/modeling/cluster/affinity_propagation.py +61 -21
- snowflake/ml/modeling/cluster/agglomerative_clustering.py +61 -21
- snowflake/ml/modeling/cluster/birch.py +61 -21
- snowflake/ml/modeling/cluster/bisecting_k_means.py +61 -21
- snowflake/ml/modeling/cluster/dbscan.py +61 -21
- snowflake/ml/modeling/cluster/feature_agglomeration.py +61 -21
- snowflake/ml/modeling/cluster/k_means.py +61 -21
- snowflake/ml/modeling/cluster/mean_shift.py +61 -21
- snowflake/ml/modeling/cluster/mini_batch_k_means.py +61 -21
- snowflake/ml/modeling/cluster/optics.py +61 -21
- snowflake/ml/modeling/cluster/spectral_biclustering.py +61 -21
- snowflake/ml/modeling/cluster/spectral_clustering.py +61 -21
- snowflake/ml/modeling/cluster/spectral_coclustering.py +61 -21
- snowflake/ml/modeling/compose/column_transformer.py +61 -21
- snowflake/ml/modeling/compose/transformed_target_regressor.py +61 -21
- snowflake/ml/modeling/covariance/elliptic_envelope.py +61 -21
- snowflake/ml/modeling/covariance/empirical_covariance.py +61 -21
- snowflake/ml/modeling/covariance/graphical_lasso.py +61 -21
- snowflake/ml/modeling/covariance/graphical_lasso_cv.py +61 -21
- snowflake/ml/modeling/covariance/ledoit_wolf.py +61 -21
- snowflake/ml/modeling/covariance/min_cov_det.py +61 -21
- snowflake/ml/modeling/covariance/oas.py +61 -21
- snowflake/ml/modeling/covariance/shrunk_covariance.py +61 -21
- snowflake/ml/modeling/decomposition/dictionary_learning.py +61 -21
- snowflake/ml/modeling/decomposition/factor_analysis.py +61 -21
- snowflake/ml/modeling/decomposition/fast_ica.py +61 -21
- snowflake/ml/modeling/decomposition/incremental_pca.py +61 -21
- snowflake/ml/modeling/decomposition/kernel_pca.py +61 -21
- snowflake/ml/modeling/decomposition/mini_batch_dictionary_learning.py +61 -21
- snowflake/ml/modeling/decomposition/mini_batch_sparse_pca.py +61 -21
- snowflake/ml/modeling/decomposition/pca.py +61 -21
- snowflake/ml/modeling/decomposition/sparse_pca.py +61 -21
- snowflake/ml/modeling/decomposition/truncated_svd.py +61 -21
- snowflake/ml/modeling/discriminant_analysis/linear_discriminant_analysis.py +61 -21
- snowflake/ml/modeling/discriminant_analysis/quadratic_discriminant_analysis.py +61 -21
- snowflake/ml/modeling/ensemble/ada_boost_classifier.py +61 -21
- snowflake/ml/modeling/ensemble/ada_boost_regressor.py +61 -21
- snowflake/ml/modeling/ensemble/bagging_classifier.py +61 -21
- snowflake/ml/modeling/ensemble/bagging_regressor.py +61 -21
- snowflake/ml/modeling/ensemble/extra_trees_classifier.py +61 -21
- snowflake/ml/modeling/ensemble/extra_trees_regressor.py +61 -21
- snowflake/ml/modeling/ensemble/gradient_boosting_classifier.py +61 -21
- snowflake/ml/modeling/ensemble/gradient_boosting_regressor.py +61 -21
- snowflake/ml/modeling/ensemble/hist_gradient_boosting_classifier.py +61 -21
- snowflake/ml/modeling/ensemble/hist_gradient_boosting_regressor.py +61 -21
- snowflake/ml/modeling/ensemble/isolation_forest.py +61 -21
- snowflake/ml/modeling/ensemble/random_forest_classifier.py +61 -21
- snowflake/ml/modeling/ensemble/random_forest_regressor.py +61 -21
- snowflake/ml/modeling/ensemble/stacking_regressor.py +61 -21
- snowflake/ml/modeling/ensemble/voting_classifier.py +61 -21
- snowflake/ml/modeling/ensemble/voting_regressor.py +61 -21
- snowflake/ml/modeling/feature_selection/generic_univariate_select.py +61 -21
- snowflake/ml/modeling/feature_selection/select_fdr.py +61 -21
- snowflake/ml/modeling/feature_selection/select_fpr.py +61 -21
- snowflake/ml/modeling/feature_selection/select_fwe.py +61 -21
- snowflake/ml/modeling/feature_selection/select_k_best.py +61 -21
- snowflake/ml/modeling/feature_selection/select_percentile.py +61 -21
- snowflake/ml/modeling/feature_selection/sequential_feature_selector.py +61 -21
- snowflake/ml/modeling/feature_selection/variance_threshold.py +61 -21
- snowflake/ml/modeling/gaussian_process/gaussian_process_classifier.py +61 -21
- snowflake/ml/modeling/gaussian_process/gaussian_process_regressor.py +61 -21
- snowflake/ml/modeling/impute/iterative_imputer.py +61 -21
- snowflake/ml/modeling/impute/knn_imputer.py +61 -21
- snowflake/ml/modeling/impute/missing_indicator.py +61 -21
- snowflake/ml/modeling/kernel_approximation/additive_chi2_sampler.py +61 -21
- snowflake/ml/modeling/kernel_approximation/nystroem.py +61 -21
- snowflake/ml/modeling/kernel_approximation/polynomial_count_sketch.py +61 -21
- snowflake/ml/modeling/kernel_approximation/rbf_sampler.py +61 -21
- snowflake/ml/modeling/kernel_approximation/skewed_chi2_sampler.py +61 -21
- snowflake/ml/modeling/kernel_ridge/kernel_ridge.py +61 -21
- snowflake/ml/modeling/lightgbm/lgbm_classifier.py +61 -21
- snowflake/ml/modeling/lightgbm/lgbm_regressor.py +61 -21
- snowflake/ml/modeling/linear_model/ard_regression.py +61 -21
- snowflake/ml/modeling/linear_model/bayesian_ridge.py +61 -21
- snowflake/ml/modeling/linear_model/elastic_net.py +61 -21
- snowflake/ml/modeling/linear_model/elastic_net_cv.py +61 -21
- snowflake/ml/modeling/linear_model/gamma_regressor.py +61 -21
- snowflake/ml/modeling/linear_model/huber_regressor.py +61 -21
- snowflake/ml/modeling/linear_model/lars.py +61 -21
- snowflake/ml/modeling/linear_model/lars_cv.py +61 -21
- snowflake/ml/modeling/linear_model/lasso.py +61 -21
- snowflake/ml/modeling/linear_model/lasso_cv.py +61 -21
- snowflake/ml/modeling/linear_model/lasso_lars.py +61 -21
- snowflake/ml/modeling/linear_model/lasso_lars_cv.py +61 -21
- snowflake/ml/modeling/linear_model/lasso_lars_ic.py +61 -21
- snowflake/ml/modeling/linear_model/linear_regression.py +61 -21
- snowflake/ml/modeling/linear_model/logistic_regression.py +61 -21
- snowflake/ml/modeling/linear_model/logistic_regression_cv.py +61 -21
- snowflake/ml/modeling/linear_model/multi_task_elastic_net.py +61 -21
- snowflake/ml/modeling/linear_model/multi_task_elastic_net_cv.py +61 -21
- snowflake/ml/modeling/linear_model/multi_task_lasso.py +61 -21
- snowflake/ml/modeling/linear_model/multi_task_lasso_cv.py +61 -21
- snowflake/ml/modeling/linear_model/orthogonal_matching_pursuit.py +61 -21
- snowflake/ml/modeling/linear_model/passive_aggressive_classifier.py +61 -21
- snowflake/ml/modeling/linear_model/passive_aggressive_regressor.py +61 -21
- snowflake/ml/modeling/linear_model/perceptron.py +61 -21
- snowflake/ml/modeling/linear_model/poisson_regressor.py +61 -21
- snowflake/ml/modeling/linear_model/ransac_regressor.py +61 -21
- snowflake/ml/modeling/linear_model/ridge.py +61 -21
- snowflake/ml/modeling/linear_model/ridge_classifier.py +61 -21
- snowflake/ml/modeling/linear_model/ridge_classifier_cv.py +61 -21
- snowflake/ml/modeling/linear_model/ridge_cv.py +61 -21
- snowflake/ml/modeling/linear_model/sgd_classifier.py +61 -21
- snowflake/ml/modeling/linear_model/sgd_one_class_svm.py +61 -21
- snowflake/ml/modeling/linear_model/sgd_regressor.py +61 -21
- snowflake/ml/modeling/linear_model/theil_sen_regressor.py +61 -21
- snowflake/ml/modeling/linear_model/tweedie_regressor.py +61 -21
- snowflake/ml/modeling/manifold/isomap.py +61 -21
- snowflake/ml/modeling/manifold/mds.py +61 -21
- snowflake/ml/modeling/manifold/spectral_embedding.py +61 -21
- snowflake/ml/modeling/manifold/tsne.py +61 -21
- snowflake/ml/modeling/metrics/metrics_utils.py +2 -2
- snowflake/ml/modeling/metrics/ranking.py +0 -3
- snowflake/ml/modeling/metrics/regression.py +0 -3
- snowflake/ml/modeling/mixture/bayesian_gaussian_mixture.py +61 -21
- snowflake/ml/modeling/mixture/gaussian_mixture.py +61 -21
- snowflake/ml/modeling/multiclass/one_vs_one_classifier.py +61 -21
- snowflake/ml/modeling/multiclass/one_vs_rest_classifier.py +61 -21
- snowflake/ml/modeling/multiclass/output_code_classifier.py +61 -21
- snowflake/ml/modeling/naive_bayes/bernoulli_nb.py +61 -21
- snowflake/ml/modeling/naive_bayes/categorical_nb.py +61 -21
- snowflake/ml/modeling/naive_bayes/complement_nb.py +61 -21
- snowflake/ml/modeling/naive_bayes/gaussian_nb.py +61 -21
- snowflake/ml/modeling/naive_bayes/multinomial_nb.py +61 -21
- snowflake/ml/modeling/neighbors/k_neighbors_classifier.py +61 -21
- snowflake/ml/modeling/neighbors/k_neighbors_regressor.py +61 -21
- snowflake/ml/modeling/neighbors/kernel_density.py +61 -21
- snowflake/ml/modeling/neighbors/local_outlier_factor.py +61 -21
- snowflake/ml/modeling/neighbors/nearest_centroid.py +61 -21
- snowflake/ml/modeling/neighbors/nearest_neighbors.py +61 -21
- snowflake/ml/modeling/neighbors/neighborhood_components_analysis.py +61 -21
- snowflake/ml/modeling/neighbors/radius_neighbors_classifier.py +61 -21
- snowflake/ml/modeling/neighbors/radius_neighbors_regressor.py +61 -21
- snowflake/ml/modeling/neural_network/bernoulli_rbm.py +61 -21
- snowflake/ml/modeling/neural_network/mlp_classifier.py +61 -21
- snowflake/ml/modeling/neural_network/mlp_regressor.py +61 -21
- snowflake/ml/modeling/parameters/disable_model_tracer.py +5 -0
- snowflake/ml/modeling/pipeline/pipeline.py +1 -13
- snowflake/ml/modeling/preprocessing/polynomial_features.py +61 -21
- snowflake/ml/modeling/semi_supervised/label_propagation.py +61 -21
- snowflake/ml/modeling/semi_supervised/label_spreading.py +61 -21
- snowflake/ml/modeling/svm/linear_svc.py +61 -21
- snowflake/ml/modeling/svm/linear_svr.py +61 -21
- snowflake/ml/modeling/svm/nu_svc.py +61 -21
- snowflake/ml/modeling/svm/nu_svr.py +61 -21
- snowflake/ml/modeling/svm/svc.py +61 -21
- snowflake/ml/modeling/svm/svr.py +61 -21
- snowflake/ml/modeling/tree/decision_tree_classifier.py +61 -21
- snowflake/ml/modeling/tree/decision_tree_regressor.py +61 -21
- snowflake/ml/modeling/tree/extra_tree_classifier.py +61 -21
- snowflake/ml/modeling/tree/extra_tree_regressor.py +61 -21
- snowflake/ml/modeling/xgboost/xgb_classifier.py +64 -23
- snowflake/ml/modeling/xgboost/xgb_regressor.py +64 -23
- snowflake/ml/modeling/xgboost/xgbrf_classifier.py +64 -23
- snowflake/ml/modeling/xgboost/xgbrf_regressor.py +64 -23
- snowflake/ml/monitoring/_client/model_monitor.py +126 -0
- snowflake/ml/monitoring/_client/model_monitor_manager.py +361 -0
- snowflake/ml/monitoring/_client/model_monitor_version.py +1 -0
- snowflake/ml/monitoring/_client/monitor_sql_client.py +1335 -0
- snowflake/ml/monitoring/_client/queries/record_count.ssql +14 -0
- snowflake/ml/monitoring/_client/queries/rmse.ssql +28 -0
- snowflake/ml/monitoring/entities/model_monitor_config.py +28 -0
- snowflake/ml/monitoring/entities/model_monitor_interval.py +46 -0
- snowflake/ml/monitoring/entities/output_score_type.py +90 -0
- snowflake/ml/registry/_manager/model_manager.py +4 -0
- snowflake/ml/registry/registry.py +166 -8
- snowflake/ml/version.py +1 -1
- {snowflake_ml_python-1.6.1.dist-info → snowflake_ml_python-1.6.3.dist-info}/METADATA +43 -9
- snowflake_ml_python-1.6.3.dist-info/RECORD +400 -0
- {snowflake_ml_python-1.6.1.dist-info → snowflake_ml_python-1.6.3.dist-info}/WHEEL +1 -1
- snowflake/ml/_internal/container_services/image_registry/credential.py +0 -84
- snowflake/ml/_internal/container_services/image_registry/http_client.py +0 -127
- snowflake/ml/_internal/container_services/image_registry/imagelib.py +0 -400
- snowflake/ml/_internal/container_services/image_registry/registry_client.py +0 -212
- snowflake/ml/_internal/utils/log_stream_processor.py +0 -30
- snowflake/ml/_internal/utils/session_token_manager.py +0 -46
- snowflake/ml/_internal/utils/spcs_attribution_utils.py +0 -122
- snowflake/ml/_internal/utils/uri.py +0 -77
- snowflake/ml/data/torch_dataset.py +0 -33
- snowflake/ml/model/_api.py +0 -568
- snowflake/ml/model/_deploy_client/image_builds/base_image_builder.py +0 -12
- snowflake/ml/model/_deploy_client/image_builds/client_image_builder.py +0 -249
- snowflake/ml/model/_deploy_client/image_builds/docker_context.py +0 -130
- snowflake/ml/model/_deploy_client/image_builds/gunicorn_run.sh +0 -36
- snowflake/ml/model/_deploy_client/image_builds/inference_server/main.py +0 -268
- snowflake/ml/model/_deploy_client/image_builds/server_image_builder.py +0 -215
- snowflake/ml/model/_deploy_client/image_builds/templates/dockerfile_template +0 -53
- snowflake/ml/model/_deploy_client/image_builds/templates/image_build_job_spec_template +0 -38
- snowflake/ml/model/_deploy_client/image_builds/templates/kaniko_shell_script_template +0 -105
- snowflake/ml/model/_deploy_client/snowservice/deploy.py +0 -611
- snowflake/ml/model/_deploy_client/snowservice/deploy_options.py +0 -116
- snowflake/ml/model/_deploy_client/snowservice/instance_types.py +0 -10
- snowflake/ml/model/_deploy_client/snowservice/templates/service_spec_template +0 -28
- snowflake/ml/model/_deploy_client/snowservice/templates/service_spec_template_with_model +0 -21
- snowflake/ml/model/_deploy_client/utils/constants.py +0 -48
- snowflake/ml/model/_deploy_client/utils/snowservice_client.py +0 -280
- snowflake/ml/model/_deploy_client/warehouse/deploy.py +0 -202
- snowflake/ml/model/_deploy_client/warehouse/infer_template.py +0 -99
- snowflake/ml/model/_packager/model_handlers/llm.py +0 -267
- snowflake/ml/model/_packager/model_meta/_core_requirements.py +0 -11
- snowflake/ml/model/deploy_platforms.py +0 -6
- snowflake/ml/model/models/llm.py +0 -104
- snowflake/ml/monitoring/monitor.py +0 -203
- snowflake/ml/registry/_initial_schema.py +0 -142
- snowflake/ml/registry/_schema.py +0 -82
- snowflake/ml/registry/_schema_upgrade_plans.py +0 -116
- snowflake/ml/registry/_schema_version_manager.py +0 -163
- snowflake/ml/registry/model_registry.py +0 -2048
- snowflake_ml_python-1.6.1.dist-info/RECORD +0 -422
- {snowflake_ml_python-1.6.1.dist-info → snowflake_ml_python-1.6.3.dist-info}/LICENSE.txt +0 -0
- {snowflake_ml_python-1.6.1.dist-info → snowflake_ml_python-1.6.3.dist-info}/top_level.txt +0 -0
@@ -1,249 +0,0 @@
|
|
1
|
-
import json
|
2
|
-
import logging
|
3
|
-
import os
|
4
|
-
import shutil
|
5
|
-
import subprocess
|
6
|
-
import tempfile
|
7
|
-
import time
|
8
|
-
from enum import Enum
|
9
|
-
from typing import List
|
10
|
-
|
11
|
-
from snowflake import snowpark
|
12
|
-
from snowflake.ml._internal.container_services.image_registry import credential
|
13
|
-
from snowflake.ml._internal.exceptions import (
|
14
|
-
error_codes,
|
15
|
-
exceptions as snowml_exceptions,
|
16
|
-
)
|
17
|
-
from snowflake.ml.model._deploy_client.image_builds import base_image_builder
|
18
|
-
|
19
|
-
logger = logging.getLogger(__name__)
|
20
|
-
|
21
|
-
|
22
|
-
class Platform(Enum):
|
23
|
-
LINUX_AMD64 = "linux/amd64"
|
24
|
-
|
25
|
-
|
26
|
-
class ClientImageBuilder(base_image_builder.ImageBuilder):
|
27
|
-
"""
|
28
|
-
Client-side image building and upload to model registry.
|
29
|
-
|
30
|
-
Usage requirements:
|
31
|
-
Requires prior installation and running of Docker with BuildKit. See installation instructions in
|
32
|
-
https://docs.docker.com/engine/install/
|
33
|
-
"""
|
34
|
-
|
35
|
-
def __init__(
|
36
|
-
self,
|
37
|
-
*,
|
38
|
-
context_dir: str,
|
39
|
-
full_image_name: str,
|
40
|
-
image_repo: str,
|
41
|
-
session: snowpark.Session,
|
42
|
-
) -> None:
|
43
|
-
"""Initialization
|
44
|
-
|
45
|
-
Args:
|
46
|
-
context_dir: Local docker context dir.
|
47
|
-
full_image_name: Full image name consists of image name and image tag.
|
48
|
-
image_repo: Path to image repository.
|
49
|
-
session: Snowpark session
|
50
|
-
"""
|
51
|
-
self.context_dir = context_dir
|
52
|
-
self.full_image_name = full_image_name
|
53
|
-
self.image_repo = image_repo
|
54
|
-
self.session = session
|
55
|
-
|
56
|
-
def build_and_upload_image(self) -> None:
|
57
|
-
"""Builds and uploads an image to the model registry.
|
58
|
-
|
59
|
-
Raises:
|
60
|
-
SnowflakeMLException: Occurs when failed to build image or push to image registry.
|
61
|
-
"""
|
62
|
-
|
63
|
-
def _setup_docker_config(docker_config_dir: str, registry_cred: str) -> None:
|
64
|
-
"""Set up a temporary docker config, which is used for running all docker commands. The format of config
|
65
|
-
is based on the format that is compatible with docker credential helper:
|
66
|
-
{
|
67
|
-
"auths": {
|
68
|
-
"https://index.docker.io/v1/": {
|
69
|
-
"auth": "<base_64_encoded_username_password>"
|
70
|
-
}
|
71
|
-
}
|
72
|
-
}
|
73
|
-
|
74
|
-
Docker will try to find cli-plugins in the config dir, and then fallback to
|
75
|
-
/usr/local/lib/docker/cli-plugins OR /usr/local/libexec/docker/cli-plugins
|
76
|
-
/usr/lib/docker/cli-plugins OR /usr/libexec/docker/cli-plugins
|
77
|
-
To prevent the case that none of them exists, we copy the cli-plugins in the current config directory,
|
78
|
-
which is defined in DOCKER_CONFIG and default to $HOME/.docker to our temp config dir.
|
79
|
-
|
80
|
-
Args:
|
81
|
-
docker_config_dir: Path to docker configuration directory, which stores the temporary session token.
|
82
|
-
registry_cred: image registry basic auth credential.
|
83
|
-
"""
|
84
|
-
orig_docker_config_dir = os.getenv("DOCKER_CONFIG", os.path.join(os.path.expanduser("~"), ".docker"))
|
85
|
-
if os.path.exists(os.path.join(orig_docker_config_dir, "cli-plugins")):
|
86
|
-
shutil.copytree(
|
87
|
-
os.path.join(orig_docker_config_dir, "cli-plugins"),
|
88
|
-
os.path.join(docker_config_dir, "cli-plugins"),
|
89
|
-
symlinks=True,
|
90
|
-
)
|
91
|
-
content = {"auths": {self.full_image_name: {"auth": registry_cred}}}
|
92
|
-
config_path = os.path.join(docker_config_dir, "config.json")
|
93
|
-
with open(config_path, "w", encoding="utf-8") as file:
|
94
|
-
json.dump(content, file)
|
95
|
-
|
96
|
-
def _cleanup_local_image(docker_config_dir: str) -> None:
|
97
|
-
try:
|
98
|
-
image_exist_command = ["docker", "image", "inspect", self.full_image_name]
|
99
|
-
self._run_docker_commands(image_exist_command)
|
100
|
-
except Exception:
|
101
|
-
# Image does not exist, probably due to failed build step
|
102
|
-
pass
|
103
|
-
else:
|
104
|
-
commands = ["docker", "--config", docker_config_dir, "rmi", self.full_image_name]
|
105
|
-
logger.debug(f"Removing local image: {self.full_image_name}")
|
106
|
-
self._run_docker_commands(commands)
|
107
|
-
|
108
|
-
self.validate_docker_client_env()
|
109
|
-
with credential.generate_image_registry_credential(
|
110
|
-
self.session
|
111
|
-
) as registry_cred, tempfile.TemporaryDirectory() as docker_config_dir:
|
112
|
-
try:
|
113
|
-
_setup_docker_config(docker_config_dir=docker_config_dir, registry_cred=registry_cred)
|
114
|
-
start = time.time()
|
115
|
-
self._build_and_tag(docker_config_dir)
|
116
|
-
end = time.time()
|
117
|
-
logger.info(f"Time taken to build the image on the client: {end - start:.2f} seconds")
|
118
|
-
|
119
|
-
except Exception as e:
|
120
|
-
raise snowml_exceptions.SnowflakeMLException(
|
121
|
-
error_code=error_codes.INTERNAL_DOCKER_ERROR,
|
122
|
-
original_exception=RuntimeError("Failed to build docker image."),
|
123
|
-
) from e
|
124
|
-
else:
|
125
|
-
try:
|
126
|
-
start = time.time()
|
127
|
-
self._upload(docker_config_dir)
|
128
|
-
end = time.time()
|
129
|
-
logger.info(f"Time taken to upload the image to image registry: {end - start:.2f} seconds")
|
130
|
-
except Exception as e:
|
131
|
-
raise snowml_exceptions.SnowflakeMLException(
|
132
|
-
error_code=error_codes.INTERNAL_DOCKER_ERROR,
|
133
|
-
original_exception=RuntimeError("Failed to upload docker image to registry."),
|
134
|
-
) from e
|
135
|
-
finally:
|
136
|
-
_cleanup_local_image(docker_config_dir)
|
137
|
-
|
138
|
-
def validate_docker_client_env(self) -> None:
|
139
|
-
"""Ensure docker client is running and BuildKit is enabled. Note that Buildx always uses BuildKit.
|
140
|
-
- Ensure docker daemon is running through the "docker info" command on shell. When docker daemon is running,
|
141
|
-
return code will be 0, else return code will be 1.
|
142
|
-
- Ensure BuildKit is enabled by checking "docker buildx version".
|
143
|
-
|
144
|
-
Raises:
|
145
|
-
SnowflakeMLException: Occurs when Docker is not installed or is not running.
|
146
|
-
|
147
|
-
"""
|
148
|
-
try:
|
149
|
-
self._run_docker_commands(["docker", "info"])
|
150
|
-
except Exception:
|
151
|
-
raise snowml_exceptions.SnowflakeMLException(
|
152
|
-
error_code=error_codes.CLIENT_DEPENDENCY_MISSING_ERROR,
|
153
|
-
original_exception=ConnectionError(
|
154
|
-
"Failed to initialize Docker client. Please ensure Docker is installed and running."
|
155
|
-
),
|
156
|
-
)
|
157
|
-
|
158
|
-
try:
|
159
|
-
self._run_docker_commands(["docker", "buildx", "version"])
|
160
|
-
except Exception:
|
161
|
-
raise snowml_exceptions.SnowflakeMLException(
|
162
|
-
error_code=error_codes.CLIENT_DEPENDENCY_MISSING_ERROR,
|
163
|
-
original_exception=ConnectionError(
|
164
|
-
"Please ensured Docker is installed with BuildKit by following "
|
165
|
-
"https://docs.docker.com/build/buildkit/#getting-started"
|
166
|
-
),
|
167
|
-
)
|
168
|
-
|
169
|
-
def _build_and_tag(self, docker_config_dir: str) -> None:
|
170
|
-
"""Constructs the Docker context directory and then builds a Docker image based on that context.
|
171
|
-
|
172
|
-
Args:
|
173
|
-
docker_config_dir: Path to docker configuration directory, which stores the temporary session token.
|
174
|
-
"""
|
175
|
-
self._build_image_from_context(docker_config_dir=docker_config_dir)
|
176
|
-
|
177
|
-
def _run_docker_commands(self, commands: List[str]) -> None:
|
178
|
-
"""Run docker commands in a new child process.
|
179
|
-
|
180
|
-
Args:
|
181
|
-
commands: List of commands to run.
|
182
|
-
|
183
|
-
Raises:
|
184
|
-
SnowflakeMLException: Occurs when docker commands failed to execute.
|
185
|
-
"""
|
186
|
-
proc = subprocess.Popen(
|
187
|
-
commands, cwd=os.getcwd(), stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, shell=False
|
188
|
-
)
|
189
|
-
output_lines = []
|
190
|
-
|
191
|
-
if proc.stdout:
|
192
|
-
for line in iter(proc.stdout.readline, ""):
|
193
|
-
output_lines.append(line)
|
194
|
-
logger.debug(line)
|
195
|
-
|
196
|
-
if proc.wait():
|
197
|
-
for line in output_lines:
|
198
|
-
logger.error(line)
|
199
|
-
|
200
|
-
raise snowml_exceptions.SnowflakeMLException(
|
201
|
-
error_code=error_codes.INTERNAL_DOCKER_ERROR,
|
202
|
-
original_exception=RuntimeError(f"Docker command failed: {' '.join(commands)}"),
|
203
|
-
)
|
204
|
-
|
205
|
-
def _build_image_from_context(self, docker_config_dir: str, *, platform: Platform = Platform.LINUX_AMD64) -> None:
|
206
|
-
"""Builds a Docker image based on provided context.
|
207
|
-
|
208
|
-
Args:
|
209
|
-
docker_config_dir: Path to docker configuration directory, which stores the temporary session token.
|
210
|
-
platform: Target platform for the build output, in the format "os[/arch[/variant]]".
|
211
|
-
"""
|
212
|
-
|
213
|
-
commands = [
|
214
|
-
"docker",
|
215
|
-
"--config",
|
216
|
-
docker_config_dir,
|
217
|
-
"buildx",
|
218
|
-
"build",
|
219
|
-
"--platform",
|
220
|
-
platform.value,
|
221
|
-
"--tag",
|
222
|
-
f"{self.full_image_name}",
|
223
|
-
self.context_dir,
|
224
|
-
]
|
225
|
-
|
226
|
-
self._run_docker_commands(commands)
|
227
|
-
|
228
|
-
def _upload(self, docker_config_dir: str) -> None:
|
229
|
-
"""
|
230
|
-
Uploads image to the image registry. This process requires a "docker login" followed by a "docker push". Remove
|
231
|
-
local image at the end of the upload operation to save up local space. Image cache is kept for more performant
|
232
|
-
built experience at the cost of small storage footprint.
|
233
|
-
|
234
|
-
By default, Docker overwrites the local Docker config file "/.docker/config.json" whenever a docker login
|
235
|
-
occurs. However, to ensure better isolation between Snowflake-managed Docker credentials and the user's own
|
236
|
-
Docker credentials, we will not use the default Docker config. Instead, we will write the username and session
|
237
|
-
token to a temporary file and use "docker --config" so that it only applies to the specific Docker command being
|
238
|
-
executed, without affecting the user's local Docker setup. The credential file will be automatically removed
|
239
|
-
at the end of upload operation.
|
240
|
-
|
241
|
-
Args:
|
242
|
-
docker_config_dir: Path to docker configuration directory, which stores the temporary session token.
|
243
|
-
"""
|
244
|
-
commands = ["docker", "--config", docker_config_dir, "login", self.full_image_name]
|
245
|
-
self._run_docker_commands(commands)
|
246
|
-
|
247
|
-
logger.debug(f"Pushing image to image repo {self.full_image_name}")
|
248
|
-
commands = ["docker", "--config", docker_config_dir, "push", self.full_image_name]
|
249
|
-
self._run_docker_commands(commands)
|
@@ -1,130 +0,0 @@
|
|
1
|
-
import os
|
2
|
-
import posixpath
|
3
|
-
import shutil
|
4
|
-
import string
|
5
|
-
from typing import Optional
|
6
|
-
|
7
|
-
import importlib_resources
|
8
|
-
|
9
|
-
from snowflake.ml._internal import file_utils
|
10
|
-
from snowflake.ml._internal.utils import identifier
|
11
|
-
from snowflake.ml.model._deploy_client import image_builds
|
12
|
-
from snowflake.ml.model._deploy_client.utils import constants
|
13
|
-
from snowflake.ml.model._packager.model_meta import model_meta
|
14
|
-
from snowflake.snowpark import FileOperation, Session
|
15
|
-
|
16
|
-
|
17
|
-
class DockerContext:
|
18
|
-
"""
|
19
|
-
Constructs the Docker context directory required for image building.
|
20
|
-
"""
|
21
|
-
|
22
|
-
def __init__(
|
23
|
-
self,
|
24
|
-
context_dir: str,
|
25
|
-
model_meta: model_meta.ModelMetadata,
|
26
|
-
session: Optional[Session] = None,
|
27
|
-
model_zip_stage_path: Optional[str] = None,
|
28
|
-
) -> None:
|
29
|
-
"""Initialization
|
30
|
-
|
31
|
-
Args:
|
32
|
-
context_dir: Path to context directory.
|
33
|
-
model_meta: Model Metadata.
|
34
|
-
session: Snowpark session.
|
35
|
-
model_zip_stage_path: Path to model zip file on stage.
|
36
|
-
"""
|
37
|
-
self.context_dir = context_dir
|
38
|
-
self.model_meta = model_meta
|
39
|
-
assert (session is None) == (model_zip_stage_path is None)
|
40
|
-
self.session = session
|
41
|
-
self.model_zip_stage_path = model_zip_stage_path
|
42
|
-
|
43
|
-
def build(self) -> None:
|
44
|
-
"""
|
45
|
-
Generates and/or moves resources into the Docker context directory.Rename the random model directory name to
|
46
|
-
constant "model_dir" instead for better readability.
|
47
|
-
"""
|
48
|
-
self._generate_inference_code()
|
49
|
-
self._copy_entrypoint_script_to_docker_context()
|
50
|
-
self._copy_model_env_dependency_to_docker_context()
|
51
|
-
self._generate_docker_file()
|
52
|
-
|
53
|
-
def _copy_entrypoint_script_to_docker_context(self) -> None:
|
54
|
-
"""Copy gunicorn_run.sh entrypoint to docker context directory."""
|
55
|
-
script_path = importlib_resources.files(image_builds).joinpath(constants.ENTRYPOINT_SCRIPT)
|
56
|
-
target_path = os.path.join(self.context_dir, constants.ENTRYPOINT_SCRIPT)
|
57
|
-
|
58
|
-
with open(script_path, encoding="utf-8") as source_file, file_utils.open_file(target_path, "w") as target_file:
|
59
|
-
target_file.write(source_file.read())
|
60
|
-
|
61
|
-
def _copy_model_env_dependency_to_docker_context(self) -> None:
|
62
|
-
"""
|
63
|
-
Convert model dependencies to files from model metadata.
|
64
|
-
"""
|
65
|
-
self.model_meta.save(self.context_dir)
|
66
|
-
|
67
|
-
def _generate_docker_file(self) -> None:
|
68
|
-
"""
|
69
|
-
Generates dockerfile based on dockerfile template.
|
70
|
-
"""
|
71
|
-
docker_file_path = os.path.join(self.context_dir, "Dockerfile")
|
72
|
-
docker_file_template = (
|
73
|
-
importlib_resources.files(image_builds).joinpath("templates/dockerfile_template").read_text("utf-8")
|
74
|
-
)
|
75
|
-
|
76
|
-
if self.model_zip_stage_path is not None:
|
77
|
-
norm_stage_path = posixpath.normpath(identifier.remove_prefix(self.model_zip_stage_path, "@"))
|
78
|
-
assert self.session
|
79
|
-
fop = FileOperation(self.session)
|
80
|
-
# The explicit download here is inefficient but a compromise.
|
81
|
-
# We could in theory reuse the download needed for metadata extraction, but it's hacky and will go away.
|
82
|
-
# Ideally, the model download should happen as part of the server side image build,
|
83
|
-
# but it requires have our own image builder since there's need to be logic downloading model
|
84
|
-
# into the context directory.
|
85
|
-
get_res_list = fop.get(stage_location=self.model_zip_stage_path, target_directory=self.context_dir)
|
86
|
-
assert len(get_res_list) == 1, f"Single zip file should be returned, but got {len(get_res_list)} files."
|
87
|
-
local_zip_file_path = os.path.basename(get_res_list[0].file)
|
88
|
-
copy_model_statement = f"COPY {local_zip_file_path} ./{norm_stage_path}"
|
89
|
-
extra_env_statement = f"ENV MODEL_ZIP_STAGE_PATH={norm_stage_path}"
|
90
|
-
else:
|
91
|
-
copy_model_statement = ""
|
92
|
-
extra_env_statement = ""
|
93
|
-
|
94
|
-
with open(docker_file_path, "w", encoding="utf-8") as dockerfile:
|
95
|
-
base_image = "mambaorg/micromamba:1.4.3"
|
96
|
-
tag = base_image.split(":")[1]
|
97
|
-
assert tag != constants.LATEST_IMAGE_TAG, (
|
98
|
-
"Base image tag should not be 'latest' as it might cause false" "positive image cache hit"
|
99
|
-
)
|
100
|
-
dockerfile_content = string.Template(docker_file_template).safe_substitute(
|
101
|
-
{
|
102
|
-
"base_image": "mambaorg/micromamba:1.4.3",
|
103
|
-
"model_env_folder": constants.MODEL_ENV_FOLDER,
|
104
|
-
"inference_server_dir": constants.INFERENCE_SERVER_DIR,
|
105
|
-
"entrypoint_script": constants.ENTRYPOINT_SCRIPT,
|
106
|
-
# Instead of omitting this ENV var when no CUDA required, we explicitly set it to empty to override
|
107
|
-
# as no CUDA is detected thus it won't be affected by the existence of CUDA in base image.
|
108
|
-
# https://conda.io/projects/conda/en/latest/user-guide/tasks/manage-virtual.html
|
109
|
-
"cuda_override_env": self.model_meta.env.cuda_version if self.model_meta.env.cuda_version else "",
|
110
|
-
"copy_model_statement": copy_model_statement,
|
111
|
-
"extra_env_statement": extra_env_statement,
|
112
|
-
}
|
113
|
-
)
|
114
|
-
dockerfile.write(dockerfile_content)
|
115
|
-
|
116
|
-
def _generate_inference_code(self) -> None:
|
117
|
-
"""
|
118
|
-
Generates inference code based on the app template and creates a folder named 'server' to house the inference
|
119
|
-
server code.
|
120
|
-
"""
|
121
|
-
with importlib_resources.as_file(
|
122
|
-
importlib_resources.files(image_builds).joinpath(constants.INFERENCE_SERVER_DIR)
|
123
|
-
) as inference_server_folder_path:
|
124
|
-
destination_folder_path = os.path.join(self.context_dir, constants.INFERENCE_SERVER_DIR)
|
125
|
-
ignore_patterns = shutil.ignore_patterns("BUILD.bazel", "*test.py", "*.\\.*", "__pycache__")
|
126
|
-
file_utils.copytree(
|
127
|
-
inference_server_folder_path,
|
128
|
-
destination_folder_path,
|
129
|
-
ignore=ignore_patterns,
|
130
|
-
)
|
@@ -1,36 +0,0 @@
|
|
1
|
-
#!/bin/bash
|
2
|
-
set -eu
|
3
|
-
|
4
|
-
OS=$(uname)
|
5
|
-
|
6
|
-
if [[ ${OS} = "Linux" ]]; then
|
7
|
-
NUM_CORES=$(nproc)
|
8
|
-
elif [[ ${OS} = "Darwin" ]]; then
|
9
|
-
# macOS
|
10
|
-
NUM_CORES=$(sysctl -n hw.ncpu)
|
11
|
-
elif [[ ${OS} = "Windows" ]]; then
|
12
|
-
NUM_CORES=$(wmic cpu get NumberOfCores | grep -Eo '[0-9]+')
|
13
|
-
else
|
14
|
-
echo "Unsupported operating system: ${OS}"
|
15
|
-
exit 1
|
16
|
-
fi
|
17
|
-
|
18
|
-
# Check if the "NUM_WORKERS" variable is set by the user
|
19
|
-
if [[ -n "${NUM_WORKERS-}" && "${NUM_WORKERS}" != "None" ]]; then
|
20
|
-
# If the user has set the "num_workers" variable, use it to overwrite the default value
|
21
|
-
FINAL_NUM_WORKERS=${NUM_WORKERS}
|
22
|
-
else
|
23
|
-
# Based on the Gunicorn documentation, set the number of workers to number_of_cores * 2 + 1. This assumption is
|
24
|
-
# based on an ideal scenario where one core is handling two processes simultaneously, while one process is dedicated to
|
25
|
-
# IO operations and the other process is performing compute tasks.
|
26
|
-
# However, in case when the model is large, we will run into OOM error as each process will need to load the model
|
27
|
-
# into memory. In such cases, we require the user to pass in "num_workers" to overwrite the default.
|
28
|
-
FINAL_NUM_WORKERS=$((NUM_CORES * 2 + 1))
|
29
|
-
fi
|
30
|
-
|
31
|
-
echo "Number of CPU cores: $NUM_CORES"
|
32
|
-
echo "Setting number of workers to $FINAL_NUM_WORKERS"
|
33
|
-
|
34
|
-
# Exclude preload option as it won't work with non-thread-safe model, and no easy way to detect whether model is
|
35
|
-
# thread-safe or not. Defer the optimization later.
|
36
|
-
exec /opt/conda/bin/gunicorn -w "$FINAL_NUM_WORKERS" -k uvicorn.workers.UvicornWorker -b 0.0.0.0:5000 --timeout 600 inference_server.main:app
|
@@ -1,268 +0,0 @@
|
|
1
|
-
import asyncio
|
2
|
-
import http
|
3
|
-
import logging
|
4
|
-
import os
|
5
|
-
import sys
|
6
|
-
import tempfile
|
7
|
-
import threading
|
8
|
-
import time
|
9
|
-
import traceback
|
10
|
-
import zipfile
|
11
|
-
from enum import Enum
|
12
|
-
from typing import Dict, List, Optional, cast
|
13
|
-
|
14
|
-
import pandas as pd
|
15
|
-
from gunicorn import arbiter
|
16
|
-
from starlette import applications, concurrency, requests, responses, routing
|
17
|
-
|
18
|
-
|
19
|
-
class _ModelLoadingState(Enum):
|
20
|
-
"""
|
21
|
-
Enum class to represent various model loading state.
|
22
|
-
"""
|
23
|
-
|
24
|
-
LOADING = "loading"
|
25
|
-
SUCCEEDED = "succeeded"
|
26
|
-
FAILED = "failed"
|
27
|
-
|
28
|
-
|
29
|
-
class CustomThread(threading.Thread):
|
30
|
-
"""
|
31
|
-
Custom Thread implementation that overrides Thread.run.
|
32
|
-
|
33
|
-
This is necessary because the default Thread implementation suppresses exceptions in child threads. The standard
|
34
|
-
behavior involves the Thread class catching exceptions and throwing a SystemExit exception, which requires
|
35
|
-
Thread.join to terminate the process. To address this, we overwrite Thread.run and use os._exit instead.
|
36
|
-
|
37
|
-
We throw specific error code "Arbiter.APP_LOAD_ERROR" such that Gunicorn Arbiter master process will be killed,
|
38
|
-
which then trigger the container to be marked as failed. This ensures the container becomes ready when all workers
|
39
|
-
loaded the model successfully.
|
40
|
-
"""
|
41
|
-
|
42
|
-
def run(self) -> None:
|
43
|
-
try:
|
44
|
-
super().run()
|
45
|
-
# Keep the daemon thread alive to avoid destroy
|
46
|
-
# state attached to thread when loading the model.
|
47
|
-
while True:
|
48
|
-
time.sleep(60)
|
49
|
-
except Exception:
|
50
|
-
logger.error(traceback.format_exc())
|
51
|
-
os._exit(arbiter.Arbiter.APP_LOAD_ERROR)
|
52
|
-
|
53
|
-
|
54
|
-
logger = logging.getLogger(__name__)
|
55
|
-
_LOADED_MODEL = None
|
56
|
-
_LOADED_META = None
|
57
|
-
_MODEL_CODE_DIR = "code"
|
58
|
-
_MODEL_LOADING_STATE = _ModelLoadingState.LOADING
|
59
|
-
_MODEL_LOADING_EVENT = threading.Event()
|
60
|
-
_CONCURRENT_REQUESTS_MAX: Optional[int] = None
|
61
|
-
_CONCURRENT_COUNTER = 0
|
62
|
-
_CONCURRENT_COUNTER_LOCK = asyncio.Lock()
|
63
|
-
TARGET_METHOD = None
|
64
|
-
|
65
|
-
|
66
|
-
def _run_setup() -> None:
|
67
|
-
"""Set up logging and load model into memory."""
|
68
|
-
# Align the application logger's handler with Gunicorn's to capture logs from all processes.
|
69
|
-
gunicorn_logger = logging.getLogger("gunicorn.error")
|
70
|
-
logger.handlers = gunicorn_logger.handlers
|
71
|
-
logger.setLevel(gunicorn_logger.level)
|
72
|
-
|
73
|
-
logger.info(f"ENV: {os.environ}")
|
74
|
-
|
75
|
-
global _LOADED_MODEL
|
76
|
-
global _LOADED_META
|
77
|
-
global _MODEL_LOADING_STATE
|
78
|
-
global _MODEL_LOADING_EVENT
|
79
|
-
global _CONCURRENT_REQUESTS_MAX
|
80
|
-
global TARGET_METHOD
|
81
|
-
|
82
|
-
try:
|
83
|
-
model_zip_stage_path = os.getenv("MODEL_ZIP_STAGE_PATH")
|
84
|
-
assert model_zip_stage_path, "Missing environment variable MODEL_ZIP_STAGE_PATH"
|
85
|
-
|
86
|
-
TARGET_METHOD = os.getenv("TARGET_METHOD")
|
87
|
-
|
88
|
-
_concurrent_requests_max_env = os.getenv("_CONCURRENT_REQUESTS_MAX", "1")
|
89
|
-
_CONCURRENT_REQUESTS_MAX = int(_concurrent_requests_max_env)
|
90
|
-
|
91
|
-
with tempfile.TemporaryDirectory() as tmp_dir:
|
92
|
-
if zipfile.is_zipfile(model_zip_stage_path):
|
93
|
-
extracted_dir = os.path.join(tmp_dir, "extracted_model_dir")
|
94
|
-
logger.info(f"Extracting model zip from {model_zip_stage_path} to {extracted_dir}")
|
95
|
-
with zipfile.ZipFile(model_zip_stage_path, "r") as model_zip:
|
96
|
-
if len(model_zip.namelist()) > 1:
|
97
|
-
model_zip.extractall(extracted_dir)
|
98
|
-
else:
|
99
|
-
raise RuntimeError(f"No model zip found at stage path: {model_zip_stage_path}")
|
100
|
-
logger.info(f"Loading model from {extracted_dir} into memory")
|
101
|
-
|
102
|
-
sys.path.insert(0, os.path.join(extracted_dir, _MODEL_CODE_DIR))
|
103
|
-
|
104
|
-
# TODO (Server-side Model Rollout):
|
105
|
-
# Keep try block only
|
106
|
-
# SPCS spec will convert all environment variables as strings.
|
107
|
-
use_gpu = os.environ.get("SNOWML_USE_GPU", "False").lower() == "true"
|
108
|
-
try:
|
109
|
-
from snowflake.ml.model._packager import model_packager
|
110
|
-
|
111
|
-
pk = model_packager.ModelPackager(extracted_dir)
|
112
|
-
pk.load(
|
113
|
-
as_custom_model=True,
|
114
|
-
meta_only=False,
|
115
|
-
options={"use_gpu": use_gpu},
|
116
|
-
)
|
117
|
-
_LOADED_MODEL = pk.model
|
118
|
-
_LOADED_META = pk.meta
|
119
|
-
except ImportError as e:
|
120
|
-
if e.name and not e.name.startswith("snowflake.ml"):
|
121
|
-
raise e
|
122
|
-
# Legacy model support
|
123
|
-
from snowflake.ml.model import ( # type: ignore[attr-defined]
|
124
|
-
_model as model_api,
|
125
|
-
)
|
126
|
-
|
127
|
-
if hasattr(model_api, "_load_model_for_deploy"):
|
128
|
-
_LOADED_MODEL, _LOADED_META = model_api._load_model_for_deploy(extracted_dir)
|
129
|
-
else:
|
130
|
-
_LOADED_MODEL, meta_LOADED_META = model_api._load(
|
131
|
-
local_dir_path=extracted_dir,
|
132
|
-
as_custom_model=True,
|
133
|
-
options={"use_gpu": use_gpu},
|
134
|
-
)
|
135
|
-
_MODEL_LOADING_STATE = _ModelLoadingState.SUCCEEDED
|
136
|
-
logger.info("Successfully loaded model into memory")
|
137
|
-
_MODEL_LOADING_EVENT.set()
|
138
|
-
except Exception as e:
|
139
|
-
_MODEL_LOADING_STATE = _ModelLoadingState.FAILED
|
140
|
-
raise e
|
141
|
-
|
142
|
-
|
143
|
-
async def ready(request: requests.Request) -> responses.JSONResponse:
|
144
|
-
"""Check if the application is ready to serve requests.
|
145
|
-
|
146
|
-
This endpoint is used to determine the readiness of the application to handle incoming requests. It returns an HTTP
|
147
|
-
200 status code only when the model has been successfully loaded into memory. If the model has not yet been loaded,
|
148
|
-
it responds with an HTTP 503 status code, which signals to the readiness probe to continue probing until the
|
149
|
-
application becomes ready or until the client's timeout is reached.
|
150
|
-
|
151
|
-
Args:
|
152
|
-
request:
|
153
|
-
The HTTP request object.
|
154
|
-
|
155
|
-
Returns:
|
156
|
-
A JSON response with status information:
|
157
|
-
- HTTP 200 status code and {"status": "ready"} when the model is loaded and the application is ready.
|
158
|
-
- HTTP 503 status code and {"status": "not ready"} when the model is not yet loaded.
|
159
|
-
|
160
|
-
"""
|
161
|
-
if _MODEL_LOADING_STATE == _ModelLoadingState.SUCCEEDED:
|
162
|
-
return responses.JSONResponse({"status": "ready"})
|
163
|
-
return responses.JSONResponse({"status": "not ready"}, status_code=http.HTTPStatus.SERVICE_UNAVAILABLE)
|
164
|
-
|
165
|
-
|
166
|
-
def _do_predict(input_json: Dict[str, List[List[object]]]) -> responses.JSONResponse:
|
167
|
-
from snowflake.ml.model.model_signature import FeatureSpec
|
168
|
-
|
169
|
-
assert _LOADED_MODEL, "model is not loaded"
|
170
|
-
assert _LOADED_META, "model metadata is not loaded"
|
171
|
-
assert TARGET_METHOD, "Missing environment variable TARGET_METHOD"
|
172
|
-
|
173
|
-
try:
|
174
|
-
features = cast(List[FeatureSpec], _LOADED_META.signatures[TARGET_METHOD].inputs)
|
175
|
-
dtype_map = {feature.name: feature.as_dtype() for feature in features}
|
176
|
-
input_cols = [spec.name for spec in features]
|
177
|
-
output_cols = [spec.name for spec in _LOADED_META.signatures[TARGET_METHOD].outputs]
|
178
|
-
assert "data" in input_json, "missing data field in the request input"
|
179
|
-
# The expression x[1:] is used to exclude the index of the data row.
|
180
|
-
input_data = [x[1] for x in input_json["data"]]
|
181
|
-
df = pd.json_normalize(input_data).astype(dtype=dtype_map)
|
182
|
-
x = df[input_cols]
|
183
|
-
assert len(input_data) != 0 and not all(not row for row in input_data), "empty data"
|
184
|
-
except Exception as e:
|
185
|
-
error_message = f"Input data malformed: {str(e)}\n{traceback.format_exc()}"
|
186
|
-
logger.error(f"Failed request with error: {error_message}")
|
187
|
-
return responses.JSONResponse({"error": error_message}, status_code=http.HTTPStatus.BAD_REQUEST)
|
188
|
-
|
189
|
-
try:
|
190
|
-
predictions_df = getattr(_LOADED_MODEL, TARGET_METHOD)(x)
|
191
|
-
predictions_df.columns = output_cols
|
192
|
-
# Use _ID to keep the order of prediction result and associated features.
|
193
|
-
_KEEP_ORDER_COL_NAME = "_ID"
|
194
|
-
if _KEEP_ORDER_COL_NAME in df.columns:
|
195
|
-
predictions_df[_KEEP_ORDER_COL_NAME] = df[_KEEP_ORDER_COL_NAME]
|
196
|
-
response = {"data": [[i, row] for i, row in enumerate(predictions_df.to_dict(orient="records"))]}
|
197
|
-
return responses.JSONResponse(response)
|
198
|
-
except Exception as e:
|
199
|
-
error_message = f"Prediction failed: {str(e)}\n{traceback.format_exc()}"
|
200
|
-
logger.error(f"Failed request with error: {error_message}")
|
201
|
-
return responses.JSONResponse({"error": error_message}, status_code=http.HTTPStatus.BAD_REQUEST)
|
202
|
-
|
203
|
-
|
204
|
-
async def predict(request: requests.Request) -> responses.JSONResponse:
|
205
|
-
"""Endpoint to make predictions based on input data.
|
206
|
-
|
207
|
-
Args:
|
208
|
-
request: The input data is expected to be in the following JSON format:
|
209
|
-
{
|
210
|
-
"data": [
|
211
|
-
[0, {'_ID': 0, 'input_feature_0': 0.0, 'input_feature_1': 1.0}],
|
212
|
-
[1, {'_ID': 1, 'input_feature_0': 2.0, 'input_feature_1': 3.0}],
|
213
|
-
}
|
214
|
-
Each row is represented as a list, where the first element denotes the index of the row.
|
215
|
-
|
216
|
-
Returns:
|
217
|
-
Two possible responses:
|
218
|
-
For success, return a JSON response
|
219
|
-
{
|
220
|
-
"data": [
|
221
|
-
[0, {'_ID': 0, 'output': 1}],
|
222
|
-
[1, {'_ID': 1, 'output': 2}]
|
223
|
-
]
|
224
|
-
},
|
225
|
-
The first element of each resulting list denotes the index of the row, and the rest of the elements
|
226
|
-
represent the prediction results for that row.
|
227
|
-
For an error, return {"error": error_message, "status_code": http_response_status_code}.
|
228
|
-
"""
|
229
|
-
_MODEL_LOADING_EVENT.wait() # Ensure model is indeed loaded into memory
|
230
|
-
|
231
|
-
global _CONCURRENT_COUNTER
|
232
|
-
global _CONCURRENT_COUNTER_LOCK
|
233
|
-
|
234
|
-
input_json = await request.json()
|
235
|
-
|
236
|
-
if _CONCURRENT_REQUESTS_MAX:
|
237
|
-
async with _CONCURRENT_COUNTER_LOCK:
|
238
|
-
if _CONCURRENT_COUNTER >= int(_CONCURRENT_REQUESTS_MAX):
|
239
|
-
return responses.JSONResponse(
|
240
|
-
{"error": "Too many requests"}, status_code=http.HTTPStatus.TOO_MANY_REQUESTS
|
241
|
-
)
|
242
|
-
|
243
|
-
async with _CONCURRENT_COUNTER_LOCK:
|
244
|
-
_CONCURRENT_COUNTER += 1
|
245
|
-
|
246
|
-
resp = await concurrency.run_in_threadpool(_do_predict, input_json)
|
247
|
-
|
248
|
-
async with _CONCURRENT_COUNTER_LOCK:
|
249
|
-
_CONCURRENT_COUNTER -= 1
|
250
|
-
|
251
|
-
return resp
|
252
|
-
|
253
|
-
|
254
|
-
def run_app() -> applications.Starlette:
|
255
|
-
# TODO[shchen]: SNOW-893654. Before SnowService supports Startup probe, or extends support for Readiness probe
|
256
|
-
# with configurable failureThreshold, we will have to load the model in a separate thread in order to prevent
|
257
|
-
# gunicorn worker timeout.
|
258
|
-
model_loading_worker = CustomThread(target=_run_setup, daemon=True)
|
259
|
-
model_loading_worker.start()
|
260
|
-
|
261
|
-
routes = [
|
262
|
-
routing.Route("/health", endpoint=ready, methods=["GET"]),
|
263
|
-
routing.Route("/predict", endpoint=predict, methods=["POST"]),
|
264
|
-
]
|
265
|
-
return applications.Starlette(routes=routes)
|
266
|
-
|
267
|
-
|
268
|
-
app = run_app()
|