MLflow John Snow Labs 模型集成：`mlflow.johnsnowlabs` Flavor 的日志、加载与部署实战指南

MLflow John Snow Labs 模型集成：mlflow.johnsnowlabsFlavor 的日志、加载与部署实战指南

【免费下载链接】mlflowThe open source AI engineering platform for agents, LLMs, and ML models. MLflow enables teams of all sizes to debug, evaluate, monitor, and optimize production-quality AI applications while controlling costs and managing access to models and data.项目地址: https://gitcode.com/GitHub_Trending/ml/mlflow

mlflow.johnsnowlabs是 MLflow 为 John Snow Labs 企业级 NLP 模型提供的官方模型 Flavor，用于在 MLflow 中记录（log）、保存（save）与加载（load）基于 Spark NLP / NLU 的模型。本指南以 mlflow.johnsnowlabs.rst 对应的 API 文档为骨架，结合 模块源码、官方示例 与 集成测试，完整讲解 License 环境变量的配置、log_model/save_model/load_model的核心用法、pyfunc 部署以及底层实现原理，帮助你快速把 John Snow Labs 的医疗、金融、法律等领域预训练模型纳入 MLflow 模型生命周期管理。

一、mlflow.johnsnowlabs是什么：双 Flavor 模型集成

mlflow.johnsnowlabs模块提供了针对Spark NLP 与 NLU 模型的日志与加载 API（模块入口见 mlflow/johnsnowlabs/init.py）。该模块每次导出模型时都会同时产生两个 Flavor：

Johnsnowlabs（原生）格式模型以 Spark MLlib 的 PipelineModel 持久化格式保存，可被加载为NluPipeline（底层即 Spark MLlib PipelineModel），用于在 Spark Session 中以 Spark Transformer 的方式做分布式打分。这是该模块的主 Flavor，总是被生成。

mlflow.pyfunc格式支持在 Spark 之外部署：加载时会实例化一个 SparkContext，并把输入数据转为 Spark DataFrame 后进行打分；同时也支持在 Spark 内以 Spark UDF 形式部署。带此 Flavor 的模型可以被加载为普通 Python 函数进行推理。该 Flavor 同样总是被生成。

借助这一 Flavor，你可以访问 John Snow Labs 提供的20000+ 覆盖 200+ 语言的 SOTA 企业级 NLP 模型（如 LLM、文本摘要、问答、命名实体识别、关系抽取、情感分析、拼写检查、图像分类、自动语音识别等），这些模型由 John Snow Labs 提供，需要其Enterprise NLP License（详见 docs/docs/classic-ml/model/index.mdx）。

二、前置条件：License 环境变量的配置

使用mlflow.johnsnowlabs前，必须配置 John Snow Labs 的许可证信息。License JSON 中必须包含以下四个 key：

配置方式：把原始 license.json 字符串写入名为JOHNSNOWLABS_LICENSE_JSON的环境变量：

import os import json # Write your raw license.json string into the 'JOHNSNOWLABS_LICENSE_JSON' env variable creds = { "AWS_ACCESS_KEY_ID": "...", "AWS_SECRET_ACCESS_KEY": "...", "SPARK_NLP_LICENSE": "...", "SECRET": "...", } os.environ["JOHNSNOWLABS_LICENSE_JSON"] = json.dumps(creds)

从源码看，_validate_env_vars()会在每次log_model/save_model/load_model调用前检查该环境变量，缺失时直接抛出Exception；随后_set_env_vars()会把 JSON 中的每一项解析后写入对应环境变量（见 mlflow/johnsnowlabs/init.py）。

此外，模块还涉及两个额外的环境变量（源码中的常量定义见 mlflow/johnsnowlabs/init.py）：

• HEALTHCARE_SECRET：用于拼接 spark-nlp-jsl（医疗版）wheel 的下载地址
• VISUAL_SECRET：用于拼接 spark-ocr（视觉版）wheel 的下载地址

在get_default_pip_requirements()中，若两个 secret 都未设置会抛出异常；设置哪个 secret，对应版本的 wheel 就会被追加进模型默认依赖列表（详见 mlflow/johnsnowlabs/init.py）。测试中还有一种更自动化的方式：通过JSL_ACCESS_KEY调用nlp.install(access_token=...)自动下载 License 与安装库，再把JslSecrets序列化为 JSON 写入环境变量（见 tests/johnsnowlabs/test_johnsnowlabs_model_export.py）。

三、用log_model记录模型到当前 Run

log_model接收一个通过nlp.load()创建的Johnsnowlabs NLUPipeline，将其作为 MLflow artifact 记录到当前 run，使用 MLlib 持久化格式，并生成带johnsnowlabsFlavor 的 MLflow Model。若当前没有活跃 run，它会自动创建一个 run 来获得 run_id。

完整示例（训练并记录一个分类器）

import os import json import pandas as pd import mlflow from johnsnowlabs import nlp # 1) Write your raw license.json string into the 'JOHNSNOWLABS_LICENSE_JSON' env variable creds = { "AWS_ACCESS_KEY_ID": "...", "AWS_SECRET_ACCESS_KEY": "...", "SPARK_NLP_LICENSE": "...", "SECRET": "...", } os.environ["JOHNSNOWLABS_LICENSE_JSON"] = json.dumps(creds) # 2) Download & Install Jars/Wheels if missing and Start a spark Session nlp.start() # 3) 加载可训练模型（trainable classifier） trainable_classifier = nlp.load("train.classifier") # 4) 构造示例训练数据 data = pd.DataFrame({ "text": ["I hate covid ", "I love covid"], "y": ["negative", "positive"], }) # 5) Fit 并得到训练好的分类器 trained_classifier = trainable_classifier.fit(data) trained_classifier.predict("He hates covid") # 6) Log it mlflow.johnsnowlabs.log_model(trained_classifier, name="my_trained_model")

（示例出处：mlflow/johnsnowlabs/init.py）

log_model核心参数说明

（参数说明详见 mlflow/johnsnowlabs/init.py）

直接保存到远程 Artifact 的优化路径

从源码看（mlflow/johnsnowlabs/init.py），log_model内部有一个优化分支：若当前 run 的 artifact URI 是远程地址，会尝试通过_maybe_save_model用 Spark 把模型直接写入 artifact 仓库；否则回退到Model.log()，经本地临时目录中转。Databricks 环境下还会根据运行环境决定是否使用mlflowdbfs协议直写（测试中对dbfs://、s3://、mlflowdbfs://多种 URI 组合均有断言，见 tests/johnsnowlabs/test_johnsnowlabs_model_export.py 被注释的历史用例）。

四、save_model：保存到本地路径

save_model将 Spark johnsnowlabs 模型保存到本地路径，默认使用 Spark MLlib 持久化机制：

from johnsnowlabs import nlp import mlflow import os import json # Write your raw license.json string into the 'JOHNSNOWLABS_LICENSE_JSON' env variable creds = { "AWS_ACCESS_KEY_ID": "...", "AWS_SECRET_ACCESS_KEY": "...", "SPARK_NLP_LICENSE": "...", "SECRET": "...", } os.environ["JOHNSNOWLABS_LICENSE_JSON"] = json.dumps(creds) # Download & Install Jars/Wheels if missing and Start a spark Session nlp.start() # load a model model = nlp.load("en.classify.bert_sequence.covid_sentiment") model.predict(["I hate covid", "I love covid"]) # Save model as pyfunc and johnsnowlabs format mlflow.johnsnowlabs.save_model(model, "saved_model") model = mlflow.johnsnowlabs.load_model("saved_model") # Predict with reloaded model model.predict(["I hate covid", "I love covid"])

（示例出处：mlflow/johnsnowlabs/init.py）

核心参数与log_model基本一致：spark_model可以是pyspark.ml.pipeline.PipelineModel或nlu.NLUPipeline（每个 johnsnowlabs 模型都是 PipelineModel，均可加载为 NLUPipeline）；path为本地保存路径；同样支持conda_env、code_paths、dfs_tmpdir、signature、input_example、pip_requirements、extra_pip_requirements、metadata、store_license。

保存时的底层处理流程

从源码可以梳理出save_model的实现链路（mlflow/johnsnowlabs/init.py）：

1. 先生成 DFS 临时路径，调用_unpack_and_save_model将模型写入临时目录：若传入的是PipelineModel直接用write().overwrite().save(dst)；若是 NLU pipe，会先执行一次predict("Init")初始化，再保存其vanilla_transformer_pipe（见 mlflow/johnsnowlabs/init.py）；
2. 若运行在 Databricks 集群且支持 DBFS FUSE，则用shutil.move把模型从 FUSE 路径移动到本地；否则通过_HadoopFileSystem.copy_to_local_file拷贝；
3. 最后调用_save_model_metadata写入MLmodel文件、conda/python 环境文件、requirements/constraints，以及模型的 jars 与（可选）license。

保存后的模型目录中，模型数据固定存放在jsl-model子目录下（常量_JOHNSNOWLABS_MODEL_PATH_SUB = "jsl-model"），jars 与 license 存放在jsl-model/jars.jsl目录（_save_jars_and_lic，见 mlflow/johnsnowlabs/init.py）。

五、load_model：加载并做推理

load_model从指定 URI 加载 johnsnowlabs MLflow 模型，支持多种 URI 形式：

• 本地路径：/Users/me/path/to/local/model、relative/path/to/local/model
• 对象存储：s3://my_bucket/path/to/model
• MLflow run：runs:/<mlflow_run_id>/run-relative/path/to/model
• 模型注册表：models:/<model_name>/<model_version>、models:/<model_name>/<stage>

import mlflow from johnsnowlabs import nlp import os import json # Write your raw license.json string into the 'JOHNSNOWLABS_LICENSE_JSON' env variable creds = { "AWS_ACCESS_KEY_ID": "...", "AWS_SECRET_ACCESS_KEY": "...", "SPARK_NLP_LICENSE": "...", "SECRET": "...", } os.environ["JOHNSNOWLABS_LICENSE_JSON"] = json.dumps(creds) # start a spark session nlp.start() # Load your MLflow Model model = mlflow.johnsnowlabs.load_model("johnsnowlabs_model") # Make predictions on test documents prediction = model.transform(["I love Covid", "I hate Covid"])

（示例出处：mlflow/johnsnowlabs/init.py）

返回值为nlu.NLUPipeline。参数方面，dfs_tmpdir控制加载时使用的 DFS/本地临时目录（默认/tmp/mlflow，对应环境变量MLFLOW_DFS_TMP，定义见 mlflow/environment_variables.py）；dst_path可指定本地下载目录（必须已存在，否则自动创建临时路径）。

加载实现要点

• load_model会先解析 root URI 与 artifact 路径、读取johnsnowlabsFlavor 配置，并把模型下载到本地（mlflow/johnsnowlabs/init.py）；
• 若运行在 Databricks 集群且有 DBFS FUSE，会走_load_model_databricks：先把模型拷贝到 FUSE 临时目录（因集群开启 passthrough 时shutil.copytree拷贝目录权限位会报 permission-denied，故使用shutil_copytree_without_file_permissions），再nlp.load（见 mlflow/johnsnowlabs/init.py）；
• 若使用mlflowdbfs协议，则直接PipelineModel.load(mlflowdbfs_path)；
• 无活跃 SparkSession 时，_get_or_create_sparksession会基于模型目录内打包的 jars 和 license 自动启动一个 Spark Session（_fetch_deps_from_path从jars.jsl目录收集.jar与license.json，见 mlflow/johnsnowlabs/init.py）。

六、通过 pyfunc 部署：本地推理与 Spark UDF

由于每次导出都同时生成pyfuncFlavor，模型可以脱离 johnsnowlabs API 直接部署。官方示例 examples/johnsnowlabs/export.py 给出了完整流程：

import json import os import pandas as pd from johnsnowlabs import nlp import mlflow from mlflow.pyfunc import spark_udf # 1) Write your raw license.json string into the 'JOHNSNOWLABS_LICENSE_JSON' env variable for MLflow creds = { "AWS_ACCESS_KEY_ID": "...", "AWS_SECRET_ACCESS_KEY": "...", "SPARK_NLP_LICENSE": "...", "SECRET": "...", } os.environ["JOHNSNOWLABS_LICENSE_JSON"] = json.dumps(creds) # 2) Install enterprise libraries nlp.install() # 3) Start a Spark session with enterprise libraries spark = nlp.start() # 4) Load a model and test it nlu_model = "en.classify.bert_sequence.covid_sentiment" model_save_path = "my_model" johnsnowlabs_model = nlp.load(nlu_model) johnsnowlabs_model.predict(["I hate COVID,", "I love COVID"]) # 5) Export model with pyfunc and johnsnowlabs flavors with mlflow.start_run(): model_info = mlflow.johnsnowlabs.log_model(johnsnowlabs_model, name=model_save_path) # 6) Load model with johnsnowlabs flavor mlflow.johnsnowlabs.load_model(model_info.model_uri) # 7) Load model with pyfunc flavor mlflow.pyfunc.load_model(model_save_path) # 8) Deploy as Spark UDF pandas_df = pd.DataFrame({"text": ["Hello World"]}) spark_df = spark.createDataFrame(pandas_df).coalesce(1) pyfunc_udf = spark_udf( spark=spark, model_uri=model_save_path, env_manager="virtualenv", result_type="string", ) new_df = spark_df.withColumn("prediction", pyfunc_udf(*pandas_df.columns))

加载后还可以直接使用mlflow models serve命令把模型作为 REST 服务对外提供。

pyfunc 的底层实现是_PyFuncModelWrapper（见 mlflow/johnsnowlabs/init.py）：它包装 NLUPipeline 与 SparkSession，predict方法接收 pandas DataFrame 并返回 JSON 格式的预测结果；get_raw_model()可拿到底层原始模型。_load_pyfunc作为pyfunc.load_model的回调入口，当 pyfunc 作为 UDF 在工作节点执行时，会显式传入 SparkContext（因为工作节点上无法自行获取），这正是spark_udf场景下必须传入spark的原因（见 mlflow/johnsnowlabs/init.py）。

七、环境与依赖管理

get_default_pip_requirements()返回该 Flavor 导出的模型默认 pip 依赖（源码见 mlflow/johnsnowlabs/init.py），最小包含：

• johnsnowlabs_for_databricks==<version>（版本取自 johnsnowlabs 库自身设置）
• 固定版本的pyspark（经_get_pinned_requirement处理）
• 设置HEALTHCARE_SECRET时追加spark-nlp-jsl医疗版 wheel（https://pypi.johnsnowlabs.com/<secret>/spark-nlp-jsl/...）
• 设置VISUAL_SECRET时追加spark-ocr视觉版 wheel（https://pypi.johnsnowlabs.com/<secret>/spark-ocr/...）

get_default_conda_env()基于上述依赖生成默认 Conda 环境。若自定义conda_env，至少应包含默认环境中的依赖，例如：

{ 'name': 'mlflow-env', 'channels': ['defaults'], 'dependencies': [ 'python=3.8.15', 'johnsnowlabs' ] }

保存模型时，_save_model_metadata会把环境统一序列化为conda.yaml、requirements.txt、constraints.txt与python_env.yaml写入模型目录（见 mlflow/johnsnowlabs/init.py），保证后续部署时环境可复现。测试test_model_export还验证了保存/加载后原模型与重载模型（含原生与 pyfunc 两种方式、UDF 打分）预测结果的一致性（见 tests/johnsnowlabs/test_johnsnowlabs_model_export.py）。

八、常见使用要点与注意事项

1. License 先行：所有导出/加载操作前必须设置JOHNSNOWLABS_LICENSE_JSON，否则直接抛异常；医疗/视觉模型还需HEALTHCARE_SECRET/VISUAL_SECRET。
2. store_license=True才能随模型携带 License：默认 License 不随模型存储，跨环境加载时需在目标环境重新配置环境变量；开启后 License 会写入jsl-model/jars.jsl/license.json，加载时自动读取。
3. Spark Session 自动管理：加载模型时若没有活跃 SparkSession，会自动用模型自带 jars 启动一个；作为 UDF 部署时必须显式传入 SparkContext。
4. 签名与示例：建议通过infer_signature与input_example记录模型签名和示例，便于在 MLflow UI 中预览与校验（对应测试见 tests/johnsnowlabs/test_johnsnowlabs_model_export.py）。
5. 环境变量MLFLOW_DFS_TMP：控制 DFS 临时目录，默认/tmp/mlflow，可在大规模集群场景下按需调整（定义见 mlflow/environment_variables.py）。

九、总结

mlflow.johnsnowlabs通过"johnsnowlabs 原生 + pyfunc"双 Flavor 设计，把 John Snow Labs 的 20000+ 企业级 NLP 模型无缝接入 MLflow 的模型记录、注册、版本管理与部署体系：原生 Flavor 保留 Spark 分布式能力，pyfunc Flavor 让模型可以在任意 MLflow 推理环境（本地、Serving、Spark UDF）中运行。配置好 License 环境变量后，从nlp.load()到log_model/save_model/load_model，再到spark_udf与mlflow models serve部署，即可完成完整的企业级 NLP 模型生命周期管理。若需继续深入，可阅读 模块源码、官方导出示例 与 模型 Flavor 集成测试。

【免费下载链接】mlflowThe open source AI engineering platform for agents, LLMs, and ML models. MLflow enables teams of all sizes to debug, evaluate, monitor, and optimize production-quality AI applications while controlling costs and managing access to models and data.项目地址: https://gitcode.com/GitHub_Trending/ml/mlflow