chore: import upstream snapshot with attribution

tests/pyfunc/sample_code/code_with_dependencies.py

@@ -0,0 +1,41 @@
+import os
+
+import mlflow
+from mlflow.models import set_model, set_retriever_schema
+from mlflow.pyfunc import PythonModel
+
+test_trace = os.environ.get("TEST_TRACE", "true").lower() == "true"
+
+
+class MyModel(PythonModel):
+    def _call_retriever(self, id):
+        return f"Retriever called with ID: {id}. Output: 42."
+
+    def predict(self, context, model_input):
+        return f"Input: {model_input}. {self._call_retriever(model_input)}"
+
+    def predict_stream(self, context, model_input, params=None):
+        yield f"Input: {model_input}. {self._call_retriever(model_input)}"
+
+
+class MyModelWithTrace(PythonModel):
+    def _call_retriever(self, id):
+        return f"Retriever called with ID: {id}. Output: 42."
+
+    @mlflow.trace
+    def predict(self, context, model_input):
+        return f"Input: {model_input}. {self._call_retriever(model_input)}"
+
+    @mlflow.trace
+    def predict_stream(self, context, model_input, params=None):
+        yield f"Input: {model_input}. {self._call_retriever(model_input)}"
+
+
+model = MyModelWithTrace() if test_trace else MyModel()
+set_model(model)
+set_retriever_schema(
+    primary_key="primary-key",
+    text_column="text-column",
+    doc_uri="doc-uri",
+    other_columns=["column1", "column2"],
+)

tests/pyfunc/sample_code/config.yml

@@ -0,0 +1,3 @@
+use_gpu: True
+temperature: 0.9
+timeout: 300

tests/pyfunc/sample_code/func_code.py

@@ -0,0 +1,8 @@
+from mlflow.models import set_model
+
+
+def predict(model_input):
+    return model_input
+
+
+set_model(predict)

tests/pyfunc/sample_code/func_code_with_config.py

@@ -0,0 +1,10 @@
+from mlflow.models import ModelConfig, set_model
+
+
+def predict(model_input: list[str]):
+    model_config = ModelConfig(development_config="tests/pyfunc/sample_code/config.yml")
+    timeout = model_config.get("timeout")
+    return f"This was the input: {model_input[0]}, timeout {timeout}"
+
+
+set_model(predict)

tests/pyfunc/sample_code/func_code_with_type_hint.py

@@ -0,0 +1,8 @@
+from mlflow.models import set_model
+
+
+def predict(model_input: list[str]):
+    return model_input
+
+
+set_model(predict)

tests/pyfunc/sample_code/python_model.py

@@ -0,0 +1,10 @@
+from mlflow.models import set_model
+from mlflow.pyfunc import PythonModel
+
+
+class MyModel(PythonModel):
+    def predict(self, context, model_input):
+        return f"This was the input: {model_input}"
+
+
+set_model(MyModel())

tests/pyfunc/sample_code/python_model_with_config.py

@@ -0,0 +1,13 @@
+from mlflow.models import ModelConfig, set_model
+from mlflow.pyfunc import PythonModel
+
+base_config = ModelConfig(development_config="tests/pyfunc/sample_code/config.yml")
+
+
+class MyModel(PythonModel):
+    def predict(self, context, model_input):
+        timeout = base_config.get("timeout")
+        return f"Predict called with input {model_input}, timeout {timeout}"
+
+
+set_model(MyModel())

tests/pyfunc/sample_code/python_model_with_utils.py

@@ -0,0 +1,12 @@
+from mlflow.models import set_model
+from mlflow.pyfunc import PythonModel
+
+
+class MyModel(PythonModel):
+    def predict(self, context, model_input):
+        from utils import my_function
+
+        return my_function(model_input)
+
+
+set_model(MyModel())

tests/pyfunc/sample_code/streamable_model_code.py

@@ -0,0 +1,17 @@
+from mlflow.models import set_model
+from mlflow.pyfunc import PythonModel
+
+
+class StreamableModel(PythonModel):
+    def __init__(self):
+        pass
+
+    def predict(self, context, model_input, params=None):
+        pass
+
+    def predict_stream(self, context, model_input, params=None):
+        yield "test1"
+        yield "test2"
+
+
+set_model(StreamableModel())

tests/pyfunc/sample_code/utils.py

@@ -0,0 +1,2 @@
+def my_function(input):
+    return f"My utils function received this input: {input}"