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.
Files changed (284) hide show
  1. snowflake/cortex/__init__.py +4 -0
  2. snowflake/cortex/_classify_text.py +2 -2
  3. snowflake/cortex/_embed_text_1024.py +37 -0
  4. snowflake/cortex/_embed_text_768.py +37 -0
  5. snowflake/cortex/_extract_answer.py +2 -2
  6. snowflake/cortex/_sentiment.py +2 -2
  7. snowflake/cortex/_summarize.py +2 -2
  8. snowflake/cortex/_translate.py +2 -2
  9. snowflake/cortex/_util.py +4 -4
  10. snowflake/ml/_internal/env_utils.py +5 -5
  11. snowflake/ml/_internal/exceptions/error_codes.py +2 -0
  12. snowflake/ml/_internal/telemetry.py +142 -20
  13. snowflake/ml/_internal/utils/db_utils.py +50 -0
  14. snowflake/ml/_internal/utils/identifier.py +48 -11
  15. snowflake/ml/_internal/utils/service_logger.py +63 -0
  16. snowflake/ml/_internal/utils/snowflake_env.py +23 -13
  17. snowflake/ml/_internal/utils/sql_identifier.py +26 -2
  18. snowflake/ml/_internal/utils/table_manager.py +19 -1
  19. snowflake/ml/data/_internal/arrow_ingestor.py +1 -11
  20. snowflake/ml/data/data_connector.py +33 -7
  21. snowflake/ml/data/ingestor_utils.py +20 -10
  22. snowflake/ml/data/torch_utils.py +68 -0
  23. snowflake/ml/dataset/dataset.py +1 -3
  24. snowflake/ml/feature_store/access_manager.py +3 -3
  25. snowflake/ml/feature_store/feature_store.py +60 -19
  26. snowflake/ml/feature_store/feature_view.py +84 -30
  27. snowflake/ml/fileset/embedded_stage_fs.py +1 -1
  28. snowflake/ml/fileset/fileset.py +1 -1
  29. snowflake/ml/fileset/sfcfs.py +9 -3
  30. snowflake/ml/fileset/stage_fs.py +2 -1
  31. snowflake/ml/lineage/lineage_node.py +7 -2
  32. snowflake/ml/model/__init__.py +1 -2
  33. snowflake/ml/model/_client/model/model_version_impl.py +96 -12
  34. snowflake/ml/model/_client/ops/model_ops.py +124 -6
  35. snowflake/ml/model/_client/ops/service_ops.py +309 -9
  36. snowflake/ml/model/_client/service/model_deployment_spec.py +8 -5
  37. snowflake/ml/model/_client/service/model_deployment_spec_schema.py +2 -2
  38. snowflake/ml/model/_client/sql/_base.py +5 -0
  39. snowflake/ml/model/_client/sql/model.py +1 -0
  40. snowflake/ml/model/_client/sql/model_version.py +9 -5
  41. snowflake/ml/model/_client/sql/service.py +121 -20
  42. snowflake/ml/model/_model_composer/model_composer.py +11 -39
  43. snowflake/ml/model/_model_composer/model_manifest/model_manifest.py +31 -11
  44. snowflake/ml/model/_packager/model_env/model_env.py +4 -38
  45. snowflake/ml/model/_packager/model_handlers/_utils.py +134 -28
  46. snowflake/ml/model/_packager/model_handlers/catboost.py +31 -30
  47. snowflake/ml/model/_packager/model_handlers/huggingface_pipeline.py +26 -18
  48. snowflake/ml/model/_packager/model_handlers/lightgbm.py +31 -58
  49. snowflake/ml/model/_packager/model_handlers/mlflow.py +3 -5
  50. snowflake/ml/model/_packager/model_handlers/model_objective_utils.py +169 -0
  51. snowflake/ml/model/_packager/model_handlers/sentence_transformers.py +15 -8
  52. snowflake/ml/model/_packager/model_handlers/sklearn.py +56 -60
  53. snowflake/ml/model/_packager/model_handlers/snowmlmodel.py +141 -9
  54. snowflake/ml/model/_packager/model_handlers/torchscript.py +2 -2
  55. snowflake/ml/model/_packager/model_handlers/xgboost.py +63 -48
  56. snowflake/ml/model/_packager/model_meta/model_meta.py +16 -42
  57. snowflake/ml/model/_packager/model_meta/model_meta_schema.py +1 -14
  58. snowflake/ml/model/_packager/model_packager.py +14 -8
  59. snowflake/ml/model/_packager/model_runtime/model_runtime.py +11 -0
  60. snowflake/ml/model/_signatures/pytorch_handler.py +1 -1
  61. snowflake/ml/model/_signatures/snowpark_handler.py +3 -2
  62. snowflake/ml/model/_signatures/utils.py +9 -0
  63. snowflake/ml/model/type_hints.py +12 -145
  64. snowflake/ml/modeling/_internal/constants.py +1 -0
  65. snowflake/ml/modeling/_internal/local_implementations/pandas_handlers.py +5 -5
  66. snowflake/ml/modeling/_internal/local_implementations/pandas_trainer.py +9 -6
  67. snowflake/ml/modeling/_internal/model_specifications.py +2 -0
  68. snowflake/ml/modeling/_internal/model_trainer.py +1 -0
  69. snowflake/ml/modeling/_internal/snowpark_implementations/distributed_hpo_trainer.py +2 -4
  70. snowflake/ml/modeling/_internal/snowpark_implementations/snowpark_handlers.py +5 -5
  71. snowflake/ml/modeling/_internal/snowpark_implementations/snowpark_trainer.py +130 -166
  72. snowflake/ml/modeling/_internal/snowpark_implementations/xgboost_external_memory_trainer.py +0 -1
  73. snowflake/ml/modeling/calibration/calibrated_classifier_cv.py +61 -21
  74. snowflake/ml/modeling/cluster/affinity_propagation.py +61 -21
  75. snowflake/ml/modeling/cluster/agglomerative_clustering.py +61 -21
  76. snowflake/ml/modeling/cluster/birch.py +61 -21
  77. snowflake/ml/modeling/cluster/bisecting_k_means.py +61 -21
  78. snowflake/ml/modeling/cluster/dbscan.py +61 -21
  79. snowflake/ml/modeling/cluster/feature_agglomeration.py +61 -21
  80. snowflake/ml/modeling/cluster/k_means.py +61 -21
  81. snowflake/ml/modeling/cluster/mean_shift.py +61 -21
  82. snowflake/ml/modeling/cluster/mini_batch_k_means.py +61 -21
  83. snowflake/ml/modeling/cluster/optics.py +61 -21
  84. snowflake/ml/modeling/cluster/spectral_biclustering.py +61 -21
  85. snowflake/ml/modeling/cluster/spectral_clustering.py +61 -21
  86. snowflake/ml/modeling/cluster/spectral_coclustering.py +61 -21
  87. snowflake/ml/modeling/compose/column_transformer.py +61 -21
  88. snowflake/ml/modeling/compose/transformed_target_regressor.py +61 -21
  89. snowflake/ml/modeling/covariance/elliptic_envelope.py +61 -21
  90. snowflake/ml/modeling/covariance/empirical_covariance.py +61 -21
  91. snowflake/ml/modeling/covariance/graphical_lasso.py +61 -21
  92. snowflake/ml/modeling/covariance/graphical_lasso_cv.py +61 -21
  93. snowflake/ml/modeling/covariance/ledoit_wolf.py +61 -21
  94. snowflake/ml/modeling/covariance/min_cov_det.py +61 -21
  95. snowflake/ml/modeling/covariance/oas.py +61 -21
  96. snowflake/ml/modeling/covariance/shrunk_covariance.py +61 -21
  97. snowflake/ml/modeling/decomposition/dictionary_learning.py +61 -21
  98. snowflake/ml/modeling/decomposition/factor_analysis.py +61 -21
  99. snowflake/ml/modeling/decomposition/fast_ica.py +61 -21
  100. snowflake/ml/modeling/decomposition/incremental_pca.py +61 -21
  101. snowflake/ml/modeling/decomposition/kernel_pca.py +61 -21
  102. snowflake/ml/modeling/decomposition/mini_batch_dictionary_learning.py +61 -21
  103. snowflake/ml/modeling/decomposition/mini_batch_sparse_pca.py +61 -21
  104. snowflake/ml/modeling/decomposition/pca.py +61 -21
  105. snowflake/ml/modeling/decomposition/sparse_pca.py +61 -21
  106. snowflake/ml/modeling/decomposition/truncated_svd.py +61 -21
  107. snowflake/ml/modeling/discriminant_analysis/linear_discriminant_analysis.py +61 -21
  108. snowflake/ml/modeling/discriminant_analysis/quadratic_discriminant_analysis.py +61 -21
  109. snowflake/ml/modeling/ensemble/ada_boost_classifier.py +61 -21
  110. snowflake/ml/modeling/ensemble/ada_boost_regressor.py +61 -21
  111. snowflake/ml/modeling/ensemble/bagging_classifier.py +61 -21
  112. snowflake/ml/modeling/ensemble/bagging_regressor.py +61 -21
  113. snowflake/ml/modeling/ensemble/extra_trees_classifier.py +61 -21
  114. snowflake/ml/modeling/ensemble/extra_trees_regressor.py +61 -21
  115. snowflake/ml/modeling/ensemble/gradient_boosting_classifier.py +61 -21
  116. snowflake/ml/modeling/ensemble/gradient_boosting_regressor.py +61 -21
  117. snowflake/ml/modeling/ensemble/hist_gradient_boosting_classifier.py +61 -21
  118. snowflake/ml/modeling/ensemble/hist_gradient_boosting_regressor.py +61 -21
  119. snowflake/ml/modeling/ensemble/isolation_forest.py +61 -21
  120. snowflake/ml/modeling/ensemble/random_forest_classifier.py +61 -21
  121. snowflake/ml/modeling/ensemble/random_forest_regressor.py +61 -21
  122. snowflake/ml/modeling/ensemble/stacking_regressor.py +61 -21
  123. snowflake/ml/modeling/ensemble/voting_classifier.py +61 -21
  124. snowflake/ml/modeling/ensemble/voting_regressor.py +61 -21
  125. snowflake/ml/modeling/feature_selection/generic_univariate_select.py +61 -21
  126. snowflake/ml/modeling/feature_selection/select_fdr.py +61 -21
  127. snowflake/ml/modeling/feature_selection/select_fpr.py +61 -21
  128. snowflake/ml/modeling/feature_selection/select_fwe.py +61 -21
  129. snowflake/ml/modeling/feature_selection/select_k_best.py +61 -21
  130. snowflake/ml/modeling/feature_selection/select_percentile.py +61 -21
  131. snowflake/ml/modeling/feature_selection/sequential_feature_selector.py +61 -21
  132. snowflake/ml/modeling/feature_selection/variance_threshold.py +61 -21
  133. snowflake/ml/modeling/gaussian_process/gaussian_process_classifier.py +61 -21
  134. snowflake/ml/modeling/gaussian_process/gaussian_process_regressor.py +61 -21
  135. snowflake/ml/modeling/impute/iterative_imputer.py +61 -21
  136. snowflake/ml/modeling/impute/knn_imputer.py +61 -21
  137. snowflake/ml/modeling/impute/missing_indicator.py +61 -21
  138. snowflake/ml/modeling/kernel_approximation/additive_chi2_sampler.py +61 -21
  139. snowflake/ml/modeling/kernel_approximation/nystroem.py +61 -21
  140. snowflake/ml/modeling/kernel_approximation/polynomial_count_sketch.py +61 -21
  141. snowflake/ml/modeling/kernel_approximation/rbf_sampler.py +61 -21
  142. snowflake/ml/modeling/kernel_approximation/skewed_chi2_sampler.py +61 -21
  143. snowflake/ml/modeling/kernel_ridge/kernel_ridge.py +61 -21
  144. snowflake/ml/modeling/lightgbm/lgbm_classifier.py +61 -21
  145. snowflake/ml/modeling/lightgbm/lgbm_regressor.py +61 -21
  146. snowflake/ml/modeling/linear_model/ard_regression.py +61 -21
  147. snowflake/ml/modeling/linear_model/bayesian_ridge.py +61 -21
  148. snowflake/ml/modeling/linear_model/elastic_net.py +61 -21
  149. snowflake/ml/modeling/linear_model/elastic_net_cv.py +61 -21
  150. snowflake/ml/modeling/linear_model/gamma_regressor.py +61 -21
  151. snowflake/ml/modeling/linear_model/huber_regressor.py +61 -21
  152. snowflake/ml/modeling/linear_model/lars.py +61 -21
  153. snowflake/ml/modeling/linear_model/lars_cv.py +61 -21
  154. snowflake/ml/modeling/linear_model/lasso.py +61 -21
  155. snowflake/ml/modeling/linear_model/lasso_cv.py +61 -21
  156. snowflake/ml/modeling/linear_model/lasso_lars.py +61 -21
  157. snowflake/ml/modeling/linear_model/lasso_lars_cv.py +61 -21
  158. snowflake/ml/modeling/linear_model/lasso_lars_ic.py +61 -21
  159. snowflake/ml/modeling/linear_model/linear_regression.py +61 -21
  160. snowflake/ml/modeling/linear_model/logistic_regression.py +61 -21
  161. snowflake/ml/modeling/linear_model/logistic_regression_cv.py +61 -21
  162. snowflake/ml/modeling/linear_model/multi_task_elastic_net.py +61 -21
  163. snowflake/ml/modeling/linear_model/multi_task_elastic_net_cv.py +61 -21
  164. snowflake/ml/modeling/linear_model/multi_task_lasso.py +61 -21
  165. snowflake/ml/modeling/linear_model/multi_task_lasso_cv.py +61 -21
  166. snowflake/ml/modeling/linear_model/orthogonal_matching_pursuit.py +61 -21
  167. snowflake/ml/modeling/linear_model/passive_aggressive_classifier.py +61 -21
  168. snowflake/ml/modeling/linear_model/passive_aggressive_regressor.py +61 -21
  169. snowflake/ml/modeling/linear_model/perceptron.py +61 -21
  170. snowflake/ml/modeling/linear_model/poisson_regressor.py +61 -21
  171. snowflake/ml/modeling/linear_model/ransac_regressor.py +61 -21
  172. snowflake/ml/modeling/linear_model/ridge.py +61 -21
  173. snowflake/ml/modeling/linear_model/ridge_classifier.py +61 -21
  174. snowflake/ml/modeling/linear_model/ridge_classifier_cv.py +61 -21
  175. snowflake/ml/modeling/linear_model/ridge_cv.py +61 -21
  176. snowflake/ml/modeling/linear_model/sgd_classifier.py +61 -21
  177. snowflake/ml/modeling/linear_model/sgd_one_class_svm.py +61 -21
  178. snowflake/ml/modeling/linear_model/sgd_regressor.py +61 -21
  179. snowflake/ml/modeling/linear_model/theil_sen_regressor.py +61 -21
  180. snowflake/ml/modeling/linear_model/tweedie_regressor.py +61 -21
  181. snowflake/ml/modeling/manifold/isomap.py +61 -21
  182. snowflake/ml/modeling/manifold/mds.py +61 -21
  183. snowflake/ml/modeling/manifold/spectral_embedding.py +61 -21
  184. snowflake/ml/modeling/manifold/tsne.py +61 -21
  185. snowflake/ml/modeling/metrics/metrics_utils.py +2 -2
  186. snowflake/ml/modeling/metrics/ranking.py +0 -3
  187. snowflake/ml/modeling/metrics/regression.py +0 -3
  188. snowflake/ml/modeling/mixture/bayesian_gaussian_mixture.py +61 -21
  189. snowflake/ml/modeling/mixture/gaussian_mixture.py +61 -21
  190. snowflake/ml/modeling/multiclass/one_vs_one_classifier.py +61 -21
  191. snowflake/ml/modeling/multiclass/one_vs_rest_classifier.py +61 -21
  192. snowflake/ml/modeling/multiclass/output_code_classifier.py +61 -21
  193. snowflake/ml/modeling/naive_bayes/bernoulli_nb.py +61 -21
  194. snowflake/ml/modeling/naive_bayes/categorical_nb.py +61 -21
  195. snowflake/ml/modeling/naive_bayes/complement_nb.py +61 -21
  196. snowflake/ml/modeling/naive_bayes/gaussian_nb.py +61 -21
  197. snowflake/ml/modeling/naive_bayes/multinomial_nb.py +61 -21
  198. snowflake/ml/modeling/neighbors/k_neighbors_classifier.py +61 -21
  199. snowflake/ml/modeling/neighbors/k_neighbors_regressor.py +61 -21
  200. snowflake/ml/modeling/neighbors/kernel_density.py +61 -21
  201. snowflake/ml/modeling/neighbors/local_outlier_factor.py +61 -21
  202. snowflake/ml/modeling/neighbors/nearest_centroid.py +61 -21
  203. snowflake/ml/modeling/neighbors/nearest_neighbors.py +61 -21
  204. snowflake/ml/modeling/neighbors/neighborhood_components_analysis.py +61 -21
  205. snowflake/ml/modeling/neighbors/radius_neighbors_classifier.py +61 -21
  206. snowflake/ml/modeling/neighbors/radius_neighbors_regressor.py +61 -21
  207. snowflake/ml/modeling/neural_network/bernoulli_rbm.py +61 -21
  208. snowflake/ml/modeling/neural_network/mlp_classifier.py +61 -21
  209. snowflake/ml/modeling/neural_network/mlp_regressor.py +61 -21
  210. snowflake/ml/modeling/parameters/disable_model_tracer.py +5 -0
  211. snowflake/ml/modeling/pipeline/pipeline.py +1 -13
  212. snowflake/ml/modeling/preprocessing/polynomial_features.py +61 -21
  213. snowflake/ml/modeling/semi_supervised/label_propagation.py +61 -21
  214. snowflake/ml/modeling/semi_supervised/label_spreading.py +61 -21
  215. snowflake/ml/modeling/svm/linear_svc.py +61 -21
  216. snowflake/ml/modeling/svm/linear_svr.py +61 -21
  217. snowflake/ml/modeling/svm/nu_svc.py +61 -21
  218. snowflake/ml/modeling/svm/nu_svr.py +61 -21
  219. snowflake/ml/modeling/svm/svc.py +61 -21
  220. snowflake/ml/modeling/svm/svr.py +61 -21
  221. snowflake/ml/modeling/tree/decision_tree_classifier.py +61 -21
  222. snowflake/ml/modeling/tree/decision_tree_regressor.py +61 -21
  223. snowflake/ml/modeling/tree/extra_tree_classifier.py +61 -21
  224. snowflake/ml/modeling/tree/extra_tree_regressor.py +61 -21
  225. snowflake/ml/modeling/xgboost/xgb_classifier.py +64 -23
  226. snowflake/ml/modeling/xgboost/xgb_regressor.py +64 -23
  227. snowflake/ml/modeling/xgboost/xgbrf_classifier.py +64 -23
  228. snowflake/ml/modeling/xgboost/xgbrf_regressor.py +64 -23
  229. snowflake/ml/monitoring/_client/model_monitor.py +126 -0
  230. snowflake/ml/monitoring/_client/model_monitor_manager.py +361 -0
  231. snowflake/ml/monitoring/_client/model_monitor_version.py +1 -0
  232. snowflake/ml/monitoring/_client/monitor_sql_client.py +1335 -0
  233. snowflake/ml/monitoring/_client/queries/record_count.ssql +14 -0
  234. snowflake/ml/monitoring/_client/queries/rmse.ssql +28 -0
  235. snowflake/ml/monitoring/entities/model_monitor_config.py +28 -0
  236. snowflake/ml/monitoring/entities/model_monitor_interval.py +46 -0
  237. snowflake/ml/monitoring/entities/output_score_type.py +90 -0
  238. snowflake/ml/registry/_manager/model_manager.py +4 -0
  239. snowflake/ml/registry/registry.py +166 -8
  240. snowflake/ml/version.py +1 -1
  241. {snowflake_ml_python-1.6.1.dist-info → snowflake_ml_python-1.6.3.dist-info}/METADATA +43 -9
  242. snowflake_ml_python-1.6.3.dist-info/RECORD +400 -0
  243. {snowflake_ml_python-1.6.1.dist-info → snowflake_ml_python-1.6.3.dist-info}/WHEEL +1 -1
  244. snowflake/ml/_internal/container_services/image_registry/credential.py +0 -84
  245. snowflake/ml/_internal/container_services/image_registry/http_client.py +0 -127
  246. snowflake/ml/_internal/container_services/image_registry/imagelib.py +0 -400
  247. snowflake/ml/_internal/container_services/image_registry/registry_client.py +0 -212
  248. snowflake/ml/_internal/utils/log_stream_processor.py +0 -30
  249. snowflake/ml/_internal/utils/session_token_manager.py +0 -46
  250. snowflake/ml/_internal/utils/spcs_attribution_utils.py +0 -122
  251. snowflake/ml/_internal/utils/uri.py +0 -77
  252. snowflake/ml/data/torch_dataset.py +0 -33
  253. snowflake/ml/model/_api.py +0 -568
  254. snowflake/ml/model/_deploy_client/image_builds/base_image_builder.py +0 -12
  255. snowflake/ml/model/_deploy_client/image_builds/client_image_builder.py +0 -249
  256. snowflake/ml/model/_deploy_client/image_builds/docker_context.py +0 -130
  257. snowflake/ml/model/_deploy_client/image_builds/gunicorn_run.sh +0 -36
  258. snowflake/ml/model/_deploy_client/image_builds/inference_server/main.py +0 -268
  259. snowflake/ml/model/_deploy_client/image_builds/server_image_builder.py +0 -215
  260. snowflake/ml/model/_deploy_client/image_builds/templates/dockerfile_template +0 -53
  261. snowflake/ml/model/_deploy_client/image_builds/templates/image_build_job_spec_template +0 -38
  262. snowflake/ml/model/_deploy_client/image_builds/templates/kaniko_shell_script_template +0 -105
  263. snowflake/ml/model/_deploy_client/snowservice/deploy.py +0 -611
  264. snowflake/ml/model/_deploy_client/snowservice/deploy_options.py +0 -116
  265. snowflake/ml/model/_deploy_client/snowservice/instance_types.py +0 -10
  266. snowflake/ml/model/_deploy_client/snowservice/templates/service_spec_template +0 -28
  267. snowflake/ml/model/_deploy_client/snowservice/templates/service_spec_template_with_model +0 -21
  268. snowflake/ml/model/_deploy_client/utils/constants.py +0 -48
  269. snowflake/ml/model/_deploy_client/utils/snowservice_client.py +0 -280
  270. snowflake/ml/model/_deploy_client/warehouse/deploy.py +0 -202
  271. snowflake/ml/model/_deploy_client/warehouse/infer_template.py +0 -99
  272. snowflake/ml/model/_packager/model_handlers/llm.py +0 -267
  273. snowflake/ml/model/_packager/model_meta/_core_requirements.py +0 -11
  274. snowflake/ml/model/deploy_platforms.py +0 -6
  275. snowflake/ml/model/models/llm.py +0 -104
  276. snowflake/ml/monitoring/monitor.py +0 -203
  277. snowflake/ml/registry/_initial_schema.py +0 -142
  278. snowflake/ml/registry/_schema.py +0 -82
  279. snowflake/ml/registry/_schema_upgrade_plans.py +0 -116
  280. snowflake/ml/registry/_schema_version_manager.py +0 -163
  281. snowflake/ml/registry/model_registry.py +0 -2048
  282. snowflake_ml_python-1.6.1.dist-info/RECORD +0 -422
  283. {snowflake_ml_python-1.6.1.dist-info → snowflake_ml_python-1.6.3.dist-info}/LICENSE.txt +0 -0
  284. {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()