Support for execution plans

amlb/__init__.py

@@ -2,22 +2,22 @@
 amlb entrypoint package.
 """
 
+from .benchmark import SetupMode
 from .logger import app_logger as log
 from .errors import AutoMLError
 from .resources import Resources
-from .benchmark import Benchmark, SetupMode
-from .runners import AWSBenchmark, DockerBenchmark, SingularityBenchmark
+from .runners import LocalBenchmark, AWSBenchmark, DockerBenchmark, SingularityBenchmark
 from .results import TaskResult
 from .__version__ import __version__
 
 __all__ = [
     "log",
     "AutoMLError",
     "Resources",
-    "Benchmark",
+    "LocalBenchmark",
+    "AWSBenchmark",
     "DockerBenchmark",
     "SingularityBenchmark",
-    "AWSBenchmark",
     "SetupMode",
     "TaskResult",
     "__version__"

amlb/plan.py

@@ -0,0 +1,31 @@
+from dataclasses import dataclass
+import logging
+from typing import List
+
+from .benchmark import Benchmark
+
+log = logging.getLogger(__name__)
+
+
+@dataclass
+class RunTask:
+    framework: str
+    benchmark: str
+    constraint: str
+    tasks: [str]
+    folds: [int]
+
+
+class ExecutionPlan:
+
+    def __init__(self, benchmark_cls, plan: List[RunTask]):
+        self.benchmark_cls = benchmark_cls
+        self.plan = plan
+
+    def run(self):
+        main_instance = Benchmark(None, None, None)
+        bench_instances = [self.benchmark_cls(rt.framework, rt.benchmark, rt.constraint) for rt in self.plan]
+        jobs = []
+        for instance, rt in zip(bench_instances, self.plan):
+            jobs.extend(instance.create_jobs(tasks=rt.tasks, folds=rt.folds))
+        return main_instance.run_jobs(jobs)

amlb/results.py

@@ -50,7 +50,6 @@ class ResultError(Exception):
     pass
 
 
-
 class Scoreboard:
 
     results_file = 'results.csv'
@@ -404,7 +403,7 @@ def compute_score(self, result=None, meta_result=None):
             version=metadata.framework_version,
             params=repr(metadata.framework_params) if metadata.framework_params else '',
             fold=self.fold,
-            mode=rconfig().run_mode,
+            mode=metadata.run_mode,
             seed=metadata.seed,
             app_version=rget().app_version,
             utc=datetime_iso(),

amlb/runners/__init__.py

@@ -2,12 +2,14 @@
 benchmark runners
 """
 
+from .local import LocalBenchmark
 from .aws import AWSBenchmark
 from .docker import DockerBenchmark
 from .singularity import SingularityBenchmark
 
 __all__ = [
+    "LocalBenchmark",
+    "AWSBenchmark",
     "DockerBenchmark",
     "SingularityBenchmark",
-    "AWSBenchmark",
 ]

amlb/runners/aws.py

@@ -62,6 +62,7 @@ class AWSBenchmark(Benchmark):
     """AWSBenchmark
     an extension of Benchmark class, to run benchmarks on AWS
     """
+    run_mode = 'aws'
 
     @classmethod
     def fetch_results(cls, instances_file, instance_selector=None):
@@ -112,7 +113,7 @@ def _on_done(job_self):
         finally:
             bench.cleanup()
 
-    def __init__(self, framework_name, benchmark_name, constraint_name, region=None):
+    def __init__(self, framework_name: str, benchmark_name: str, constraint_name: str, region: str = None):
         """
 
         :param framework_name:
@@ -183,41 +184,40 @@ def cleanup(self):
         if rconfig().aws.s3.temporary is True:
             self._delete_s3_bucket()
 
-    def run(self, task_name=None, fold=None):
-        task_defs = self._get_task_defs(task_name)  # validates tasks
-        self._exec_start()
-        self._monitoring_start()
+    def run(self, tasks=None, folds=None):
+        task_defs = self._get_task_defs(tasks)  # validates tasks
+        jobs = None
         if self.parallel_jobs > 1:
             if rconfig().aws.minimize_instances:
                 # use one instance per task: all folds executed on same instance
-                try:
-                    jobs = flatten([self._make_aws_job([task_def.name], fold) for task_def in task_defs])
-                    results = self._run_jobs(jobs)
-                    return self._process_results(results, task_name=task_name)
-                finally:
-                    self.cleanup()
-            else:
-                # use one instance per fold per task
-                return super().run(task_name, fold)
+                jobs = flatten([self._make_aws_job([task_def.name], folds) for task_def in task_defs])
         else:
             # use one instance for all
-            try:
-                task_names = None if task_name is None else [task_def.name for task_def in task_defs]
-                job = self._make_aws_job(task_names, fold)
-                results = self._run_jobs([job])
-                return self._process_results(results, task_name=task_name)
-            finally:
-                self.cleanup()
+            task_names = None if tasks is None else [task_def.name for task_def in task_defs]
+            jobs = [self._make_aws_job(task_names, folds)]
 
-    def _create_job_runner(self, jobs):
-        if self.parallel_jobs == 1:
-            return SimpleJobRunner(jobs)
+        if jobs is not None:
+            self.run_jobs(jobs)
         else:
-            queueing_strategy = MultiThreadingJobRunner.QueueingStrategy.enforce_job_priority
-            return MultiThreadingJobRunner(jobs, self.parallel_jobs,
+            # use one instance per fold per task
+            return super().run(tasks, folds)
+
+    def _create_job_runner(self, jobs):
+        jr = (SimpleJobRunner(jobs) if self.parallel_jobs == 1
+              else MultiThreadingJobRunner(jobs, self.parallel_jobs,
                                            delay_secs=rconfig().job_scheduler.delay_between_jobs,
                                            done_async=True,
-                                           queueing_strategy=queueing_strategy)
+                                           queueing_strategy=MultiThreadingJobRunner.QueueingStrategy.enforce_job_priority))
+
+        super_on_state = jr._on_state
+
+        def _on_state(_self, state):
+            if state is JobState.starting:
+                self._exec_start()
+                self._monitoring_start()
+            super_on_state(state)
+
+        jr._on_state = _on_state.__get__(jr)
 
     def _make_job(self, task_def, fold=int):
         return self._make_aws_job([task_def.name], [fold])
@@ -274,8 +274,8 @@ def _reset_retry(self):
         for j in self.jobs:
             j.ext.retry = None
 
-    def _make_aws_job(self, task_names=None, folds=None):
-        task_names = [] if task_names is None else task_names
+    def _make_aws_job(self, tasks=None, folds=None):
+        task_names = [] if tasks is None else tasks
         folds = [] if folds is None else [str(f) for f in folds]
         task_def = (self._get_task_def(task_names[0]) if len(task_names) >= 1
                     else self._get_task_def('__defaults__', include_disabled=True, fail_on_missing=False) or ns(name='all'))
@@ -888,8 +888,7 @@ def download_file(obj, dest, dest_display_path=None):
                     download_file(obj, dest_path)
                     # if obj.key == result_key:
                     if is_result and not success:
-                        if rconfig().results.save:
-                            self._exec_send(lambda path: self._append(Scoreboard.load_df(path)), dest_path)
+                        self._exec_send(lambda path: self._append(Scoreboard.load_df(path)), dest_path)
                         success = True
                 except Exception as e:
                     if is_result:

amlb/runners/container.py

@@ -6,14 +6,12 @@
 """
 from abc import abstractmethod
 import logging
-import os
 import re
 
 from ..benchmark import Benchmark, SetupMode
 from ..errors import InvalidStateError
 from ..job import Job
 from ..resources import config as rconfig, get as rget
-from ..utils import dir_of, run_cmd
 from ..__version__ import __version__, _dev_version as dev
 
 
@@ -40,7 +38,7 @@ def image_name(cls, framework_def, label=None, **kwargs):
         return f"{author}/{image}:{tag}"
 
     @abstractmethod
-    def __init__(self, framework_name, benchmark_name, constraint_name):
+    def __init__(self, framework_name: str, benchmark_name: str, constraint_name: str):
         """
 
         :param framework_name:
@@ -82,17 +80,13 @@ def cleanup(self):
         # TODO: remove generated script? anything else?
         pass
 
-    def run(self, task_name=None, fold=None):
-        self._get_task_defs(task_name)  # validates tasks
+    def run(self, tasks=None, folds=None):
+        self._get_task_defs(tasks)  # validates tasks
         if self.parallel_jobs > 1 or not self.minimize_instances:
-            return super().run(task_name, fold)
+            return super().run(tasks, folds)
         else:
-            job = self._make_container_job(task_name, fold)
-            try:
-                results = self._run_jobs([job])
-                return self._process_results(results, task_name=task_name)
-            finally:
-                self.cleanup()
+            job = self._make_container_job(tasks, folds)
+            return self.run_jobs([job])
 
     def _make_job(self, task_def, fold=int):
         return self._make_container_job([task_def.name], [fold])

amlb/runners/docker.py

@@ -21,7 +21,9 @@ class DockerBenchmark(ContainerBenchmark):
     an extension of ContainerBenchmark to run benchmarks inside docker.
     """
 
-    def __init__(self, framework_name, benchmark_name, constraint_name):
+    run_mode = 'docker'
+
+    def __init__(self, framework_name: str, benchmark_name: str, constraint_name: str):
         """
 
         :param framework_name: