optimized tqdm to finish everything in a single loop

Author: jjmachan

nbs/project/experiments.ipynb

@@ -378,25 +378,50 @@
     "            if name_prefix:\n",
     "                name = f\"{name_prefix}-{name}\"\n",
     "\n",
-    "            # Create tasks for all items\n",
-    "            tasks = []\n",
-    "            for item in dataset:\n",
-    "                tasks.append(wrapped_experiment(item))\n",
-    "\n",
-    "            # Use as_completed with tqdm for progress tracking\n",
-    "            results = []\n",
-    "            for future in tqdm(asyncio.as_completed(tasks), total=len(tasks)):\n",
-    "                result = await future\n",
-    "                # Add each result to experiment view as it completes\n",
-    "                if result is not None:\n",
-    "                    results.append(result)\n",
-    "\n",
-    "            # upload results to experiment view\n",
-    "            experiment_view = self.create_experiment(name=name, model=experiment_model)\n",
-    "            for result in results:\n",
-    "                experiment_view.append(result)\n",
-    "\n",
-    "            return experiment_view\n",
+    "            experiment_view = None\n",
+    "            try:\n",
+    "                # Create the experiment view upfront\n",
+    "                experiment_view = self.create_experiment(name=name, model=experiment_model)\n",
+    "                \n",
+    "                # Create tasks for all items\n",
+    "                tasks = []\n",
+    "                for item in dataset:\n",
+    "                    tasks.append(wrapped_experiment(item))\n",
+    "\n",
+    "                # Calculate total operations (processing + appending)\n",
+    "                total_operations = len(tasks) * 2  # Each item requires processing and appending\n",
+    "                \n",
+    "                # Use tqdm for combined progress tracking\n",
+    "                results = []\n",
+    "                progress_bar = tqdm(total=total_operations, desc=\"Running experiment\")\n",
+    "                \n",
+    "                # Process all items\n",
+    "                for future in asyncio.as_completed(tasks):\n",
+    "                    result = await future\n",
+    "                    if result is not None:\n",
+    "                        results.append(result)\n",
+    "                    progress_bar.update(1)  # Update for task completion\n",
+    "                \n",
+    "                # Append results to experiment view\n",
+    "                for result in results:\n",
+    "                    experiment_view.append(result)\n",
+    "                    progress_bar.update(1)  # Update for append operation\n",
+    "                    \n",
+    "                progress_bar.close()\n",
+    "                return experiment_view\n",
+    "                \n",
+    "            except Exception as e:\n",
+    "                # Clean up the experiment if there was an error and it was created\n",
+    "                if experiment_view is not None:\n",
+    "                    try:\n",
+    "                        # Delete the experiment (you might need to implement this method)\n",
+    "                        sync_version = async_to_sync(self._ragas_api_client.delete_experiment)\n",
+    "                        sync_version(project_id=self.project_id, experiment_id=experiment_view.experiment_id)\n",
+    "                    except Exception as cleanup_error:\n",
+    "                        print(f\"Failed to clean up experiment after error: {cleanup_error}\")\n",
+    "                \n",
+    "                # Re-raise the original exception\n",
+    "                raise e\n",
     "\n",
     "        wrapped_experiment.__setattr__(\"run_async\", run_async)\n",
     "        return t.cast(ExperimentProtocol, wrapped_experiment)\n",
@@ -431,7 +456,6 @@
     "# create a test experiment function\n",
     "@p.experiment(TextExperimentModel)\n",
     "async def test_experiment(item: TestModel):\n",
-    "    print(item)\n",
     "    return TextExperimentModel(**item.model_dump(), response=\"test response\", is_correct=\"yes\")\n"
    ]
   },
@@ -444,22 +468,13 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "100%|██████████| 3/3 [00:00<00:00, 7752.87it/s]\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "name='test item 2' description='test item 2 description' price=200.0\n",
-      "name='test item 1' description='test item 1 description' price=100.0\n",
-      "name='test item 3' description='test item 3 description' price=300.0\n"
+      "Running experiment: 100%|██████████| 6/6 [00:01<00:00,  3.84it/s]\n"
      ]
     },
     {
      "data": {
       "text/plain": [
-       "Experiment(name=keen_backus, model=TextExperimentModel)"
+       "Experiment(name=gallant_torvalds, model=TextExperimentModel)"
       ]
      },
      "execution_count": null,

ragas_experimental/project/experiments.py

@@ -162,25 +162,50 @@ async def run_async(dataset: Dataset, name: t.Optional[str] = None):
             if name_prefix:
                 name = f"{name_prefix}-{name}"
 
-            # Create tasks for all items
-            tasks = []
-            for item in dataset:
-                tasks.append(wrapped_experiment(item))
-
-            # Use as_completed with tqdm for progress tracking
-            results = []
-            for future in tqdm(asyncio.as_completed(tasks), total=len(tasks)):
-                result = await future
-                # Add each result to experiment view as it completes
-                if result is not None:
-                    results.append(result)
-
-            # upload results to experiment view
-            experiment_view = self.create_experiment(name=name, model=experiment_model)
-            for result in results:
-                experiment_view.append(result)
-
-            return experiment_view
+            experiment_view = None
+            try:
+                # Create the experiment view upfront
+                experiment_view = self.create_experiment(name=name, model=experiment_model)
+                
+                # Create tasks for all items
+                tasks = []
+                for item in dataset:
+                    tasks.append(wrapped_experiment(item))
+
+                # Calculate total operations (processing + appending)
+                total_operations = len(tasks) * 2  # Each item requires processing and appending
+                
+                # Use tqdm for combined progress tracking
+                results = []
+                progress_bar = tqdm(total=total_operations, desc="Running experiment")
+                
+                # Process all items
+                for future in asyncio.as_completed(tasks):
+                    result = await future
+                    if result is not None:
+                        results.append(result)
+                    progress_bar.update(1)  # Update for task completion
+                
+                # Append results to experiment view
+                for result in results:
+                    experiment_view.append(result)
+                    progress_bar.update(1)  # Update for append operation
+                    
+                progress_bar.close()
+                return experiment_view
+                
+            except Exception as e:
+                # Clean up the experiment if there was an error and it was created
+                if experiment_view is not None:
+                    try:
+                        # Delete the experiment (you might need to implement this method)
+                        sync_version = async_to_sync(self._ragas_api_client.delete_experiment)
+                        sync_version(project_id=self.project_id, experiment_id=experiment_view.experiment_id)
+                    except Exception as cleanup_error:
+                        print(f"Failed to clean up experiment after error: {cleanup_error}")
+                
+                # Re-raise the original exception
+                raise e
 
         wrapped_experiment.__setattr__("run_async", run_async)
         return t.cast(ExperimentProtocol, wrapped_experiment)