feat: MLflow experiment tracking for Crusoe Managed Inference

mlflow-on-crusoe/.gitignore

@@ -0,0 +1,6 @@
+.env
+__pycache__/
+*.pyc
+.DS_Store
+mlruns/
+response.txt

mlflow-on-crusoe/README.md

@@ -0,0 +1,74 @@
+# MLflow × Crusoe AI
+
+Track LLM experiments on [Crusoe Managed Inference](https://www.crusoe.ai/cloud/managed-inference) using MLflow — log parameters, latency metrics, and model outputs across runs.
+
+## What this does
+
+Runs 3 LLM inference experiments and logs everything to MLflow:
+
+- **Parameters** — model name, temperature, max tokens, prompt
+- **Metrics** — latency, output word count, words per second
+- **Artifacts** — full model response saved per run
+
+## Prerequisites
+
+- Python 3.10+
+- A Crusoe Cloud account → [console.crusoecloud.com](https://console.crusoecloud.com)
+- Inference API key (Intelligence API Keys section under Security)
+
+## Setup
+
+```bash
+pip install -r requirements.txt
+export CRUSOE_API_KEY="your-api-key"
+```
+
+## Run the experiments
+
+```bash
+python train_and_log.py
+```
+
+## View results in MLflow UI
+
+```bash
+mlflow ui
+```
+
+Then open http://localhost:5000 — you'll see all runs under the `crusoe-llm-experiments` experiment with latency and throughput metrics side by side.
+
+## Local testing (no Crusoe account needed)
+
+The script automatically falls back to Groq if `CRUSOE_API_KEY` is not set:
+
+```bash
+pip install langchain-groq
+export GROQ_API_KEY="your-groq-key"  # free at console.groq.com
+python train_and_log.py
+```
+
+## What gets logged per run
+
+| Type | Details |
+|------|---------|
+| Parameters | model, temperature, max_tokens, prompt |
+| Metrics | latency_seconds, output_word_count, words_per_second |
+| Artifacts | full response saved as response.txt |
+
+## Extend it
+
+Add your own prompts to the list in `train_and_log.py`:
+
+```python
+prompts = [
+    ("my_task", "Your prompt here"),
+]
+```
+
+Each entry becomes a named MLflow run, making it easy to compare outputs across models or temperatures.
+
+## Related
+
+- [langchain-crusoe](../langchain-crusoe/) — LangChain integration for Crusoe Managed Inference
+- [langgraph-crusoe](../langgraph-crusoe/) — Multi-node agentic pipelines on Crusoe
+- [Crusoe Managed Inference Docs](https://docs.crusoecloud.com/managed-inference/overview)

mlflow-on-crusoe/requirements.txt

@@ -0,0 +1,4 @@
+mlflow>=2.13.0
+langchain-crusoe>=0.1.0
+langchain-groq>=1.1.0
+python-dotenv

mlflow-on-crusoe/train_and_log.py

@@ -0,0 +1,86 @@
+"""
+MLflow experiment tracking with Crusoe Managed Inference.
+Logs model parameters, metrics, and LLM responses to MLflow.
+Tested locally with Groq as a drop-in replacement for Crusoe.
+"""
+import os
+import time
+import mlflow
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+def get_llm():
+    if os.getenv("CRUSOE_API_KEY"):
+        from langchain_crusoe import ChatCrusoe
+        return ChatCrusoe(
+            model="meta-llama/Llama-3.3-70B-Instruct",
+            temperature=0.3,
+            max_tokens=512,
+        )
+    else:
+        from langchain_groq import ChatGroq
+        return ChatGroq(
+            model="llama-3.3-70b-versatile",
+            temperature=0.3,
+            max_tokens=512,
+        )
+
+
+def run_experiment(prompt: str, temperature: float, max_tokens: int, run_name: str):
+    """Run a single LLM call and log everything to MLflow."""
+    mlflow.set_experiment("crusoe-llm-experiments")
+
+    with mlflow.start_run(run_name=run_name):
+        # Log parameters
+        mlflow.log_param("model", "Llama-3.3-70B-Instruct")
+        mlflow.log_param("temperature", temperature)
+        mlflow.log_param("max_tokens", max_tokens)
+        mlflow.log_param("prompt", prompt)
+
+        # Run inference
+        llm = get_llm()
+        start = time.time()
+        from langchain_core.messages import HumanMessage
+        response = llm.invoke([HumanMessage(content=prompt)])
+        latency = time.time() - start
+
+        output = response.content
+        word_count = len(output.split())
+
+        # Log metrics
+        mlflow.log_metric("latency_seconds", round(latency, 3))
+        mlflow.log_metric("output_word_count", word_count)
+        mlflow.log_metric("words_per_second", round(word_count / latency, 2))
+
+        # Log output as artifact
+        with open("response.txt", "w") as f:
+            f.write(f"PROMPT:\n{prompt}\n\nRESPONSE:\n{output}")
+        mlflow.log_artifact("response.txt")
+        os.remove("response.txt")
+
+        print(f"\nRun: {run_name}")
+        print(f"Latency: {latency:.2f}s | Words: {word_count} | WPS: {word_count/latency:.1f}")
+        print(f"Response preview: {output[:200]}...")
+
+        return {"latency": latency, "word_count": word_count, "output": output}
+
+
+if __name__ == "__main__":
+    prompts = [
+        ("summarization", "Summarize the key advantages of GPU cluster computing for AI workloads in 3 bullet points."),
+        ("reasoning",     "Explain the tradeoff between model quantization and inference accuracy."),
+        ("generation",    "Write a Python function that retries an API call up to 3 times with exponential backoff."),
+    ]
+
+    print("Starting MLflow experiment tracking with Crusoe Managed Inference...")
+    print("=" * 60)
+
+    for run_name, prompt in prompts:
+        run_experiment(prompt, temperature=0.3, max_tokens=512, run_name=run_name)
+
+    print("\n" + "=" * 60)
+    print("All runs logged. Launch the MLflow UI with:")
+    print("    mlflow ui")
+    print("Then open http://localhost:5000")