genesis-flow 1.0.0__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 (645) hide show
  1. genesis_flow-1.0.0.dist-info/METADATA +822 -0
  2. genesis_flow-1.0.0.dist-info/RECORD +645 -0
  3. genesis_flow-1.0.0.dist-info/WHEEL +5 -0
  4. genesis_flow-1.0.0.dist-info/entry_points.txt +19 -0
  5. genesis_flow-1.0.0.dist-info/licenses/LICENSE.txt +202 -0
  6. genesis_flow-1.0.0.dist-info/top_level.txt +1 -0
  7. mlflow/__init__.py +367 -0
  8. mlflow/__main__.py +3 -0
  9. mlflow/ag2/__init__.py +56 -0
  10. mlflow/ag2/ag2_logger.py +294 -0
  11. mlflow/anthropic/__init__.py +40 -0
  12. mlflow/anthropic/autolog.py +129 -0
  13. mlflow/anthropic/chat.py +144 -0
  14. mlflow/artifacts/__init__.py +268 -0
  15. mlflow/autogen/__init__.py +144 -0
  16. mlflow/autogen/chat.py +142 -0
  17. mlflow/azure/__init__.py +26 -0
  18. mlflow/azure/auth_handler.py +257 -0
  19. mlflow/azure/client.py +319 -0
  20. mlflow/azure/config.py +120 -0
  21. mlflow/azure/connection_factory.py +340 -0
  22. mlflow/azure/exceptions.py +27 -0
  23. mlflow/azure/stores.py +327 -0
  24. mlflow/azure/utils.py +183 -0
  25. mlflow/bedrock/__init__.py +45 -0
  26. mlflow/bedrock/_autolog.py +202 -0
  27. mlflow/bedrock/chat.py +122 -0
  28. mlflow/bedrock/stream.py +160 -0
  29. mlflow/bedrock/utils.py +43 -0
  30. mlflow/cli.py +707 -0
  31. mlflow/client.py +12 -0
  32. mlflow/config/__init__.py +56 -0
  33. mlflow/crewai/__init__.py +79 -0
  34. mlflow/crewai/autolog.py +253 -0
  35. mlflow/crewai/chat.py +29 -0
  36. mlflow/data/__init__.py +75 -0
  37. mlflow/data/artifact_dataset_sources.py +170 -0
  38. mlflow/data/code_dataset_source.py +40 -0
  39. mlflow/data/dataset.py +123 -0
  40. mlflow/data/dataset_registry.py +168 -0
  41. mlflow/data/dataset_source.py +110 -0
  42. mlflow/data/dataset_source_registry.py +219 -0
  43. mlflow/data/delta_dataset_source.py +167 -0
  44. mlflow/data/digest_utils.py +108 -0
  45. mlflow/data/evaluation_dataset.py +562 -0
  46. mlflow/data/filesystem_dataset_source.py +81 -0
  47. mlflow/data/http_dataset_source.py +145 -0
  48. mlflow/data/huggingface_dataset.py +258 -0
  49. mlflow/data/huggingface_dataset_source.py +118 -0
  50. mlflow/data/meta_dataset.py +104 -0
  51. mlflow/data/numpy_dataset.py +223 -0
  52. mlflow/data/pandas_dataset.py +231 -0
  53. mlflow/data/polars_dataset.py +352 -0
  54. mlflow/data/pyfunc_dataset_mixin.py +31 -0
  55. mlflow/data/schema.py +76 -0
  56. mlflow/data/sources.py +1 -0
  57. mlflow/data/spark_dataset.py +406 -0
  58. mlflow/data/spark_dataset_source.py +74 -0
  59. mlflow/data/spark_delta_utils.py +118 -0
  60. mlflow/data/tensorflow_dataset.py +350 -0
  61. mlflow/data/uc_volume_dataset_source.py +81 -0
  62. mlflow/db.py +27 -0
  63. mlflow/dspy/__init__.py +17 -0
  64. mlflow/dspy/autolog.py +197 -0
  65. mlflow/dspy/callback.py +398 -0
  66. mlflow/dspy/constant.py +1 -0
  67. mlflow/dspy/load.py +93 -0
  68. mlflow/dspy/save.py +393 -0
  69. mlflow/dspy/util.py +109 -0
  70. mlflow/dspy/wrapper.py +226 -0
  71. mlflow/entities/__init__.py +104 -0
  72. mlflow/entities/_mlflow_object.py +52 -0
  73. mlflow/entities/assessment.py +545 -0
  74. mlflow/entities/assessment_error.py +80 -0
  75. mlflow/entities/assessment_source.py +141 -0
  76. mlflow/entities/dataset.py +92 -0
  77. mlflow/entities/dataset_input.py +51 -0
  78. mlflow/entities/dataset_summary.py +62 -0
  79. mlflow/entities/document.py +48 -0
  80. mlflow/entities/experiment.py +109 -0
  81. mlflow/entities/experiment_tag.py +35 -0
  82. mlflow/entities/file_info.py +45 -0
  83. mlflow/entities/input_tag.py +35 -0
  84. mlflow/entities/lifecycle_stage.py +35 -0
  85. mlflow/entities/logged_model.py +228 -0
  86. mlflow/entities/logged_model_input.py +26 -0
  87. mlflow/entities/logged_model_output.py +32 -0
  88. mlflow/entities/logged_model_parameter.py +46 -0
  89. mlflow/entities/logged_model_status.py +74 -0
  90. mlflow/entities/logged_model_tag.py +33 -0
  91. mlflow/entities/metric.py +200 -0
  92. mlflow/entities/model_registry/__init__.py +29 -0
  93. mlflow/entities/model_registry/_model_registry_entity.py +13 -0
  94. mlflow/entities/model_registry/model_version.py +243 -0
  95. mlflow/entities/model_registry/model_version_deployment_job_run_state.py +44 -0
  96. mlflow/entities/model_registry/model_version_deployment_job_state.py +70 -0
  97. mlflow/entities/model_registry/model_version_search.py +25 -0
  98. mlflow/entities/model_registry/model_version_stages.py +25 -0
  99. mlflow/entities/model_registry/model_version_status.py +35 -0
  100. mlflow/entities/model_registry/model_version_tag.py +35 -0
  101. mlflow/entities/model_registry/prompt.py +73 -0
  102. mlflow/entities/model_registry/prompt_version.py +244 -0
  103. mlflow/entities/model_registry/registered_model.py +175 -0
  104. mlflow/entities/model_registry/registered_model_alias.py +35 -0
  105. mlflow/entities/model_registry/registered_model_deployment_job_state.py +39 -0
  106. mlflow/entities/model_registry/registered_model_search.py +25 -0
  107. mlflow/entities/model_registry/registered_model_tag.py +35 -0
  108. mlflow/entities/multipart_upload.py +74 -0
  109. mlflow/entities/param.py +49 -0
  110. mlflow/entities/run.py +97 -0
  111. mlflow/entities/run_data.py +84 -0
  112. mlflow/entities/run_info.py +188 -0
  113. mlflow/entities/run_inputs.py +59 -0
  114. mlflow/entities/run_outputs.py +43 -0
  115. mlflow/entities/run_status.py +41 -0
  116. mlflow/entities/run_tag.py +36 -0
  117. mlflow/entities/source_type.py +31 -0
  118. mlflow/entities/span.py +774 -0
  119. mlflow/entities/span_event.py +96 -0
  120. mlflow/entities/span_status.py +102 -0
  121. mlflow/entities/trace.py +317 -0
  122. mlflow/entities/trace_data.py +71 -0
  123. mlflow/entities/trace_info.py +220 -0
  124. mlflow/entities/trace_info_v2.py +162 -0
  125. mlflow/entities/trace_location.py +173 -0
  126. mlflow/entities/trace_state.py +39 -0
  127. mlflow/entities/trace_status.py +68 -0
  128. mlflow/entities/view_type.py +51 -0
  129. mlflow/environment_variables.py +866 -0
  130. mlflow/evaluation/__init__.py +16 -0
  131. mlflow/evaluation/assessment.py +369 -0
  132. mlflow/evaluation/evaluation.py +411 -0
  133. mlflow/evaluation/evaluation_tag.py +61 -0
  134. mlflow/evaluation/fluent.py +48 -0
  135. mlflow/evaluation/utils.py +201 -0
  136. mlflow/exceptions.py +213 -0
  137. mlflow/experiments.py +140 -0
  138. mlflow/gemini/__init__.py +81 -0
  139. mlflow/gemini/autolog.py +186 -0
  140. mlflow/gemini/chat.py +261 -0
  141. mlflow/genai/__init__.py +71 -0
  142. mlflow/genai/datasets/__init__.py +67 -0
  143. mlflow/genai/datasets/evaluation_dataset.py +131 -0
  144. mlflow/genai/evaluation/__init__.py +3 -0
  145. mlflow/genai/evaluation/base.py +411 -0
  146. mlflow/genai/evaluation/constant.py +23 -0
  147. mlflow/genai/evaluation/utils.py +244 -0
  148. mlflow/genai/judges/__init__.py +21 -0
  149. mlflow/genai/judges/databricks.py +404 -0
  150. mlflow/genai/label_schemas/__init__.py +153 -0
  151. mlflow/genai/label_schemas/label_schemas.py +209 -0
  152. mlflow/genai/labeling/__init__.py +159 -0
  153. mlflow/genai/labeling/labeling.py +250 -0
  154. mlflow/genai/optimize/__init__.py +13 -0
  155. mlflow/genai/optimize/base.py +198 -0
  156. mlflow/genai/optimize/optimizers/__init__.py +4 -0
  157. mlflow/genai/optimize/optimizers/base_optimizer.py +38 -0
  158. mlflow/genai/optimize/optimizers/dspy_mipro_optimizer.py +221 -0
  159. mlflow/genai/optimize/optimizers/dspy_optimizer.py +91 -0
  160. mlflow/genai/optimize/optimizers/utils/dspy_mipro_callback.py +76 -0
  161. mlflow/genai/optimize/optimizers/utils/dspy_mipro_utils.py +18 -0
  162. mlflow/genai/optimize/types.py +75 -0
  163. mlflow/genai/optimize/util.py +30 -0
  164. mlflow/genai/prompts/__init__.py +206 -0
  165. mlflow/genai/scheduled_scorers.py +431 -0
  166. mlflow/genai/scorers/__init__.py +26 -0
  167. mlflow/genai/scorers/base.py +492 -0
  168. mlflow/genai/scorers/builtin_scorers.py +765 -0
  169. mlflow/genai/scorers/scorer_utils.py +138 -0
  170. mlflow/genai/scorers/validation.py +165 -0
  171. mlflow/genai/utils/data_validation.py +146 -0
  172. mlflow/genai/utils/enum_utils.py +23 -0
  173. mlflow/genai/utils/trace_utils.py +211 -0
  174. mlflow/groq/__init__.py +42 -0
  175. mlflow/groq/_groq_autolog.py +74 -0
  176. mlflow/johnsnowlabs/__init__.py +888 -0
  177. mlflow/langchain/__init__.py +24 -0
  178. mlflow/langchain/api_request_parallel_processor.py +330 -0
  179. mlflow/langchain/autolog.py +147 -0
  180. mlflow/langchain/chat_agent_langgraph.py +340 -0
  181. mlflow/langchain/constant.py +1 -0
  182. mlflow/langchain/constants.py +1 -0
  183. mlflow/langchain/databricks_dependencies.py +444 -0
  184. mlflow/langchain/langchain_tracer.py +597 -0
  185. mlflow/langchain/model.py +919 -0
  186. mlflow/langchain/output_parsers.py +142 -0
  187. mlflow/langchain/retriever_chain.py +153 -0
  188. mlflow/langchain/runnables.py +527 -0
  189. mlflow/langchain/utils/chat.py +402 -0
  190. mlflow/langchain/utils/logging.py +671 -0
  191. mlflow/langchain/utils/serialization.py +36 -0
  192. mlflow/legacy_databricks_cli/__init__.py +0 -0
  193. mlflow/legacy_databricks_cli/configure/__init__.py +0 -0
  194. mlflow/legacy_databricks_cli/configure/provider.py +482 -0
  195. mlflow/litellm/__init__.py +175 -0
  196. mlflow/llama_index/__init__.py +22 -0
  197. mlflow/llama_index/autolog.py +55 -0
  198. mlflow/llama_index/chat.py +43 -0
  199. mlflow/llama_index/constant.py +1 -0
  200. mlflow/llama_index/model.py +577 -0
  201. mlflow/llama_index/pyfunc_wrapper.py +332 -0
  202. mlflow/llama_index/serialize_objects.py +188 -0
  203. mlflow/llama_index/tracer.py +561 -0
  204. mlflow/metrics/__init__.py +479 -0
  205. mlflow/metrics/base.py +39 -0
  206. mlflow/metrics/genai/__init__.py +25 -0
  207. mlflow/metrics/genai/base.py +101 -0
  208. mlflow/metrics/genai/genai_metric.py +771 -0
  209. mlflow/metrics/genai/metric_definitions.py +450 -0
  210. mlflow/metrics/genai/model_utils.py +371 -0
  211. mlflow/metrics/genai/prompt_template.py +68 -0
  212. mlflow/metrics/genai/prompts/__init__.py +0 -0
  213. mlflow/metrics/genai/prompts/v1.py +422 -0
  214. mlflow/metrics/genai/utils.py +6 -0
  215. mlflow/metrics/metric_definitions.py +619 -0
  216. mlflow/mismatch.py +34 -0
  217. mlflow/mistral/__init__.py +34 -0
  218. mlflow/mistral/autolog.py +71 -0
  219. mlflow/mistral/chat.py +135 -0
  220. mlflow/ml_package_versions.py +452 -0
  221. mlflow/models/__init__.py +97 -0
  222. mlflow/models/auth_policy.py +83 -0
  223. mlflow/models/cli.py +354 -0
  224. mlflow/models/container/__init__.py +294 -0
  225. mlflow/models/container/scoring_server/__init__.py +0 -0
  226. mlflow/models/container/scoring_server/nginx.conf +39 -0
  227. mlflow/models/dependencies_schemas.py +287 -0
  228. mlflow/models/display_utils.py +158 -0
  229. mlflow/models/docker_utils.py +211 -0
  230. mlflow/models/evaluation/__init__.py +23 -0
  231. mlflow/models/evaluation/_shap_patch.py +64 -0
  232. mlflow/models/evaluation/artifacts.py +194 -0
  233. mlflow/models/evaluation/base.py +1811 -0
  234. mlflow/models/evaluation/calibration_curve.py +109 -0
  235. mlflow/models/evaluation/default_evaluator.py +996 -0
  236. mlflow/models/evaluation/deprecated.py +23 -0
  237. mlflow/models/evaluation/evaluator_registry.py +80 -0
  238. mlflow/models/evaluation/evaluators/classifier.py +704 -0
  239. mlflow/models/evaluation/evaluators/default.py +233 -0
  240. mlflow/models/evaluation/evaluators/regressor.py +96 -0
  241. mlflow/models/evaluation/evaluators/shap.py +296 -0
  242. mlflow/models/evaluation/lift_curve.py +178 -0
  243. mlflow/models/evaluation/utils/metric.py +123 -0
  244. mlflow/models/evaluation/utils/trace.py +179 -0
  245. mlflow/models/evaluation/validation.py +434 -0
  246. mlflow/models/flavor_backend.py +93 -0
  247. mlflow/models/flavor_backend_registry.py +53 -0
  248. mlflow/models/model.py +1639 -0
  249. mlflow/models/model_config.py +150 -0
  250. mlflow/models/notebook_resources/agent_evaluation_template.html +235 -0
  251. mlflow/models/notebook_resources/eval_with_dataset_example.py +22 -0
  252. mlflow/models/notebook_resources/eval_with_synthetic_example.py +22 -0
  253. mlflow/models/python_api.py +369 -0
  254. mlflow/models/rag_signatures.py +128 -0
  255. mlflow/models/resources.py +321 -0
  256. mlflow/models/signature.py +662 -0
  257. mlflow/models/utils.py +2054 -0
  258. mlflow/models/wheeled_model.py +280 -0
  259. mlflow/openai/__init__.py +57 -0
  260. mlflow/openai/_agent_tracer.py +364 -0
  261. mlflow/openai/api_request_parallel_processor.py +131 -0
  262. mlflow/openai/autolog.py +509 -0
  263. mlflow/openai/constant.py +1 -0
  264. mlflow/openai/model.py +824 -0
  265. mlflow/openai/utils/chat_schema.py +367 -0
  266. mlflow/optuna/__init__.py +3 -0
  267. mlflow/optuna/storage.py +646 -0
  268. mlflow/plugins/__init__.py +72 -0
  269. mlflow/plugins/base.py +358 -0
  270. mlflow/plugins/builtin/__init__.py +24 -0
  271. mlflow/plugins/builtin/pytorch_plugin.py +150 -0
  272. mlflow/plugins/builtin/sklearn_plugin.py +158 -0
  273. mlflow/plugins/builtin/transformers_plugin.py +187 -0
  274. mlflow/plugins/cli.py +321 -0
  275. mlflow/plugins/discovery.py +340 -0
  276. mlflow/plugins/manager.py +465 -0
  277. mlflow/plugins/registry.py +316 -0
  278. mlflow/plugins/templates/framework_plugin_template.py +329 -0
  279. mlflow/prompt/constants.py +20 -0
  280. mlflow/prompt/promptlab_model.py +197 -0
  281. mlflow/prompt/registry_utils.py +248 -0
  282. mlflow/promptflow/__init__.py +495 -0
  283. mlflow/protos/__init__.py +0 -0
  284. mlflow/protos/assessments_pb2.py +174 -0
  285. mlflow/protos/databricks_artifacts_pb2.py +489 -0
  286. mlflow/protos/databricks_filesystem_service_pb2.py +196 -0
  287. mlflow/protos/databricks_managed_catalog_messages_pb2.py +95 -0
  288. mlflow/protos/databricks_managed_catalog_service_pb2.py +86 -0
  289. mlflow/protos/databricks_pb2.py +267 -0
  290. mlflow/protos/databricks_trace_server_pb2.py +374 -0
  291. mlflow/protos/databricks_uc_registry_messages_pb2.py +1249 -0
  292. mlflow/protos/databricks_uc_registry_service_pb2.py +170 -0
  293. mlflow/protos/facet_feature_statistics_pb2.py +296 -0
  294. mlflow/protos/internal_pb2.py +77 -0
  295. mlflow/protos/mlflow_artifacts_pb2.py +336 -0
  296. mlflow/protos/model_registry_pb2.py +1073 -0
  297. mlflow/protos/scalapb/__init__.py +0 -0
  298. mlflow/protos/scalapb/scalapb_pb2.py +104 -0
  299. mlflow/protos/service_pb2.py +2600 -0
  300. mlflow/protos/unity_catalog_oss_messages_pb2.py +457 -0
  301. mlflow/protos/unity_catalog_oss_service_pb2.py +130 -0
  302. mlflow/protos/unity_catalog_prompt_messages_pb2.py +447 -0
  303. mlflow/protos/unity_catalog_prompt_messages_pb2_grpc.py +24 -0
  304. mlflow/protos/unity_catalog_prompt_service_pb2.py +164 -0
  305. mlflow/protos/unity_catalog_prompt_service_pb2_grpc.py +785 -0
  306. mlflow/py.typed +0 -0
  307. mlflow/pydantic_ai/__init__.py +57 -0
  308. mlflow/pydantic_ai/autolog.py +173 -0
  309. mlflow/pyfunc/__init__.py +3844 -0
  310. mlflow/pyfunc/_mlflow_pyfunc_backend_predict.py +61 -0
  311. mlflow/pyfunc/backend.py +523 -0
  312. mlflow/pyfunc/context.py +78 -0
  313. mlflow/pyfunc/dbconnect_artifact_cache.py +144 -0
  314. mlflow/pyfunc/loaders/__init__.py +7 -0
  315. mlflow/pyfunc/loaders/chat_agent.py +117 -0
  316. mlflow/pyfunc/loaders/chat_model.py +125 -0
  317. mlflow/pyfunc/loaders/code_model.py +31 -0
  318. mlflow/pyfunc/loaders/responses_agent.py +112 -0
  319. mlflow/pyfunc/mlserver.py +46 -0
  320. mlflow/pyfunc/model.py +1473 -0
  321. mlflow/pyfunc/scoring_server/__init__.py +604 -0
  322. mlflow/pyfunc/scoring_server/app.py +7 -0
  323. mlflow/pyfunc/scoring_server/client.py +146 -0
  324. mlflow/pyfunc/spark_model_cache.py +48 -0
  325. mlflow/pyfunc/stdin_server.py +44 -0
  326. mlflow/pyfunc/utils/__init__.py +3 -0
  327. mlflow/pyfunc/utils/data_validation.py +224 -0
  328. mlflow/pyfunc/utils/environment.py +22 -0
  329. mlflow/pyfunc/utils/input_converter.py +47 -0
  330. mlflow/pyfunc/utils/serving_data_parser.py +11 -0
  331. mlflow/pytorch/__init__.py +1171 -0
  332. mlflow/pytorch/_lightning_autolog.py +580 -0
  333. mlflow/pytorch/_pytorch_autolog.py +50 -0
  334. mlflow/pytorch/pickle_module.py +35 -0
  335. mlflow/rfunc/__init__.py +42 -0
  336. mlflow/rfunc/backend.py +134 -0
  337. mlflow/runs.py +89 -0
  338. mlflow/server/__init__.py +302 -0
  339. mlflow/server/auth/__init__.py +1224 -0
  340. mlflow/server/auth/__main__.py +4 -0
  341. mlflow/server/auth/basic_auth.ini +6 -0
  342. mlflow/server/auth/cli.py +11 -0
  343. mlflow/server/auth/client.py +537 -0
  344. mlflow/server/auth/config.py +34 -0
  345. mlflow/server/auth/db/__init__.py +0 -0
  346. mlflow/server/auth/db/cli.py +18 -0
  347. mlflow/server/auth/db/migrations/__init__.py +0 -0
  348. mlflow/server/auth/db/migrations/alembic.ini +110 -0
  349. mlflow/server/auth/db/migrations/env.py +76 -0
  350. mlflow/server/auth/db/migrations/versions/8606fa83a998_initial_migration.py +51 -0
  351. mlflow/server/auth/db/migrations/versions/__init__.py +0 -0
  352. mlflow/server/auth/db/models.py +67 -0
  353. mlflow/server/auth/db/utils.py +37 -0
  354. mlflow/server/auth/entities.py +165 -0
  355. mlflow/server/auth/logo.py +14 -0
  356. mlflow/server/auth/permissions.py +65 -0
  357. mlflow/server/auth/routes.py +18 -0
  358. mlflow/server/auth/sqlalchemy_store.py +263 -0
  359. mlflow/server/graphql/__init__.py +0 -0
  360. mlflow/server/graphql/autogenerated_graphql_schema.py +353 -0
  361. mlflow/server/graphql/graphql_custom_scalars.py +24 -0
  362. mlflow/server/graphql/graphql_errors.py +15 -0
  363. mlflow/server/graphql/graphql_no_batching.py +89 -0
  364. mlflow/server/graphql/graphql_schema_extensions.py +74 -0
  365. mlflow/server/handlers.py +3217 -0
  366. mlflow/server/prometheus_exporter.py +17 -0
  367. mlflow/server/validation.py +30 -0
  368. mlflow/shap/__init__.py +691 -0
  369. mlflow/sklearn/__init__.py +1994 -0
  370. mlflow/sklearn/utils.py +1041 -0
  371. mlflow/smolagents/__init__.py +66 -0
  372. mlflow/smolagents/autolog.py +139 -0
  373. mlflow/smolagents/chat.py +29 -0
  374. mlflow/store/__init__.py +10 -0
  375. mlflow/store/_unity_catalog/__init__.py +1 -0
  376. mlflow/store/_unity_catalog/lineage/__init__.py +1 -0
  377. mlflow/store/_unity_catalog/lineage/constants.py +2 -0
  378. mlflow/store/_unity_catalog/registry/__init__.py +6 -0
  379. mlflow/store/_unity_catalog/registry/prompt_info.py +75 -0
  380. mlflow/store/_unity_catalog/registry/rest_store.py +1740 -0
  381. mlflow/store/_unity_catalog/registry/uc_oss_rest_store.py +507 -0
  382. mlflow/store/_unity_catalog/registry/utils.py +121 -0
  383. mlflow/store/artifact/__init__.py +0 -0
  384. mlflow/store/artifact/artifact_repo.py +472 -0
  385. mlflow/store/artifact/artifact_repository_registry.py +154 -0
  386. mlflow/store/artifact/azure_blob_artifact_repo.py +275 -0
  387. mlflow/store/artifact/azure_data_lake_artifact_repo.py +295 -0
  388. mlflow/store/artifact/cli.py +141 -0
  389. mlflow/store/artifact/cloud_artifact_repo.py +332 -0
  390. mlflow/store/artifact/databricks_artifact_repo.py +729 -0
  391. mlflow/store/artifact/databricks_artifact_repo_resources.py +301 -0
  392. mlflow/store/artifact/databricks_logged_model_artifact_repo.py +93 -0
  393. mlflow/store/artifact/databricks_models_artifact_repo.py +216 -0
  394. mlflow/store/artifact/databricks_sdk_artifact_repo.py +134 -0
  395. mlflow/store/artifact/databricks_sdk_models_artifact_repo.py +97 -0
  396. mlflow/store/artifact/dbfs_artifact_repo.py +240 -0
  397. mlflow/store/artifact/ftp_artifact_repo.py +132 -0
  398. mlflow/store/artifact/gcs_artifact_repo.py +296 -0
  399. mlflow/store/artifact/hdfs_artifact_repo.py +209 -0
  400. mlflow/store/artifact/http_artifact_repo.py +218 -0
  401. mlflow/store/artifact/local_artifact_repo.py +142 -0
  402. mlflow/store/artifact/mlflow_artifacts_repo.py +94 -0
  403. mlflow/store/artifact/models_artifact_repo.py +259 -0
  404. mlflow/store/artifact/optimized_s3_artifact_repo.py +356 -0
  405. mlflow/store/artifact/presigned_url_artifact_repo.py +173 -0
  406. mlflow/store/artifact/r2_artifact_repo.py +70 -0
  407. mlflow/store/artifact/runs_artifact_repo.py +265 -0
  408. mlflow/store/artifact/s3_artifact_repo.py +330 -0
  409. mlflow/store/artifact/sftp_artifact_repo.py +141 -0
  410. mlflow/store/artifact/uc_volume_artifact_repo.py +76 -0
  411. mlflow/store/artifact/unity_catalog_models_artifact_repo.py +168 -0
  412. mlflow/store/artifact/unity_catalog_oss_models_artifact_repo.py +168 -0
  413. mlflow/store/artifact/utils/__init__.py +0 -0
  414. mlflow/store/artifact/utils/models.py +148 -0
  415. mlflow/store/db/__init__.py +0 -0
  416. mlflow/store/db/base_sql_model.py +3 -0
  417. mlflow/store/db/db_types.py +10 -0
  418. mlflow/store/db/utils.py +314 -0
  419. mlflow/store/db_migrations/__init__.py +0 -0
  420. mlflow/store/db_migrations/alembic.ini +74 -0
  421. mlflow/store/db_migrations/env.py +84 -0
  422. mlflow/store/db_migrations/versions/0584bdc529eb_add_cascading_deletion_to_datasets_from_experiments.py +88 -0
  423. mlflow/store/db_migrations/versions/0a8213491aaa_drop_duplicate_killed_constraint.py +49 -0
  424. mlflow/store/db_migrations/versions/0c779009ac13_add_deleted_time_field_to_runs_table.py +24 -0
  425. mlflow/store/db_migrations/versions/181f10493468_allow_nulls_for_metric_values.py +35 -0
  426. mlflow/store/db_migrations/versions/27a6a02d2cf1_add_model_version_tags_table.py +38 -0
  427. mlflow/store/db_migrations/versions/2b4d017a5e9b_add_model_registry_tables_to_db.py +77 -0
  428. mlflow/store/db_migrations/versions/2d6e25af4d3e_increase_max_param_val_length.py +33 -0
  429. mlflow/store/db_migrations/versions/3500859a5d39_add_model_aliases_table.py +50 -0
  430. mlflow/store/db_migrations/versions/39d1c3be5f05_add_is_nan_constraint_for_metrics_tables_if_necessary.py +41 -0
  431. mlflow/store/db_migrations/versions/400f98739977_add_logged_model_tables.py +123 -0
  432. mlflow/store/db_migrations/versions/4465047574b1_increase_max_dataset_schema_size.py +38 -0
  433. mlflow/store/db_migrations/versions/451aebb31d03_add_metric_step.py +35 -0
  434. mlflow/store/db_migrations/versions/5b0e9adcef9c_add_cascade_deletion_to_trace_tables_fk.py +40 -0
  435. mlflow/store/db_migrations/versions/6953534de441_add_step_to_inputs_table.py +25 -0
  436. mlflow/store/db_migrations/versions/728d730b5ebd_add_registered_model_tags_table.py +38 -0
  437. mlflow/store/db_migrations/versions/7ac759974ad8_update_run_tags_with_larger_limit.py +36 -0
  438. mlflow/store/db_migrations/versions/7f2a7d5fae7d_add_datasets_inputs_input_tags_tables.py +82 -0
  439. mlflow/store/db_migrations/versions/84291f40a231_add_run_link_to_model_version.py +26 -0
  440. mlflow/store/db_migrations/versions/867495a8f9d4_add_trace_tables.py +90 -0
  441. mlflow/store/db_migrations/versions/89d4b8295536_create_latest_metrics_table.py +169 -0
  442. mlflow/store/db_migrations/versions/90e64c465722_migrate_user_column_to_tags.py +64 -0
  443. mlflow/store/db_migrations/versions/97727af70f4d_creation_time_last_update_time_experiments.py +25 -0
  444. mlflow/store/db_migrations/versions/__init__.py +0 -0
  445. mlflow/store/db_migrations/versions/a8c4a736bde6_allow_nulls_for_run_id.py +27 -0
  446. mlflow/store/db_migrations/versions/acf3f17fdcc7_add_storage_location_field_to_model_.py +29 -0
  447. mlflow/store/db_migrations/versions/bd07f7e963c5_create_index_on_run_uuid.py +26 -0
  448. mlflow/store/db_migrations/versions/bda7b8c39065_increase_model_version_tag_value_limit.py +38 -0
  449. mlflow/store/db_migrations/versions/c48cb773bb87_reset_default_value_for_is_nan_in_metrics_table_for_mysql.py +41 -0
  450. mlflow/store/db_migrations/versions/cbc13b556ace_add_v3_trace_schema_columns.py +31 -0
  451. mlflow/store/db_migrations/versions/cc1f77228345_change_param_value_length_to_500.py +34 -0
  452. mlflow/store/db_migrations/versions/cfd24bdc0731_update_run_status_constraint_with_killed.py +78 -0
  453. mlflow/store/db_migrations/versions/df50e92ffc5e_add_experiment_tags_table.py +38 -0
  454. mlflow/store/db_migrations/versions/f5a4f2784254_increase_run_tag_value_limit.py +36 -0
  455. mlflow/store/entities/__init__.py +3 -0
  456. mlflow/store/entities/paged_list.py +18 -0
  457. mlflow/store/model_registry/__init__.py +10 -0
  458. mlflow/store/model_registry/abstract_store.py +1081 -0
  459. mlflow/store/model_registry/base_rest_store.py +44 -0
  460. mlflow/store/model_registry/databricks_workspace_model_registry_rest_store.py +37 -0
  461. mlflow/store/model_registry/dbmodels/__init__.py +0 -0
  462. mlflow/store/model_registry/dbmodels/models.py +206 -0
  463. mlflow/store/model_registry/file_store.py +1091 -0
  464. mlflow/store/model_registry/rest_store.py +481 -0
  465. mlflow/store/model_registry/sqlalchemy_store.py +1286 -0
  466. mlflow/store/tracking/__init__.py +23 -0
  467. mlflow/store/tracking/abstract_store.py +816 -0
  468. mlflow/store/tracking/dbmodels/__init__.py +0 -0
  469. mlflow/store/tracking/dbmodels/initial_models.py +243 -0
  470. mlflow/store/tracking/dbmodels/models.py +1073 -0
  471. mlflow/store/tracking/file_store.py +2438 -0
  472. mlflow/store/tracking/postgres_managed_identity.py +146 -0
  473. mlflow/store/tracking/rest_store.py +1131 -0
  474. mlflow/store/tracking/sqlalchemy_store.py +2785 -0
  475. mlflow/system_metrics/__init__.py +61 -0
  476. mlflow/system_metrics/metrics/__init__.py +0 -0
  477. mlflow/system_metrics/metrics/base_metrics_monitor.py +32 -0
  478. mlflow/system_metrics/metrics/cpu_monitor.py +23 -0
  479. mlflow/system_metrics/metrics/disk_monitor.py +21 -0
  480. mlflow/system_metrics/metrics/gpu_monitor.py +71 -0
  481. mlflow/system_metrics/metrics/network_monitor.py +34 -0
  482. mlflow/system_metrics/metrics/rocm_monitor.py +123 -0
  483. mlflow/system_metrics/system_metrics_monitor.py +198 -0
  484. mlflow/tracing/__init__.py +16 -0
  485. mlflow/tracing/assessment.py +356 -0
  486. mlflow/tracing/client.py +531 -0
  487. mlflow/tracing/config.py +125 -0
  488. mlflow/tracing/constant.py +105 -0
  489. mlflow/tracing/destination.py +81 -0
  490. mlflow/tracing/display/__init__.py +40 -0
  491. mlflow/tracing/display/display_handler.py +196 -0
  492. mlflow/tracing/export/async_export_queue.py +186 -0
  493. mlflow/tracing/export/inference_table.py +138 -0
  494. mlflow/tracing/export/mlflow_v3.py +137 -0
  495. mlflow/tracing/export/utils.py +70 -0
  496. mlflow/tracing/fluent.py +1417 -0
  497. mlflow/tracing/processor/base_mlflow.py +199 -0
  498. mlflow/tracing/processor/inference_table.py +175 -0
  499. mlflow/tracing/processor/mlflow_v3.py +47 -0
  500. mlflow/tracing/processor/otel.py +73 -0
  501. mlflow/tracing/provider.py +487 -0
  502. mlflow/tracing/trace_manager.py +200 -0
  503. mlflow/tracing/utils/__init__.py +616 -0
  504. mlflow/tracing/utils/artifact_utils.py +28 -0
  505. mlflow/tracing/utils/copy.py +55 -0
  506. mlflow/tracing/utils/environment.py +55 -0
  507. mlflow/tracing/utils/exception.py +21 -0
  508. mlflow/tracing/utils/once.py +35 -0
  509. mlflow/tracing/utils/otlp.py +63 -0
  510. mlflow/tracing/utils/processor.py +54 -0
  511. mlflow/tracing/utils/search.py +292 -0
  512. mlflow/tracing/utils/timeout.py +250 -0
  513. mlflow/tracing/utils/token.py +19 -0
  514. mlflow/tracing/utils/truncation.py +124 -0
  515. mlflow/tracing/utils/warning.py +76 -0
  516. mlflow/tracking/__init__.py +39 -0
  517. mlflow/tracking/_model_registry/__init__.py +1 -0
  518. mlflow/tracking/_model_registry/client.py +764 -0
  519. mlflow/tracking/_model_registry/fluent.py +853 -0
  520. mlflow/tracking/_model_registry/registry.py +67 -0
  521. mlflow/tracking/_model_registry/utils.py +251 -0
  522. mlflow/tracking/_tracking_service/__init__.py +0 -0
  523. mlflow/tracking/_tracking_service/client.py +883 -0
  524. mlflow/tracking/_tracking_service/registry.py +56 -0
  525. mlflow/tracking/_tracking_service/utils.py +275 -0
  526. mlflow/tracking/artifact_utils.py +179 -0
  527. mlflow/tracking/client.py +5900 -0
  528. mlflow/tracking/context/__init__.py +0 -0
  529. mlflow/tracking/context/abstract_context.py +35 -0
  530. mlflow/tracking/context/databricks_cluster_context.py +15 -0
  531. mlflow/tracking/context/databricks_command_context.py +15 -0
  532. mlflow/tracking/context/databricks_job_context.py +49 -0
  533. mlflow/tracking/context/databricks_notebook_context.py +41 -0
  534. mlflow/tracking/context/databricks_repo_context.py +43 -0
  535. mlflow/tracking/context/default_context.py +51 -0
  536. mlflow/tracking/context/git_context.py +32 -0
  537. mlflow/tracking/context/registry.py +98 -0
  538. mlflow/tracking/context/system_environment_context.py +15 -0
  539. mlflow/tracking/default_experiment/__init__.py +1 -0
  540. mlflow/tracking/default_experiment/abstract_context.py +43 -0
  541. mlflow/tracking/default_experiment/databricks_notebook_experiment_provider.py +44 -0
  542. mlflow/tracking/default_experiment/registry.py +75 -0
  543. mlflow/tracking/fluent.py +3595 -0
  544. mlflow/tracking/metric_value_conversion_utils.py +93 -0
  545. mlflow/tracking/multimedia.py +206 -0
  546. mlflow/tracking/registry.py +86 -0
  547. mlflow/tracking/request_auth/__init__.py +0 -0
  548. mlflow/tracking/request_auth/abstract_request_auth_provider.py +34 -0
  549. mlflow/tracking/request_auth/registry.py +60 -0
  550. mlflow/tracking/request_header/__init__.py +0 -0
  551. mlflow/tracking/request_header/abstract_request_header_provider.py +36 -0
  552. mlflow/tracking/request_header/databricks_request_header_provider.py +38 -0
  553. mlflow/tracking/request_header/default_request_header_provider.py +17 -0
  554. mlflow/tracking/request_header/registry.py +79 -0
  555. mlflow/transformers/__init__.py +2982 -0
  556. mlflow/transformers/flavor_config.py +258 -0
  557. mlflow/transformers/hub_utils.py +83 -0
  558. mlflow/transformers/llm_inference_utils.py +468 -0
  559. mlflow/transformers/model_io.py +301 -0
  560. mlflow/transformers/peft.py +51 -0
  561. mlflow/transformers/signature.py +183 -0
  562. mlflow/transformers/torch_utils.py +55 -0
  563. mlflow/types/__init__.py +21 -0
  564. mlflow/types/agent.py +270 -0
  565. mlflow/types/chat.py +240 -0
  566. mlflow/types/llm.py +935 -0
  567. mlflow/types/responses.py +139 -0
  568. mlflow/types/responses_helpers.py +416 -0
  569. mlflow/types/schema.py +1505 -0
  570. mlflow/types/type_hints.py +647 -0
  571. mlflow/types/utils.py +753 -0
  572. mlflow/utils/__init__.py +283 -0
  573. mlflow/utils/_capture_modules.py +256 -0
  574. mlflow/utils/_capture_transformers_modules.py +75 -0
  575. mlflow/utils/_spark_utils.py +201 -0
  576. mlflow/utils/_unity_catalog_oss_utils.py +97 -0
  577. mlflow/utils/_unity_catalog_utils.py +479 -0
  578. mlflow/utils/annotations.py +218 -0
  579. mlflow/utils/arguments_utils.py +16 -0
  580. mlflow/utils/async_logging/__init__.py +1 -0
  581. mlflow/utils/async_logging/async_artifacts_logging_queue.py +258 -0
  582. mlflow/utils/async_logging/async_logging_queue.py +366 -0
  583. mlflow/utils/async_logging/run_artifact.py +38 -0
  584. mlflow/utils/async_logging/run_batch.py +58 -0
  585. mlflow/utils/async_logging/run_operations.py +49 -0
  586. mlflow/utils/autologging_utils/__init__.py +737 -0
  587. mlflow/utils/autologging_utils/client.py +432 -0
  588. mlflow/utils/autologging_utils/config.py +33 -0
  589. mlflow/utils/autologging_utils/events.py +294 -0
  590. mlflow/utils/autologging_utils/logging_and_warnings.py +328 -0
  591. mlflow/utils/autologging_utils/metrics_queue.py +71 -0
  592. mlflow/utils/autologging_utils/safety.py +1104 -0
  593. mlflow/utils/autologging_utils/versioning.py +95 -0
  594. mlflow/utils/checkpoint_utils.py +206 -0
  595. mlflow/utils/class_utils.py +6 -0
  596. mlflow/utils/cli_args.py +257 -0
  597. mlflow/utils/conda.py +354 -0
  598. mlflow/utils/credentials.py +231 -0
  599. mlflow/utils/data_utils.py +17 -0
  600. mlflow/utils/databricks_utils.py +1436 -0
  601. mlflow/utils/docstring_utils.py +477 -0
  602. mlflow/utils/doctor.py +133 -0
  603. mlflow/utils/download_cloud_file_chunk.py +43 -0
  604. mlflow/utils/env_manager.py +16 -0
  605. mlflow/utils/env_pack.py +131 -0
  606. mlflow/utils/environment.py +1009 -0
  607. mlflow/utils/exception_utils.py +14 -0
  608. mlflow/utils/file_utils.py +978 -0
  609. mlflow/utils/git_utils.py +77 -0
  610. mlflow/utils/gorilla.py +797 -0
  611. mlflow/utils/import_hooks/__init__.py +363 -0
  612. mlflow/utils/lazy_load.py +51 -0
  613. mlflow/utils/logging_utils.py +168 -0
  614. mlflow/utils/mime_type_utils.py +58 -0
  615. mlflow/utils/mlflow_tags.py +103 -0
  616. mlflow/utils/model_utils.py +486 -0
  617. mlflow/utils/name_utils.py +346 -0
  618. mlflow/utils/nfs_on_spark.py +62 -0
  619. mlflow/utils/openai_utils.py +164 -0
  620. mlflow/utils/os.py +12 -0
  621. mlflow/utils/oss_registry_utils.py +29 -0
  622. mlflow/utils/plugins.py +17 -0
  623. mlflow/utils/process.py +182 -0
  624. mlflow/utils/promptlab_utils.py +146 -0
  625. mlflow/utils/proto_json_utils.py +743 -0
  626. mlflow/utils/pydantic_utils.py +54 -0
  627. mlflow/utils/request_utils.py +279 -0
  628. mlflow/utils/requirements_utils.py +704 -0
  629. mlflow/utils/rest_utils.py +673 -0
  630. mlflow/utils/search_logged_model_utils.py +127 -0
  631. mlflow/utils/search_utils.py +2111 -0
  632. mlflow/utils/secure_loading.py +221 -0
  633. mlflow/utils/security_validation.py +384 -0
  634. mlflow/utils/server_cli_utils.py +61 -0
  635. mlflow/utils/spark_utils.py +15 -0
  636. mlflow/utils/string_utils.py +138 -0
  637. mlflow/utils/thread_utils.py +63 -0
  638. mlflow/utils/time.py +54 -0
  639. mlflow/utils/timeout.py +42 -0
  640. mlflow/utils/uri.py +572 -0
  641. mlflow/utils/validation.py +662 -0
  642. mlflow/utils/virtualenv.py +458 -0
  643. mlflow/utils/warnings_utils.py +25 -0
  644. mlflow/utils/yaml_utils.py +179 -0
  645. mlflow/version.py +24 -0
@@ -0,0 +1,704 @@
1
+ """
2
+ This module provides a set of utilities for interpreting and creating requirements files
3
+ (e.g. pip's `requirements.txt`), which is useful for managing ML software environments.
4
+ """
5
+
6
+ import importlib.metadata
7
+ import json
8
+ import logging
9
+ import os
10
+ import re
11
+ import subprocess
12
+ import sys
13
+ import tempfile
14
+ from collections import namedtuple
15
+ from itertools import chain, filterfalse
16
+ from pathlib import Path
17
+ from threading import Timer
18
+ from typing import NamedTuple, Optional
19
+
20
+ import importlib_metadata
21
+ from packaging.requirements import Requirement
22
+ from packaging.version import InvalidVersion, Version
23
+
24
+ import mlflow
25
+ from mlflow.environment_variables import (
26
+ _MLFLOW_IN_CAPTURE_MODULE_PROCESS,
27
+ MLFLOW_REQUIREMENTS_INFERENCE_RAISE_ERRORS,
28
+ MLFLOW_REQUIREMENTS_INFERENCE_TIMEOUT,
29
+ )
30
+ from mlflow.exceptions import MlflowException
31
+ from mlflow.tracking.artifact_utils import _download_artifact_from_uri
32
+ from mlflow.utils.autologging_utils.versioning import _strip_dev_version_suffix
33
+ from mlflow.utils.databricks_utils import (
34
+ get_databricks_env_vars,
35
+ is_in_databricks_runtime,
36
+ )
37
+
38
+ _logger = logging.getLogger(__name__)
39
+
40
+
41
+ def _is_comment(line):
42
+ return line.startswith("#")
43
+
44
+
45
+ def _is_empty(line):
46
+ return line == ""
47
+
48
+
49
+ def _strip_inline_comment(line):
50
+ return line[: line.find(" #")].rstrip() if " #" in line else line
51
+
52
+
53
+ def _is_requirements_file(line):
54
+ return line.startswith("-r ") or line.startswith("--requirement ")
55
+
56
+
57
+ def _is_constraints_file(line):
58
+ return line.startswith("-c ") or line.startswith("--constraint ")
59
+
60
+
61
+ def _join_continued_lines(lines):
62
+ """
63
+ Joins lines ending with '\\'.
64
+
65
+ >>> _join_continued_lines["a\\", "b\\", "c"]
66
+ >>> "abc"
67
+ """
68
+ continued_lines = []
69
+
70
+ for line in lines:
71
+ if line.endswith("\\"):
72
+ continued_lines.append(line.rstrip("\\"))
73
+ else:
74
+ continued_lines.append(line)
75
+ yield "".join(continued_lines)
76
+ continued_lines.clear()
77
+
78
+ # The last line ends with '\'
79
+ if continued_lines:
80
+ yield "".join(continued_lines)
81
+
82
+
83
+ # Represents a pip requirement.
84
+ #
85
+ # :param req_str: A requirement string (e.g. "scikit-learn == 0.24.2").
86
+ # :param is_constraint: A boolean indicating whether this requirement is a constraint.
87
+ _Requirement = namedtuple("_Requirement", ["req_str", "is_constraint"])
88
+
89
+
90
+ def _parse_requirements(requirements, is_constraint, base_dir=None):
91
+ """A simplified version of `pip._internal.req.parse_requirements` which performs the following
92
+ operations on the given requirements file and yields the parsed requirements.
93
+
94
+ - Remove comments and blank lines
95
+ - Join continued lines
96
+ - Resolve requirements file references (e.g. '-r requirements.txt')
97
+ - Resolve constraints file references (e.g. '-c constraints.txt')
98
+
99
+ Args:
100
+ requirements: A string path to a requirements file on the local filesystem or
101
+ an iterable of pip requirement strings.
102
+ is_constraint: Indicates the parsed requirements file is a constraint file.
103
+ base_dir: If specified, resolve relative file references (e.g. '-r requirements.txt')
104
+ against the specified directory.
105
+
106
+ Returns:
107
+ A list of ``_Requirement`` instances.
108
+
109
+ References:
110
+ - `pip._internal.req.parse_requirements`:
111
+ https://github.com/pypa/pip/blob/7a77484a492c8f1e1f5ef24eaf71a43df9ea47eb/src/pip/_internal/req/req_file.py#L118
112
+ - Requirements File Format:
113
+ https://pip.pypa.io/en/stable/cli/pip_install/#requirements-file-format
114
+ - Constraints Files:
115
+ https://pip.pypa.io/en/stable/user_guide/#constraints-files
116
+ """
117
+ if base_dir is None:
118
+ if isinstance(requirements, (str, Path)):
119
+ base_dir = os.path.dirname(requirements)
120
+ with open(requirements) as f:
121
+ requirements = f.read().splitlines()
122
+ else:
123
+ base_dir = os.getcwd()
124
+
125
+ lines = map(str.strip, requirements)
126
+ lines = map(_strip_inline_comment, lines)
127
+ lines = _join_continued_lines(lines)
128
+ lines = filterfalse(_is_comment, lines)
129
+ lines = filterfalse(_is_empty, lines)
130
+
131
+ for line in lines:
132
+ if _is_requirements_file(line):
133
+ req_file = line.split(maxsplit=1)[1]
134
+ # If `req_file` is an absolute path, `os.path.join` returns `req_file`:
135
+ # https://docs.python.org/3/library/os.path.html#os.path.join
136
+ abs_path = os.path.join(base_dir, req_file)
137
+ yield from _parse_requirements(abs_path, is_constraint=False)
138
+ elif _is_constraints_file(line):
139
+ req_file = line.split(maxsplit=1)[1]
140
+ abs_path = os.path.join(base_dir, req_file)
141
+ yield from _parse_requirements(abs_path, is_constraint=True)
142
+ else:
143
+ yield _Requirement(line, is_constraint)
144
+
145
+
146
+ def _flatten(iterable):
147
+ return chain.from_iterable(iterable)
148
+
149
+
150
+ # https://www.python.org/dev/peps/pep-0508/#names
151
+ _PACKAGE_NAME_REGEX = re.compile(r"^(\w+|\w+[\w._-]*\w+)")
152
+
153
+
154
+ def _get_package_name(requirement):
155
+ m = _PACKAGE_NAME_REGEX.match(requirement)
156
+ return m and m.group(1)
157
+
158
+
159
+ _NORMALIZE_REGEX = re.compile(r"[-_.]+")
160
+
161
+
162
+ def _normalize_package_name(pkg_name):
163
+ """
164
+ Normalizes a package name using the rule defined in PEP 503:
165
+ https://www.python.org/dev/peps/pep-0503/#normalized-names
166
+ """
167
+ return _NORMALIZE_REGEX.sub("-", pkg_name).lower()
168
+
169
+
170
+ def _iter_requires(name: str):
171
+ """
172
+ Iterates over the requirements of the specified package.
173
+
174
+ Args:
175
+ name: The name of the package.
176
+
177
+ Yields:
178
+ The names of the required packages.
179
+ """
180
+ try:
181
+ reqs = importlib.metadata.requires(name)
182
+ except importlib.metadata.PackageNotFoundError:
183
+ return
184
+
185
+ if reqs is None:
186
+ return
187
+
188
+ for req in reqs:
189
+ # Skip extra dependencies
190
+ semi_colon_idx = req.find(";")
191
+ if (semi_colon_idx != -1) and req[semi_colon_idx:].startswith("; extra =="):
192
+ continue
193
+
194
+ req = Requirement(req)
195
+ # Skip the requirement if the environment marker is not satisfied
196
+ if req.marker and not req.marker.evaluate():
197
+ continue
198
+
199
+ yield req.name
200
+
201
+
202
+ def _get_requires(pkg_name):
203
+ norm_pkg_name = _normalize_package_name(pkg_name)
204
+ for req in _iter_requires(norm_pkg_name):
205
+ yield _normalize_package_name(req)
206
+
207
+
208
+ def _get_requires_recursive(pkg_name, seen_before=None):
209
+ """
210
+ Recursively yields both direct and transitive dependencies of the specified
211
+ package.
212
+ """
213
+ norm_pkg_name = _normalize_package_name(pkg_name)
214
+ seen_before = seen_before or {norm_pkg_name}
215
+ for req in _get_requires(pkg_name):
216
+ # Prevent infinite recursion due to cyclic dependencies
217
+ if req in seen_before:
218
+ continue
219
+ seen_before.add(req)
220
+ yield req
221
+ yield from _get_requires_recursive(req, seen_before)
222
+
223
+
224
+ def _prune_packages(packages):
225
+ """
226
+ Prunes packages required by other packages. For example, `["scikit-learn", "numpy"]` is pruned
227
+ to `["scikit-learn"]`.
228
+ """
229
+ packages = set(packages)
230
+ requires = set(_flatten(map(_get_requires_recursive, packages)))
231
+
232
+ # LlamaIndex have one root "llama-index" package that bundles many sub-packages such as
233
+ # llama-index-llms-openai. Many of those sub-packages are optional, but some are defined
234
+ # as dependencies of the root package. However, the root package does not pin the versions
235
+ # for those sub-packages, resulting in non-deterministic behavior when loading the model
236
+ # later. To address this issue, we keep all sub-packages within the requirements.
237
+ # Ref: https://github.com/run-llama/llama_index/issues/14788#issuecomment-2232107585
238
+ requires = {req for req in requires if not req.startswith("llama-index-")}
239
+
240
+ # Do not exclude mlflow's dependencies
241
+ # Do not exclude databricks-connect since it conflicts with pyspark during execution time,
242
+ # and we need to determine if pyspark needs to be stripped based on the inferred packages
243
+ return packages - (requires - set(_get_requires("mlflow")) - {"databricks-connect"})
244
+
245
+
246
+ def _run_command(cmd, timeout_seconds, env=None):
247
+ """
248
+ Runs the specified command. If it exits with non-zero status, `MlflowException` is raised.
249
+ """
250
+ proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env)
251
+ timer = Timer(timeout_seconds, proc.kill)
252
+ try:
253
+ timer.start()
254
+ stdout, stderr = proc.communicate()
255
+ stdout = stdout.decode("utf-8")
256
+ stderr = stderr.decode("utf-8")
257
+ if proc.returncode != 0:
258
+ msg = "\n".join(
259
+ [
260
+ f"Encountered an unexpected error while running {cmd}",
261
+ f"exit status: {proc.returncode}",
262
+ f"stdout: {stdout}",
263
+ f"stderr: {stderr}",
264
+ ]
265
+ )
266
+ raise MlflowException(msg)
267
+ finally:
268
+ if timer.is_alive():
269
+ timer.cancel()
270
+
271
+
272
+ def _get_installed_version(package: str, module: Optional[str] = None) -> str:
273
+ """
274
+ Obtains the installed package version using `importlib_metadata.version`. If it fails, use
275
+ `__import__(module or package).__version__`.
276
+ """
277
+ if package == "mlflow":
278
+ # `importlib.metadata.version` may return an incorrect version of MLflow when it's
279
+ # installed in editable mode (e.g. `pip install -e .`).
280
+ return mlflow.__version__
281
+
282
+ try:
283
+ version = importlib_metadata.version(package)
284
+ except importlib_metadata.PackageNotFoundError:
285
+ # Note `importlib_metadata.version(package)` is not necessarily equal to
286
+ # `__import__(package).__version__`. See the example for pytorch below.
287
+ #
288
+ # Example
289
+ # -------
290
+ # $ pip install torch==1.9.0
291
+ # $ python -c "import torch; print(torch.__version__)"
292
+ # 1.9.0+cu102
293
+ # $ python -c "import importlib_metadata; print(importlib_metadata.version('torch'))"
294
+ # 1.9.0
295
+ version = __import__(module or package).__version__
296
+
297
+ # Strip the suffix from `dev` versions of PySpark, which are not available for installation
298
+ # from Anaconda or PyPI
299
+ if package == "pyspark":
300
+ version = _strip_dev_version_suffix(version)
301
+
302
+ return version
303
+
304
+
305
+ def _capture_imported_modules(model_uri, flavor, record_full_module=False, extra_env_vars=None):
306
+ """Runs `_capture_modules.py` in a subprocess and captures modules imported during the model
307
+ loading procedure.
308
+ If flavor is `transformers`, `_capture_transformers_modules.py` is run instead.
309
+
310
+ Args:
311
+ model_uri: The URI of the model.
312
+ flavor: The flavor name of the model.
313
+ record_full_module: Whether to capture top level modules for inferring python
314
+ package purpose. Default to False.
315
+ extra_env_vars: A dictionary of extra environment variables to pass to the subprocess.
316
+ Default to None.
317
+
318
+ Returns:
319
+ A list of captured modules.
320
+
321
+ """
322
+ local_model_path = _download_artifact_from_uri(model_uri)
323
+
324
+ process_timeout = MLFLOW_REQUIREMENTS_INFERENCE_TIMEOUT.get()
325
+ raise_on_error = MLFLOW_REQUIREMENTS_INFERENCE_RAISE_ERRORS.get()
326
+ extra_env_vars = extra_env_vars or {}
327
+
328
+ # Run `_capture_modules.py` to capture modules imported during the loading procedure
329
+ with tempfile.TemporaryDirectory() as tmpdir:
330
+ output_file = os.path.join(tmpdir, "imported_modules.txt")
331
+ # Pass the main environment variables to the subprocess for environment variable mapping
332
+ main_env = os.environ.copy()
333
+ # Reset the path variable from the main process so that the subprocess retains all
334
+ # main process configuration that a user has.
335
+ # See: ``https://github.com/mlflow/mlflow/issues/6905`` for context on minio configuration
336
+ # resolution in a subprocess based on PATH entries.
337
+ main_env["PATH"] = "/usr/sbin:/sbin:" + main_env["PATH"]
338
+ # Add databricks env, for langchain models loading we might need CLI configurations
339
+ if is_in_databricks_runtime():
340
+ main_env.update(get_databricks_env_vars(mlflow.get_tracking_uri()))
341
+
342
+ record_full_module_args = ["--record-full-module"] if record_full_module else []
343
+
344
+ if flavor == mlflow.transformers.FLAVOR_NAME:
345
+ # Lazily import `_capture_transformers_module` here to avoid circular imports.
346
+ from mlflow.utils import _capture_transformers_modules
347
+
348
+ for module_to_throw in ["tensorflow", "torch"]:
349
+ # NB: Setting USE_TF or USE_TORCH here as Transformers only checks these env
350
+ # variable on the first import of the library, which could happen anytime during
351
+ # the model loading process (or even mlflow import). When these variables are not
352
+ # set, Transformers import some torch/tensorflow modules even if they are not
353
+ # used by the model, resulting in false positives in the captured modules.
354
+ transformer_env = (
355
+ {"USE_TF": "TRUE"} if module_to_throw == "torch" else {"USE_TORCH": "TRUE"}
356
+ )
357
+ try:
358
+ _run_command(
359
+ [
360
+ sys.executable,
361
+ _capture_transformers_modules.__file__,
362
+ "--model-path",
363
+ local_model_path,
364
+ "--flavor",
365
+ flavor,
366
+ "--output-file",
367
+ output_file,
368
+ "--sys-path",
369
+ json.dumps(sys.path),
370
+ "--module-to-throw",
371
+ module_to_throw,
372
+ *record_full_module_args,
373
+ ],
374
+ timeout_seconds=process_timeout,
375
+ env={
376
+ **main_env,
377
+ **transformer_env,
378
+ _MLFLOW_IN_CAPTURE_MODULE_PROCESS.name: "true",
379
+ **extra_env_vars,
380
+ },
381
+ )
382
+ with open(output_file) as f:
383
+ return f.read().splitlines()
384
+
385
+ except MlflowException:
386
+ pass
387
+
388
+ # Lazily import `_capture_module` here to avoid circular imports.
389
+ from mlflow.utils import _capture_modules
390
+
391
+ error_file = os.path.join(tmpdir, "error.txt")
392
+ _run_command(
393
+ [
394
+ sys.executable,
395
+ _capture_modules.__file__,
396
+ "--model-path",
397
+ local_model_path,
398
+ "--flavor",
399
+ flavor,
400
+ "--output-file",
401
+ output_file,
402
+ "--error-file",
403
+ error_file,
404
+ "--sys-path",
405
+ json.dumps(sys.path),
406
+ *record_full_module_args,
407
+ ],
408
+ timeout_seconds=process_timeout,
409
+ env={
410
+ **main_env,
411
+ _MLFLOW_IN_CAPTURE_MODULE_PROCESS.name: "true",
412
+ **extra_env_vars,
413
+ },
414
+ )
415
+
416
+ if os.path.exists(error_file):
417
+ with open(error_file) as f:
418
+ errors = f.read()
419
+ if errors:
420
+ if raise_on_error:
421
+ raise MlflowException(
422
+ f"Encountered an error while capturing imported modules: {errors}"
423
+ )
424
+ _logger.warning(errors)
425
+
426
+ with open(output_file) as f:
427
+ return f.read().splitlines()
428
+
429
+
430
+ DATABRICKS_MODULES_TO_PACKAGES = {
431
+ "databricks.automl": ["databricks-automl-runtime"],
432
+ "databricks.automl_runtime": ["databricks-automl-runtime"],
433
+ "databricks.model_monitoring": ["databricks-model-monitoring"],
434
+ }
435
+ MLFLOW_MODULES_TO_PACKAGES = {
436
+ "mlflow.gateway": ["mlflow[gateway]"],
437
+ }
438
+ _MODULES_TO_PACKAGES = None
439
+ _PACKAGES_TO_MODULES = None
440
+
441
+
442
+ def _init_modules_to_packages_map():
443
+ global _MODULES_TO_PACKAGES
444
+ if _MODULES_TO_PACKAGES is None:
445
+ # Note `importlib_metadata.packages_distributions` only captures packages installed into
446
+ # Python's site-packages directory via tools such as pip:
447
+ # https://importlib-metadata.readthedocs.io/en/latest/using.html#using-importlib-metadata
448
+ _MODULES_TO_PACKAGES = importlib_metadata.packages_distributions()
449
+
450
+ # Add mapping for MLflow extras
451
+ _MODULES_TO_PACKAGES.update(MLFLOW_MODULES_TO_PACKAGES)
452
+
453
+ # Multiple packages populate the `databricks` module namespace on Databricks; to avoid
454
+ # bundling extraneous Databricks packages into model dependencies, we scope each module
455
+ # to its relevant package
456
+ _MODULES_TO_PACKAGES.update(DATABRICKS_MODULES_TO_PACKAGES)
457
+ if "databricks" in _MODULES_TO_PACKAGES:
458
+ _MODULES_TO_PACKAGES["databricks"] = [
459
+ package
460
+ for package in _MODULES_TO_PACKAGES["databricks"]
461
+ if package not in _flatten(DATABRICKS_MODULES_TO_PACKAGES.values())
462
+ ]
463
+
464
+ # In Databricks, `_MODULES_TO_PACKAGES` doesn't contain pyspark since it's not installed
465
+ # via pip or conda. To work around this issue, manually add pyspark.
466
+ if is_in_databricks_runtime():
467
+ _MODULES_TO_PACKAGES.update({"pyspark": ["pyspark"]})
468
+
469
+
470
+ def _init_packages_to_modules_map():
471
+ _init_modules_to_packages_map()
472
+ global _PACKAGES_TO_MODULES
473
+ _PACKAGES_TO_MODULES = {}
474
+ for module, pkg_list in _MODULES_TO_PACKAGES.items():
475
+ for pkg_name in pkg_list:
476
+ _PACKAGES_TO_MODULES[pkg_name] = module
477
+
478
+
479
+ def _infer_requirements(model_uri, flavor, raise_on_error=False, extra_env_vars=None):
480
+ """Infers the pip requirements of the specified model by creating a subprocess and loading
481
+ the model in it to determine which packages are imported.
482
+
483
+ Args:
484
+ model_uri: The URI of the model.
485
+ flavor: The flavor name of the model.
486
+ raise_on_error: If True, raise an exception if an unrecognized package is encountered.
487
+ extra_env_vars: A dictionary of extra environment variables to pass to the subprocess.
488
+ Default to None.
489
+
490
+ Returns:
491
+ A list of inferred pip requirements.
492
+
493
+ """
494
+ _init_modules_to_packages_map()
495
+
496
+ modules = _capture_imported_modules(model_uri, flavor, extra_env_vars=extra_env_vars)
497
+ packages = _flatten([_MODULES_TO_PACKAGES.get(module, []) for module in modules])
498
+ packages = map(_normalize_package_name, packages)
499
+ packages = _prune_packages(packages)
500
+ excluded_packages = [
501
+ # Certain packages (e.g. scikit-learn 0.24.2) imports `setuptools` or `pkg_resources`
502
+ # (a module provided by `setuptools`) to process or interact with package metadata.
503
+ # It should be safe to exclude `setuptools` because it's rare to encounter a python
504
+ # environment where `setuptools` is not pre-installed.
505
+ "setuptools",
506
+ # Exclude a package that provides the mlflow module (e.g. mlflow, mlflow-skinny).
507
+ # Certain flavors (e.g. pytorch) import mlflow while loading a model, but mlflow should
508
+ # not be counted as a model requirement.
509
+ *_MODULES_TO_PACKAGES.get("mlflow", []),
510
+ ]
511
+ packages = packages - set(excluded_packages)
512
+
513
+ # Handle pandas incompatibility issue with numpy 2.x https://github.com/pandas-dev/pandas/issues/55519
514
+ # pandas == 2.2.*: compatible with numpy >= 2
515
+ # pandas >= 2.1.2: incompatible with numpy >= 2, but it pins numpy < 2
516
+ # pandas < 2.1.2: incompatible with numpy >= 2 and doesn't pin numpy, so we need to pin numpy
517
+ if any(
518
+ package == "pandas"
519
+ and Version(_get_pinned_requirement(package).split("==")[1]) < Version("2.1.2")
520
+ for package in packages
521
+ ):
522
+ packages.add("numpy")
523
+
524
+ return sorted(map(_get_pinned_requirement, packages))
525
+
526
+
527
+ def _get_local_version_label(version):
528
+ """Extracts a local version label from `version`.
529
+
530
+ Args:
531
+ version: A version string.
532
+ """
533
+ try:
534
+ return Version(version).local
535
+ except InvalidVersion:
536
+ return None
537
+
538
+
539
+ def _strip_local_version_label(version):
540
+ """Strips a local version label in `version`.
541
+
542
+ Local version identifiers:
543
+ https://www.python.org/dev/peps/pep-0440/#local-version-identifiers
544
+
545
+ Args:
546
+ version: A version string to strip.
547
+ """
548
+
549
+ class IgnoreLocal(Version):
550
+ @property
551
+ def local(self):
552
+ return None
553
+
554
+ try:
555
+ return str(IgnoreLocal(version))
556
+ except InvalidVersion:
557
+ return version
558
+
559
+
560
+ def _get_pinned_requirement(req_str, version=None, module=None):
561
+ """Returns a string representing a pinned pip requirement to install the specified package and
562
+ version (e.g. 'mlflow==1.2.3').
563
+
564
+ Args:
565
+ req_str: The package requirement string (e.g. "mlflow" or "mlflow[gateway]").
566
+ version: The version of the package. If None, defaults to the installed version.
567
+ module: The name of the top-level module provided by the package . For example,
568
+ if `package` is 'scikit-learn', `module` should be 'sklearn'. If None, defaults
569
+ to `package`.
570
+ extras: A list of extra names for the package.
571
+
572
+ """
573
+ req = Requirement(req_str)
574
+ package = req.name
575
+ if version is None:
576
+ version_raw = _get_installed_version(package, module)
577
+ local_version_label = _get_local_version_label(version_raw)
578
+ if local_version_label:
579
+ version = _strip_local_version_label(version_raw)
580
+ if not (is_in_databricks_runtime() and package in ("torch", "torchvision")):
581
+ msg = (
582
+ f"Found {package} version ({version_raw}) contains a local version label "
583
+ f"(+{local_version_label}). MLflow logged a pip requirement for this package "
584
+ f"as '{package}=={version}' without the local version label to make it "
585
+ "installable from PyPI. To specify pip requirements containing local version "
586
+ "labels, please use `conda_env` or `pip_requirements`."
587
+ )
588
+ _logger.warning(msg)
589
+
590
+ else:
591
+ version = version_raw
592
+
593
+ if req.extras:
594
+ return f"{package}[{','.join(req.extras)}]=={version}"
595
+ return f"{package}=={version}"
596
+
597
+
598
+ class _MismatchedPackageInfo(NamedTuple):
599
+ package_name: str
600
+ installed_version: Optional[str]
601
+ requirement: str
602
+
603
+ def __str__(self):
604
+ current_status = self.installed_version if self.installed_version else "uninstalled"
605
+ return f"{self.package_name} (current: {current_status}, required: {self.requirement})"
606
+
607
+
608
+ def _check_requirement_satisfied(requirement_str):
609
+ """
610
+ Checks whether the current python environment satisfies the given requirement if it is parsable
611
+ as a package name and a set of version specifiers, and returns a `_MismatchedPackageInfo`
612
+ object containing the mismatched package name, installed version, and requirement if the
613
+ requirement is not satisfied. Otherwise, returns None.
614
+ """
615
+ _init_packages_to_modules_map()
616
+ try:
617
+ req = Requirement(requirement_str)
618
+ except Exception:
619
+ # We reach here if the requirement string is a file path or a URL.
620
+ # Extracting the package name from the requirement string is not trivial,
621
+ # so we skip the check.
622
+ return
623
+ pkg_name = req.name
624
+
625
+ try:
626
+ installed_version = _get_installed_version(pkg_name, _PACKAGES_TO_MODULES.get(pkg_name))
627
+ except ModuleNotFoundError:
628
+ return _MismatchedPackageInfo(
629
+ package_name=pkg_name,
630
+ installed_version=None,
631
+ requirement=requirement_str,
632
+ )
633
+
634
+ if pkg_name == "mlflow" and "gateway" in req.extras:
635
+ try:
636
+ from mlflow import gateway # noqa: F401
637
+ except ModuleNotFoundError:
638
+ return _MismatchedPackageInfo(
639
+ package_name="mlflow[gateway]",
640
+ installed_version=None,
641
+ requirement=requirement_str,
642
+ )
643
+
644
+ if pkg_name == "mlflow" and Version(installed_version).is_devrelease:
645
+ return None
646
+
647
+ if len(req.specifier) > 0 and not req.specifier.contains(installed_version):
648
+ return _MismatchedPackageInfo(
649
+ package_name=pkg_name,
650
+ installed_version=installed_version,
651
+ requirement=requirement_str,
652
+ )
653
+
654
+ return None
655
+
656
+
657
+ def warn_dependency_requirement_mismatches(model_requirements: list[str]):
658
+ """
659
+ Inspects the model's dependencies and prints a warning if the current Python environment
660
+ doesn't satisfy them.
661
+ """
662
+ # Suppress databricks-feature-lookup warning for feature store cases
663
+ # Suppress databricks-chains, databricks-rag, and databricks-agents warnings for RAG
664
+ # Studio cases
665
+ # NB: When a final name has been decided for GA for the aforementioned
666
+ # "Databricks RAG Studio" product, remove unrelated names from this listing.
667
+ _DATABRICKS_FEATURE_LOOKUP = "databricks-feature-lookup"
668
+ _DATABRICKS_AGENTS = "databricks-agents"
669
+
670
+ # List of packages to ignore
671
+ packages_to_ignore = [
672
+ _DATABRICKS_FEATURE_LOOKUP,
673
+ _DATABRICKS_AGENTS,
674
+ ]
675
+
676
+ # Normalize package names and create ignore list
677
+ ignore_packages = list(map(_normalize_package_name, packages_to_ignore))
678
+
679
+ try:
680
+ mismatch_infos = []
681
+ for req in model_requirements:
682
+ mismatch_info = _check_requirement_satisfied(req)
683
+ if mismatch_info is not None:
684
+ if _normalize_package_name(mismatch_info.package_name) in ignore_packages:
685
+ continue
686
+ mismatch_infos.append(str(mismatch_info))
687
+
688
+ if len(mismatch_infos) > 0:
689
+ mismatch_str = " - " + "\n - ".join(mismatch_infos)
690
+ warning_msg = (
691
+ "Detected one or more mismatches between the model's dependencies and the current "
692
+ f"Python environment:\n{mismatch_str}\n"
693
+ "To fix the mismatches, call `mlflow.pyfunc.get_model_dependencies(model_uri)` "
694
+ "to fetch the model's environment and install dependencies using the resulting "
695
+ "environment file."
696
+ )
697
+ _logger.warning(warning_msg)
698
+
699
+ except Exception as e:
700
+ _logger.warning(
701
+ f"Encountered an unexpected error ({e!r}) while detecting model dependency "
702
+ "mismatches. Set logging level to DEBUG to see the full traceback."
703
+ )
704
+ _logger.debug("", exc_info=True)