[fix] catch errors

Author: huyiwen

ml_scheduler/exp/exp.py

@@ -2,7 +2,7 @@
 import subprocess
 import threading
 from logging import getLogger
-from typing import Any, Dict, List, Set
+from typing import Any, Dict, List, Optional, Set
 
 from ..pools.base import BaseAllocator, BaseResources
 from .runner import BaseRunner
@@ -61,6 +61,8 @@ def run_in_thread(**popen_kwargs):
 
         return stdout
 
-    async def report(self, metrics: Dict[str, Any], **kwargs):
+    async def report(self, metrics: Optional[Dict[str, Any]] = None, **kwargs):
+        if metrics is None:
+            metrics = {}
         metrics.update(kwargs)
         await self.runner._report(self.uuid, metrics)

ml_scheduler/exp/func.py

@@ -1,15 +1,22 @@
 import inspect
+from logging import getLogger
+from traceback import format_exc
 from typing import Any, Tuple
 
 from .exp import Exp
 from .runner.csv import CSVRunner
 
+logger = getLogger(__name__)
+
 
 class ExpFunc:
 
     def __init__(self, exp_func) -> None:
         self.exp_func = exp_func
-        self.run_csv = CSVRunner.set(self).run
+
+        csvrunner = CSVRunner.set(self)
+        self.run_csv = csvrunner.run
+        self.arun_csv = csvrunner.arun
 
     async def __call__(self, exp: Exp, **kwargs) -> Tuple[Exp, Any]:
         assert isinstance(exp, Exp)
@@ -20,7 +27,13 @@ async def __call__(self, exp: Exp, **kwargs) -> Tuple[Exp, Any]:
         ba = inspect.signature(self.exp_func).bind(exp, **filtered_kwargs)
         ba.apply_defaults()
 
-        results = await self.exp_func(**ba.arguments)
+        try:
+            results = await self.exp_func(**ba.arguments)
+        except Exception as e:
+            logger.warning(format_exc())
+            logger.error(f"Error in {self.exp_func.__name__}: {e}")
+            results = ""
+
         await exp.cleanup()
         return (exp, results)
 

ml_scheduler/exp/runner/csv.py

@@ -22,6 +22,7 @@ def run(
         read_csv_kwargs: Optional[dict] = None,
         uuid_column: str = ":uuid:",
         retval_column: Optional[str] = ":retval:",
+        extra_kwargs: Optional[Dict[str, Any]] = None,
     ):
         """Run experiments from a csv file
 
@@ -32,14 +33,16 @@ def run(
             read_csv_kwargs (`Optional[dict]`, optional): Additional kwargs passed to `pandas.read_csv`. Defaults to None.
             uuid_column (`str`, optional): The column name for the uuid. Defaults to `":uuid:"`.
             retval_column (`Optional[str]`, optional): The column name for the return value. None for not saving the return value. Defaults to `":retval:"`.
+            extra_kwargs (`Optional[Dict[str, Any]]`, optional): Extra kwargs passed to exp_func.
         """
         kwargs = {
             "csv_path": csv_path,
             "continue_cols": continue_cols,
             "force_rerun": force_rerun,
-            "read_csv_kwargs": read_csv_kwargs or {},
+            "read_csv_kwargs": read_csv_kwargs,
             "uuid_column": uuid_column,
             "retval_column": retval_column,
+            "extra_kwargs": extra_kwargs,
         }
         return asyncio.run(self.arun(**kwargs))
 
@@ -68,15 +71,22 @@ def submit_from_csv(
         if not force_rerun:
             rows = False
             for col in self.continue_cols:
-                rows |= df[col].isnull()
-            added = int(rows.sum())
-            logger.info(f"Adding {added} tasks ({len(df) - added} skipped).")
+                if col in df.columns:
+                    rows |= df[col].isnull()
+            if isinstance(rows, bool):
+                rows = slice(None)
+                logger.info(f"Adding {len(df)} tasks.")
+            else:
+                added = int(rows.sum())
+                logger.info(
+                    f"Adding {added} tasks ({len(df) - added} skipped).")
         else:
             rows = slice(None)
             logger.info(f"Adding {len(df)} tasks.")
 
         tasks = [
-            self.create_task(uuid, **row) for uuid, row in df[rows].iterrows()
+            self.create_task(uuid, **row, **self.extra_kwargs)
+            for uuid, row in df[rows].iterrows()
         ]
 
         return tasks
@@ -112,13 +122,15 @@ async def arun(
         read_csv_kwargs: Optional[dict] = None,
         uuid_column: str = ":uuid:",
         retval_column: Optional[str] = ":retval:",
+        extra_kwargs: Optional[Dict[str, Any]] = None,
     ):
         """Async run experiments from a csv file"""
 
         self.csv_path = csv_path
         self.continue_cols = continue_cols
-        self.read_csv_kwargs = read_csv_kwargs
+        self.read_csv_kwargs = read_csv_kwargs or {}
         self.uuid_column = uuid_column
+        self.extra_kwargs = extra_kwargs or {}
 
         tasks = self.submit_from_csv(force_rerun)