diff --git a/mlflow-on-crusoe/.gitignore b/mlflow-on-crusoe/.gitignore new file mode 100644 index 0000000..c0b8dba --- /dev/null +++ b/mlflow-on-crusoe/.gitignore @@ -0,0 +1,6 @@ +.env +__pycache__/ +*.pyc +.DS_Store +mlruns/ +response.txt diff --git a/mlflow-on-crusoe/README.md b/mlflow-on-crusoe/README.md new file mode 100644 index 0000000..3bfef91 --- /dev/null +++ b/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) diff --git a/mlflow-on-crusoe/requirements.txt b/mlflow-on-crusoe/requirements.txt new file mode 100644 index 0000000..a947ada --- /dev/null +++ b/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 diff --git a/mlflow-on-crusoe/train_and_log.py b/mlflow-on-crusoe/train_and_log.py new file mode 100644 index 0000000..521ecf2 --- /dev/null +++ b/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")