Merge pull request #23 from explodinggradients/mlflow-git

feat: add git support for mlflow

Author: jjmachan

nbs/project/experiments.ipynb

@@ -909,7 +909,7 @@
     "\n",
     "@patch\n",
     "def mlflow_experiment(\n",
-    "    self: Project, experiment_model, name_prefix: str = \"\"\n",
+    "    self: Project, experiment_model, name_prefix: str = \"\",save_to_git: bool = True, stage_all: bool = True\n",
     "):\n",
     "    \"\"\"Decorator for creating experiment functions with mlflow integration.\n",
     "\n",
@@ -922,33 +922,19 @@
     "    \"\"\"\n",
     "\n",
     "    def decorator(func: t.Callable) -> ExperimentProtocol:\n",
-    "        # First, create a base experiment wrapper\n",
-    "        base_experiment = self.experiment(experiment_model, name_prefix)(func)\n",
-    "\n",
-    "        # Override the wrapped function to add mlflow observation\n",
+    "        \n",
     "        @wraps(func)\n",
-    "        async def wrapped_with_mlflow(*args, **kwargs):\n",
-    "            # wrap the function with mlflow observation\n",
-    "            observed_func = trace(name=f\"{name_prefix}-{func.__name__}\")(func)\n",
+    "        async def mlflow_wrapped_func(*args, **kwargs):\n",
+    "            # Apply mlflow observation directly here\n",
+    "            trace_name = f\"{name_prefix}-{func.__name__}\" if name_prefix else func.__name__\n",
+    "            observed_func = trace(name=trace_name)(func)\n",
     "            return await observed_func(*args, **kwargs)\n",
-    "\n",
-    "        # Replace the async function to use mlflow\n",
-    "        original_run_async = base_experiment.run_async\n",
-    "\n",
-    "        # Use the original run_async but with the mlflow-wrapped function\n",
-    "        async def run_async_with_mlflow(\n",
-    "            dataset: Dataset, name: t.Optional[str] = None\n",
-    "        ):\n",
-    "            # Override the internal wrapped_experiment with our mlflow version\n",
-    "            base_experiment.__wrapped__ = wrapped_with_mlflow\n",
-    "\n",
-    "            # Call the original run_async which will now use our mlflow-wrapped function\n",
-    "            return await original_run_async(dataset, name)\n",
-    "\n",
-    "        # Replace the run_async method\n",
-    "        base_experiment.__setattr__(\"run_async\", run_async_with_mlflow)\n",
-    "\n",
-    "        return t.cast(ExperimentProtocol, base_experiment)\n",
+    "        \n",
+    "        # Now create the experiment wrapper with our already-observed function\n",
+    "        experiment_wrapper = self.experiment(experiment_model, name_prefix, save_to_git, stage_all)(mlflow_wrapped_func)\n",
+    "        \n",
+    "        return t.cast(ExperimentProtocol, experiment_wrapper)\n",
+    "        \n",
     "\n",
     "    return decorator"
    ]

ragas_experimental/project/experiments.py

@@ -125,7 +125,7 @@ def get_experiment(self: Project, experiment_name: str, model) -> Dataset:
 
 # %% ../../nbs/project/experiments.ipynb 17
 def find_git_root(
-    start_path: t.Union[str, Path, None] = None,  # starting path to search from
+    start_path: t.Union[str, Path, None] = None  # starting path to search from
 ) -> Path:
     """Find the root directory of a git repository by traversing up from the start path."""
     # Start from the current directory if no path is provided
@@ -447,7 +447,13 @@ async def langfuse_wrapped_func(*args, **kwargs):
 
 
 @patch
-def mlflow_experiment(self: Project, experiment_model, name_prefix: str = ""):
+def mlflow_experiment(
+    self: Project,
+    experiment_model,
+    name_prefix: str = "",
+    save_to_git: bool = True,
+    stage_all: bool = True,
+):
     """Decorator for creating experiment functions with mlflow integration.
 
     Args:
@@ -459,31 +465,22 @@ def mlflow_experiment(self: Project, experiment_model, name_prefix: str = ""):
     """
 
     def decorator(func: t.Callable) -> ExperimentProtocol:
-        # First, create a base experiment wrapper
-        base_experiment = self.experiment(experiment_model, name_prefix)(func)
 
-        # Override the wrapped function to add mlflow observation
         @wraps(func)
-        async def wrapped_with_mlflow(*args, **kwargs):
-            # wrap the function with mlflow observation
-            observed_func = trace(name=f"{name_prefix}-{func.__name__}")(func)
+        async def mlflow_wrapped_func(*args, **kwargs):
+            # Apply mlflow observation directly here
+            trace_name = (
+                f"{name_prefix}-{func.__name__}" if name_prefix else func.__name__
+            )
+            observed_func = trace(name=trace_name)(func)
             return await observed_func(*args, **kwargs)
 
-        # Replace the async function to use mlflow
-        original_run_async = base_experiment.run_async
-
-        # Use the original run_async but with the mlflow-wrapped function
-        async def run_async_with_mlflow(dataset: Dataset, name: t.Optional[str] = None):
-            # Override the internal wrapped_experiment with our mlflow version
-            base_experiment.__wrapped__ = wrapped_with_mlflow
-
-            # Call the original run_async which will now use our mlflow-wrapped function
-            return await original_run_async(dataset, name)
-
-        # Replace the run_async method
-        base_experiment.__setattr__("run_async", run_async_with_mlflow)
+        # Now create the experiment wrapper with our already-observed function
+        experiment_wrapper = self.experiment(
+            experiment_model, name_prefix, save_to_git, stage_all
+        )(mlflow_wrapped_func)
 
-        return t.cast(ExperimentProtocol, base_experiment)
+        return t.cast(ExperimentProtocol, experiment_wrapper)
 
     return decorator