[BUG] Cannot save prompt with partial variables.

[minkj1992]

Issues Policy acknowledgement

• I have read and agree to submit bug reports in accordance with the issues policy

Where did you encounter this bug?

Local machine

Willingness to contribute

Yes. I can contribute a fix for this bug independently.

MLflow version

mlflow --version
mlflow, version 2.12.1

System information

system_profiler SPSoftwareDataType SPHardwareDataType

Software:

System Software Overview:

  System Version: macOS 14.4.1 (23E224)
  Kernel Version: Darwin 23.4.0
  Secure Virtual Memory: Enabled
  System Integrity Protection: Enabled

Hardware:

Hardware Overview:
  Model Name: MacBook Pro
  Chip: Apple M1 Max
  System Firmware Version: 10151.101.3
  OS Loader Version: 10151.101.3

python --version
Python 3.12.2

Describe the problem

When setting up a chain with LCEL as described below and using partial_prompt as the prompt, I encountered the error MlflowException: Failed to save the runnable sequence.
It seems that the saving process is being attempted before the prompt is evaluated.

Tracking information

REPLACE_ME

Code to reproduce issue

import mlflow
from langchain_community.llms import Ollama
from langchain_core.output_parsers import PydanticOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.pydantic_v1 import BaseModel, Field

coherence_template = """You are the judge evaluating coherence/logic when you receive the debate topic and the content of the discussion.

...skip...

## Input
  TOPIC: {topic},
  Debate: {debate}

Below are the criteria for deducting points based on coherence/logic.

## Output Format

{format_instructions}
"""

topic = "Mocking topic"
debate = "Mocking debate"


class LLMResult(BaseModel):
    a_score: int = Field(description="Score of Team A")
    a_reason: str = Field(description="Reason for deduction for Team A")
    b_score: int = Field(description="Score of Team B")
    b_reason: str = Field(description="Reason for deduction for Team B")


llm = Ollama(model="llama3", temperature=0, verbose=True)
parser = PydanticOutputParser(pydantic_object=LLMResult)
coherence_prompt = PromptTemplate(
    template=coherence_template,
    input_variables=[
        "topic",
        "debate",
    ],
    partial_variables={"format_instructions": parser.get_format_instructions()},
)

coherence_chain = coherence_prompt | llm | parser

mlflow.set_experiment("Test")
with mlflow.start_run():
    model_info = mlflow.langchain.log_model(coherence_chain, "langchain_model")

Stack trace

  warnings.warn(
Traceback (most recent call last):
  File "/Users/minwook/code/personal/eval/main.py", line 48, in <module>
    model_info = mlflow.langchain.log_model(coherence_chain, "langchain_model")
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/mlops/lib/python3.12/site-packages/mlflow/langchain/__init__.py", line 519, in log_model
    return Model.log(
           ^^^^^^^^^^
  File "/opt/miniconda3/envs/mlops/lib/python3.12/site-packages/mlflow/models/model.py", line 625, in log
    flavor.save_model(path=local_path, mlflow_model=mlflow_model, **kwargs)
  File "/opt/miniconda3/envs/mlops/lib/python3.12/site-packages/mlflow/langchain/__init__.py", line 311, in save_model
    model_data_kwargs = _save_model(lc_model, path, loader_fn, persist_dir)
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/mlops/lib/python3.12/site-packages/mlflow/langchain/__init__.py", line 549, in _save_model
    return _save_runnables(model, path, loader_fn=loader_fn, persist_dir=persist_dir)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/miniconda3/envs/mlops/lib/python3.12/site-packages/mlflow/langchain/runnables.py", line 400, in _save_runnables
    _save_runnable_with_steps(
  File "/opt/miniconda3/envs/mlops/lib/python3.12/site-packages/mlflow/langchain/runnables.py", line 321, in _save_runnable_with_steps
    raise MlflowException(f"Failed to save runnable sequence: {unsaved_runnables}.")
mlflow.exceptions.MlflowException: Failed to save runnable sequence: {'0': 'input_variables=[\'debate\', \'topic\'] partial_variables={\'format_instructions\': \'The output should be formatted as a JSON instance that conforms to the JSON schema below.\\n\\nAs an example, for the schema {"properties": {"foo": {"title": "Foo", "description": "a list of strings", "type": "array", "items": {"type": "string"}}}, "required": ["foo"]}\\nthe object {"foo": ["bar", "baz"]} is a well-formatted instance of the schema. The object {"properties": {"foo": ["bar", "baz"]}} is not well-formatted.\\n\\nHere is the output schema:\\n```\\n{"properties": {"a_score": {"title": "A Score", "description": "Score of Team A", "type": "integer"}, "a_reason": {"title": "A Reason", "description": "Reason for deduction for Team A", "type": "string"}, "b_score": {"title": "B Score", "description": "Score of Team B", "type": "integer"}, "b_reason": {"title": "B Reason", "description": "Reason for deduction for Team B", "type": "string"}}, "required": ["a_score", "a_reason", "b_score", "b_reason"]}\\n```\'} template=\'You are the judge evaluating coherence/logic when you receive the debate topic and the content of the discussion.\\n\\n...skip...\\n\\n## Input\\n  TOPIC: {topic},\\n  Debate: {debate}\\n\\nBelow are the criteria for deducting points based on coherence/logic.\\n\\n## Output Format\\n\\n{format_instructions}\\n\' -- Cannot save prompt with partial variables.'}.

Other info / logs

REPLACE_ME

What component(s) does this bug affect?

• area/artifacts: Artifact stores and artifact logging
• area/build: Build and test infrastructure for MLflow
• area/deployments: MLflow Deployments client APIs, server, and third-party Deployments integrations
• area/docs: MLflow documentation pages
• area/examples: Example code
• area/model-registry: Model Registry service, APIs, and the fluent client calls for Model Registry
• area/models: MLmodel format, model serialization/deserialization, flavors
• area/recipes: Recipes, Recipe APIs, Recipe configs, Recipe Templates
• area/projects: MLproject format, project running backends
• area/scoring: MLflow Model server, model deployment tools, Spark UDFs
• area/server-infra: MLflow Tracking server backend
• area/tracking: Tracking Service, tracking client APIs, autologging

What interface(s) does this bug affect?

• area/uiux: Front-end, user experience, plotting, JavaScript, JavaScript dev server
• area/docker: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
• area/sqlalchemy: Use of SQLAlchemy in the Tracking Service or Model Registry
• area/windows: Windows support

What language(s) does this bug affect?

• language/r: R APIs and clients
• language/java: Java APIs and clients
• language/new: Proposals for new client languages

What integration(s) does this bug affect?

• integrations/azure: Azure and Azure ML integrations
• integrations/sagemaker: SageMaker integrations
• integrations/databricks: Databricks integrations

[WeichenXu123]

CC @serena-ruan

[serena-ruan]

@minkj1992 We can try cloudpickle on PromptTemplate, as langchain doesn't natively support saving it with partial variables. If you're willing to contribute, could you add PromptTemplate into picklable_runnable_types and add a test to see if it works?

[github-actions]

@mlflow/mlflow-team Please assign a maintainer and start triaging this issue.