From cae8e47da5b8f3bc99a023ce95cc6905c1068f4e Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Fri, 12 Jun 2026 19:59:59 +0200 Subject: [PATCH 01/13] Introducing yaml spec for rose api --- pyproject.toml | 2 + rose/__init__.py | 4 + rose/spec/__init__.py | 57 ++++ rose/spec/adapters.py | 100 +++++++ rose/spec/builder.py | 107 ++++++++ rose/spec/schema.py | 155 +++++++++++ tests/integration/spec/__init__.py | 0 tests/integration/spec/helpers.py | 17 ++ tests/integration/spec/test_yaml_workflow.py | 73 +++++ tests/unit/spec/__init__.py | 0 tests/unit/spec/test_adapters.py | 96 +++++++ tests/unit/spec/test_schema.py | 267 +++++++++++++++++++ 12 files changed, 878 insertions(+) create mode 100644 rose/spec/__init__.py create mode 100644 rose/spec/adapters.py create mode 100644 rose/spec/builder.py create mode 100644 rose/spec/schema.py create mode 100644 tests/integration/spec/__init__.py create mode 100644 tests/integration/spec/helpers.py create mode 100644 tests/integration/spec/test_yaml_workflow.py create mode 100644 tests/unit/spec/__init__.py create mode 100644 tests/unit/spec/test_adapters.py create mode 100644 tests/unit/spec/test_schema.py diff --git a/pyproject.toml b/pyproject.toml index e8fcca05..f57c4226 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,6 +36,8 @@ mlflow = ["mlflow>=2.0"] clearml = ["clearml>=1.14"] tracking = ["mlflow>=2.0", "clearml>=1.14"] +spec = ["pyyaml>=6.0"] + # All test deps dev = [ "pytest", diff --git a/rose/__init__.py b/rose/__init__.py index 0038267b..db05a14b 100644 --- a/rose/__init__.py +++ b/rose/__init__.py @@ -2,6 +2,7 @@ from rose.learner import IterationState, Learner, LearnerConfig, TaskConfig from rose.metrics import * # noqa: F403 from rose.rl import reinforcement_learner +from rose.spec import WorkflowSpec, load_spec from rose.tracking import PipelineManifest, TrackerBase from rose.uq import uq_active_learner, uq_learner, uq_scorer @@ -21,4 +22,7 @@ # Tracking "TrackerBase", "PipelineManifest", + # YAML spec layer + "load_spec", + "WorkflowSpec", ] diff --git a/rose/spec/__init__.py b/rose/spec/__init__.py new file mode 100644 index 00000000..6852be15 --- /dev/null +++ b/rose/spec/__init__.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Any, Callable, Coroutine + +from .builder import LearnerBuilder +from .schema import SpecConfig + +__all__ = ["load_spec", "WorkflowSpec"] + + +def load_spec(path: str | Path) -> "WorkflowSpec": + """Load and validate a YAML workflow spec. Raises ValueError on schema errors.""" + return WorkflowSpec(SpecConfig.from_yaml(path)) + + +class WorkflowSpec: + """Validated workflow spec that produces a coroutine compatible with service_utils.run().""" + + def __init__(self, config: SpecConfig) -> None: + self.config = config + + @property + def workflow(self) -> Callable[..., Coroutine[Any, Any, None]]: + cfg = self.config + + async def _workflow(bridge_url: str, edge_name: str) -> None: + import rhapsody + from radical.asyncflow import WorkflowEngine + + engine = await rhapsody.get_backend( + "edge", bridge_url=bridge_url, edge_name=edge_name + ) + asyncflow = await WorkflowEngine.create(engine) + + builder = LearnerBuilder(cfg, asyncflow) + learner = builder.build() + + start_kwargs: dict[str, Any] = {"max_iter": cfg.learner.max_iter} + if cfg.learner.type == "parallel_active_learner": + lc = builder.build_learner_configs() + if lc is not None: + start_kwargs["parallel_learners"] = len(lc) + start_kwargs["learner_configs"] = lc + else: + start_kwargs["parallel_learners"] = cfg.learner.parallel_learners + + try: + async for state in learner.start(**start_kwargs): + print( + f"[iter {state.iteration}] metric={state.metric_value}", + flush=True, + ) + finally: + await asyncflow.shutdown() + + return _workflow diff --git a/rose/spec/adapters.py b/rose/spec/adapters.py new file mode 100644 index 00000000..dee7e9c5 --- /dev/null +++ b/rose/spec/adapters.py @@ -0,0 +1,100 @@ +from __future__ import annotations + +from typing import Callable + +from .schema import RemoteConfig, TaskDef + + +class TaskAdapterFactory: + @staticmethod + def make_closure(task_def: TaskDef, remote: RemoteConfig) -> Callable: + if task_def.type == "shell": + return _make_shell_closure(task_def.command) + return _make_python_closure(task_def.function, remote.pythonpath) + + @staticmethod + def as_executable(task_def: TaskDef) -> bool: + return task_def.type == "shell" + + @staticmethod + def make_dispatch_closure( + slot_name: str, + task_defs: list[TaskDef], + remote: RemoteConfig, + ) -> Callable: + """Dispatch closure for parallel learners: routes per-candidate based on learner_id kwarg.""" + if task_defs[0].type == "shell": + return _make_shell_dispatch( + {i: td.command for i, td in enumerate(task_defs)}, slot_name + ) + return _make_python_dispatch( + {i: td.function for i, td in enumerate(task_defs)}, + list(remote.pythonpath), + slot_name, + ) + + +def _make_shell_closure(command: str) -> Callable: + _cmd = command + + async def _task(*args, **kwargs) -> str: + return _cmd + + _task.__name__ = "shell_task" + return _task + + +def _make_python_closure(spec: str, remote_paths: list[str]) -> Callable: + _spec = spec + _paths = list(remote_paths) + + async def _task(*args, **kwargs): + import importlib as _il + import inspect as _ins + import sys as _sys + + for p in _paths: + if p not in _sys.path: + _sys.path.insert(0, p) + mod_path, fn_name = _spec.rsplit(":", 1) + fn = getattr(_il.import_module(mod_path), fn_name) + result = fn(*args, **kwargs) + return (await result) if _ins.iscoroutine(result) else result + + _task.__name__ = spec.split(":")[-1] + return _task + + +def _make_shell_dispatch(cmds: dict[int, str], slot_name: str) -> Callable: + _cmds = dict(cmds) + + async def _dispatch(*args, **kwargs) -> str: + return _cmds[kwargs.get("learner_id", 0)] + + _dispatch.__name__ = f"{slot_name}_dispatch" + return _dispatch + + +def _make_python_dispatch( + specs: dict[int, str], remote_paths: list[str], slot_name: str +) -> Callable: + _specs = dict(specs) + _paths = list(remote_paths) + + async def _dispatch(*args, **kwargs): + import importlib as _il + import inspect as _ins + import sys as _sys + + lid = kwargs.pop("learner_id", 0) # strip internal routing key before calling user fn + spec = _specs[lid] + for p in _paths: + if p not in _sys.path: + _sys.path.insert(0, p) + mod_path, fn_name = spec.rsplit(":", 1) + fn = getattr(_il.import_module(mod_path), fn_name) + result = fn(*args, **kwargs) + return (await result) if _ins.iscoroutine(result) else result + + _dispatch.__name__ = f"{slot_name}_dispatch" + return _dispatch diff --git a/rose/spec/builder.py b/rose/spec/builder.py new file mode 100644 index 00000000..316569fd --- /dev/null +++ b/rose/spec/builder.py @@ -0,0 +1,107 @@ +from __future__ import annotations + +from rose.learner import Learner + +from .adapters import TaskAdapterFactory +from .schema import SpecConfig, TrackingConfig, _REQUIRED_SLOTS + + +def _get_learner_class(learner_type: str): + if learner_type == "sequential_active_learner": + from rose.al.active_learner import SequentialActiveLearner + + return SequentialActiveLearner + if learner_type == "parallel_active_learner": + from rose.al.active_learner import ParallelActiveLearner + + return ParallelActiveLearner + if learner_type == "sequential_reinforcement_learner": + from rose.rl.reinforcement_learner import SequentialReinforcementLearner + + return SequentialReinforcementLearner + if learner_type == "uq_active_learner": + from rose.uq.uq_active_learner import SeqUQLearner + + return SeqUQLearner + raise ValueError(f"No learner class registered for type '{learner_type}'") + + +class LearnerBuilder: + def __init__(self, config: SpecConfig, asyncflow) -> None: + self.config = config + self.asyncflow = asyncflow + + def build(self) -> Learner: + cfg = self.config + learner = _get_learner_class(cfg.learner.type)(self.asyncflow) + + if cfg.candidates is not None: + self._register_dispatched_tasks(learner, cfg) + else: + self._register_flat_tasks(learner, cfg) + + crit = cfg.stop_criterion + c_closure = TaskAdapterFactory.make_closure(crit.evaluator, cfg.remote) + c_as_exec = TaskAdapterFactory.as_executable(crit.evaluator) + learner.as_stop_criterion( + metric_name=crit.metric, + threshold=crit.threshold, + operator=crit.operator, + as_executable=c_as_exec, + )(c_closure) + + _attach_tracker(learner, cfg.tracking) + return learner + + def _register_flat_tasks(self, learner: Learner, cfg: SpecConfig) -> None: + for slot_name, task_def in cfg.tasks.items(): + closure = TaskAdapterFactory.make_closure(task_def, cfg.remote) + as_exec = TaskAdapterFactory.as_executable(task_def) + getattr(learner, f"{slot_name}_task")(as_executable=as_exec)(closure) + + def _register_dispatched_tasks(self, learner: Learner, cfg: SpecConfig) -> None: + required = _REQUIRED_SLOTS[cfg.learner.type] + for slot_name in required: + task_defs = [c.tasks[slot_name] for c in cfg.candidates] + as_exec = task_defs[0].type == "shell" + closure = TaskAdapterFactory.make_dispatch_closure(slot_name, task_defs, cfg.remote) + getattr(learner, f"{slot_name}_task")(as_executable=as_exec)(closure) + + def build_learner_configs(self): + """Return auto-generated LearnerConfig list for parallel candidates, or None.""" + cfg = self.config + if cfg.candidates is None: + return None + from rose.learner import LearnerConfig, TaskConfig + + required = _REQUIRED_SLOTS[cfg.learner.type] + configs = [] + max_iter = cfg.learner.max_iter + for i, _candidate in enumerate(cfg.candidates): + schedule = {n: TaskConfig(kwargs={"learner_id": i}) for n in range(max_iter + 1)} + schedule[-1] = TaskConfig(kwargs={"learner_id": i}) + configs.append( + LearnerConfig(**{slot: schedule for slot in required}, criterion=schedule) + ) + return configs + + +def _attach_tracker(learner: Learner, tracking: TrackingConfig) -> None: + if tracking.backend == "mlflow": + from rose.integrations.mlflow_tracker import MLflowTracker + + learner.add_tracker( + MLflowTracker( + experiment_name=tracking.experiment, + run_name=tracking.run_name, + ) + ) + elif tracking.backend == "clearml": + from rose.integrations.clearml_tracker import ClearMLTracker + + learner.add_tracker( + ClearMLTracker( + project_name=tracking.experiment, + task_name=tracking.run_name or "rose-spec-run", + ) + ) diff --git a/rose/spec/schema.py b/rose/spec/schema.py new file mode 100644 index 00000000..3f21fb90 --- /dev/null +++ b/rose/spec/schema.py @@ -0,0 +1,155 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Literal + +from pydantic import BaseModel, Field, model_validator + + +# ── Task definition ─────────────────────────────────────────────────────────── +class TaskDef(BaseModel): + type: Literal["shell", "python"] + command: str | None = None # required when type=="shell" + function: str | None = None # required when type=="python"; "module:callable" + + model_config = {"extra": "forbid"} + + @model_validator(mode="after") + def _check_fields(self) -> "TaskDef": + if self.type == "shell" and not self.command: + raise ValueError("shell task requires 'command'") + if self.type == "python": + if not self.function: + raise ValueError("python task requires 'function'") + if ":" not in self.function: + raise ValueError("'function' must use 'module:callable' syntax") + return self + + +# ── Stop criterion ──────────────────────────────────────────────────────────── +class StopCriterionDef(BaseModel): + metric: str + threshold: float + operator: Literal["<", ">", "==", "<=", ">="] = "<" + evaluator: TaskDef + + model_config = {"extra": "forbid"} + + +# ── Learner spec (YAML key: "learner") ─────────────────────────────────────── +_REQUIRED_SLOTS: dict[str, frozenset[str]] = { + "sequential_active_learner": frozenset({"simulation", "training", "active_learn"}), + "parallel_active_learner": frozenset({"simulation", "training", "active_learn"}), + "sequential_reinforcement_learner": frozenset({"environment", "update"}), + "uq_active_learner": frozenset({"simulation", "training", "prediction", + "active_learn", "uncertainty"}), +} +_ALL_SLOTS: frozenset[str] = frozenset().union(*_REQUIRED_SLOTS.values()) + + +class LearnerSpec(BaseModel): + type: str + max_iter: int = 0 + parallel_learners: int = 2 # used only when candidates is absent for parallel types + + model_config = {"extra": "forbid"} + + +# ── Per-candidate definition (parallel learners only) ───────────────────────── +class CandidateDef(BaseModel): + label: str = "" + simulation: TaskDef | None = None + training: TaskDef | None = None + active_learn: TaskDef | None = None + environment: TaskDef | None = None + update: TaskDef | None = None + prediction: TaskDef | None = None + uncertainty: TaskDef | None = None + + model_config = {"extra": "forbid"} + + @property + def tasks(self) -> dict[str, TaskDef]: + return {s: getattr(self, s) for s in _ALL_SLOTS if getattr(self, s) is not None} + + +# ── Remote / tracking ──────────────────────────────────────────────────────── +class RemoteConfig(BaseModel): + pythonpath: list[str] = [] + + model_config = {"extra": "forbid"} + + +class TrackingConfig(BaseModel): + backend: Literal["mlflow", "clearml", "none"] = "none" + experiment: str = "ROSE-Spec" + run_name: str | None = None + + model_config = {"extra": "forbid"} + + +# ── Top-level spec ──────────────────────────────────────────────────────────── +class SpecConfig(BaseModel): + learner: LearnerSpec + # Task slots — explicit fields keep extra="forbid" and enable IDE autocomplete. + # Access non-None slots as a dict via the .tasks property. + simulation: TaskDef | None = None + training: TaskDef | None = None + active_learn: TaskDef | None = None + environment: TaskDef | None = None + update: TaskDef | None = None + prediction: TaskDef | None = None + uncertainty: TaskDef | None = None + candidates: list[CandidateDef] | None = None + stop_criterion: StopCriterionDef + remote: RemoteConfig = Field(default_factory=RemoteConfig) + tracking: TrackingConfig = Field(default_factory=TrackingConfig) + + model_config = {"extra": "forbid"} + + @property + def tasks(self) -> dict[str, TaskDef]: + return {s: getattr(self, s) for s in _ALL_SLOTS if getattr(self, s) is not None} + + @model_validator(mode="after") + def _validate_task_slots(self) -> "SpecConfig": + ltype = self.learner.type + required = _REQUIRED_SLOTS.get(ltype) + if required is None: + raise ValueError( + f"Unknown learner type '{ltype}'. " + f"Supported: {sorted(_REQUIRED_SLOTS.keys())}" + ) + is_parallel = ltype == "parallel_active_learner" + + if is_parallel and self.candidates is not None: + for c in self.candidates: + missing = required - set(c.tasks.keys()) + if missing: + raise ValueError(f"Candidate '{c.label}' missing: {sorted(missing)}") + extra = set(c.tasks.keys()) - required + if extra: + raise ValueError(f"Candidate '{c.label}' unexpected fields: {sorted(extra)}") + for slot in required: + types = {c.tasks[slot].type for c in self.candidates} + if len(types) > 1: + raise ValueError( + f"Slot '{slot}' has mixed types across candidates: {types}. " + "All candidates must use the same type for a given slot." + ) + else: + present = set(self.tasks.keys()) + missing = required - present + if missing: + raise ValueError(f"learner type '{ltype}' requires: {sorted(missing)}") + extra = present - required + if extra: + raise ValueError(f"Unexpected task fields for '{ltype}': {sorted(extra)}") + return self + + @classmethod + def from_yaml(cls, path: str | Path) -> "SpecConfig": + import yaml + + raw = yaml.safe_load(Path(path).read_text()) + return cls.model_validate(raw) diff --git a/tests/integration/spec/__init__.py b/tests/integration/spec/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/integration/spec/helpers.py b/tests/integration/spec/helpers.py new file mode 100644 index 00000000..d8ef3090 --- /dev/null +++ b/tests/integration/spec/helpers.py @@ -0,0 +1,17 @@ +"""Lightweight task helpers used by the YAML spec integration tests.""" + + +async def sim(*args, **kwargs): + return [1.0, 2.0, 3.0] + + +async def train(data, **kwargs): + return {"mean": sum(data) / len(data)} + + +async def active_learn(sim_result, model, **kwargs): + return abs(model["mean"] - 2.0) + + +async def criterion(*args, **kwargs): + return 0.05 diff --git a/tests/integration/spec/test_yaml_workflow.py b/tests/integration/spec/test_yaml_workflow.py new file mode 100644 index 00000000..084de7c2 --- /dev/null +++ b/tests/integration/spec/test_yaml_workflow.py @@ -0,0 +1,73 @@ +"""Integration tests for the YAML spec layer — full AL loop, no HPC required.""" +import textwrap +from concurrent.futures import ThreadPoolExecutor + +import pytest +from radical.asyncflow import WorkflowEngine +from rhapsody.backends import ConcurrentExecutionBackend + +from rose.spec import load_spec +from rose.spec.builder import LearnerBuilder +from rose.spec.schema import SpecConfig + + +# ── Sequential AL via YAML ──────────────────────────────────────────────────── + +SEQ_YAML = textwrap.dedent("""\ + learner: + type: sequential_active_learner + max_iter: 3 + + simulation: + type: python + function: tests.integration.spec.helpers:sim + + training: + type: python + function: tests.integration.spec.helpers:train + + active_learn: + type: python + function: tests.integration.spec.helpers:active_learn + + stop_criterion: + metric: mse + threshold: 0.01 + operator: "<" + evaluator: + type: python + function: tests.integration.spec.helpers:criterion +""") + + +@pytest.mark.asyncio +@pytest.mark.integration +async def test_yaml_sequential_workflow(tmp_path): + p = tmp_path / "spec.yaml" + p.write_text(SEQ_YAML) + + engine = await ConcurrentExecutionBackend(ThreadPoolExecutor()) + asyncflow = await WorkflowEngine.create(engine) + + cfg = SpecConfig.from_yaml(p) + builder = LearnerBuilder(cfg, asyncflow) + learner = builder.build() + + states = [] + async for state in learner.start(max_iter=3): + states.append(state) + + assert len(states) == 3 + await asyncflow.shutdown() + + +# ── load_spec convenience wrapper ───────────────────────────────────────────── + +@pytest.mark.asyncio +@pytest.mark.integration +async def test_load_spec_returns_workflow_spec(tmp_path): + p = tmp_path / "spec.yaml" + p.write_text(SEQ_YAML) + spec = load_spec(p) + assert hasattr(spec, "workflow") + assert callable(spec.workflow) diff --git a/tests/unit/spec/__init__.py b/tests/unit/spec/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/spec/test_adapters.py b/tests/unit/spec/test_adapters.py new file mode 100644 index 00000000..3c05dee0 --- /dev/null +++ b/tests/unit/spec/test_adapters.py @@ -0,0 +1,96 @@ +"""Unit tests for rose.spec.adapters — closure factories, no HPC needed.""" +import asyncio + +import pytest + +from rose.spec.adapters import TaskAdapterFactory +from rose.spec.schema import RemoteConfig, TaskDef + + +def _remote() -> RemoteConfig: + return RemoteConfig() + + +# ── Shell closure ───────────────────────────────────────────────────────────── + +def test_shell_closure_returns_command(): + td = TaskDef(type="shell", command="echo hello") + fn = TaskAdapterFactory.make_closure(td, _remote()) + result = asyncio.run(fn()) + assert result == "echo hello" + + + +def test_shell_as_executable_true(): + td = TaskDef(type="shell", command="x") + assert TaskAdapterFactory.as_executable(td) is True + + +# ── Python closure ──────────────────────────────────────────────────────────── + +def test_python_closure_calls_sync_function(): + td = TaskDef(type="python", function="os.path:join") + fn = TaskAdapterFactory.make_closure(td, _remote()) + result = asyncio.run(fn("a", "b")) + assert result == "a/b" + + +def test_python_closure_calls_async_function(): + async def _async_fn(x): + return x * 2 + + import types + mod = types.ModuleType("_test_async_mod") + mod.double = _async_fn + import sys + sys.modules["_test_async_mod"] = mod + + td = TaskDef(type="python", function="_test_async_mod:double") + fn = TaskAdapterFactory.make_closure(td, _remote()) + result = asyncio.run(fn(21)) + assert result == 42 + + del sys.modules["_test_async_mod"] + + +def test_python_as_executable_false(): + td = TaskDef(type="python", function="os:getcwd") + assert TaskAdapterFactory.as_executable(td) is False + + +def test_python_closure_injects_remote_path(tmp_path): + (tmp_path / "myutil.py").write_text("def add(a, b): return a + b\n") + td = TaskDef(type="python", function="myutil:add") + fn = TaskAdapterFactory.make_closure(td, RemoteConfig(pythonpath=[str(tmp_path)])) + result = asyncio.run(fn(3, 4)) + assert result == 7 + + +# ── Shell dispatch ──────────────────────────────────────────────────────────── + +def test_shell_dispatch_routes_by_learner_id(): + tds = [ + TaskDef(type="shell", command="cmd_0"), + TaskDef(type="shell", command="cmd_1"), + ] + fn = TaskAdapterFactory.make_dispatch_closure("simulation", tds, _remote()) + assert asyncio.run(fn(learner_id=0)) == "cmd_0" + assert asyncio.run(fn(learner_id=1)) == "cmd_1" + + +def test_shell_dispatch_defaults_to_zero(): + tds = [TaskDef(type="shell", command="cmd_default")] + fn = TaskAdapterFactory.make_dispatch_closure("simulation", tds, _remote()) + assert asyncio.run(fn()) == "cmd_default" + + +# ── Python dispatch ─────────────────────────────────────────────────────────── + +def test_python_dispatch_routes_by_learner_id(): + tds = [ + TaskDef(type="python", function="os.path:basename"), + TaskDef(type="python", function="os.path:dirname"), + ] + fn = TaskAdapterFactory.make_dispatch_closure("training", tds, _remote()) + assert asyncio.run(fn("/a/b/c", learner_id=0)) == "c" + assert asyncio.run(fn("/a/b/c", learner_id=1)) == "/a/b" diff --git a/tests/unit/spec/test_schema.py b/tests/unit/spec/test_schema.py new file mode 100644 index 00000000..9a4e3626 --- /dev/null +++ b/tests/unit/spec/test_schema.py @@ -0,0 +1,267 @@ +"""Unit tests for rose.spec.schema — YAML validation without any HPC machinery.""" +import textwrap + +import pytest + +from rose.spec.schema import SpecConfig, TaskDef + + +# ── TaskDef ─────────────────────────────────────────────────────────────────── + +def test_taskdef_shell_valid(): + t = TaskDef(type="shell", command="python sim.py") + assert t.type == "shell" + assert t.command == "python sim.py" + + +def test_taskdef_python_valid(): + t = TaskDef(type="python", function="mymod.sub:fn") + assert t.function == "mymod.sub:fn" + + +def test_taskdef_shell_missing_command(): + with pytest.raises(ValueError, match="shell task requires 'command'"): + TaskDef(type="shell") + + +def test_taskdef_python_missing_function(): + with pytest.raises(ValueError, match="python task requires 'function'"): + TaskDef(type="python") + + +def test_taskdef_python_bad_syntax(): + with pytest.raises(ValueError, match="module:callable"): + TaskDef(type="python", function="mymod.fn") + + +def test_taskdef_extra_field_rejected(): + with pytest.raises(Exception): + TaskDef(type="shell", command="x", unknown_field="y") + + +# ── SpecConfig — sequential learner ────────────────────────────────────────── + +SEQ_YAML = textwrap.dedent("""\ + learner: + type: sequential_active_learner + max_iter: 5 + + simulation: + type: shell + command: python sim.py + + training: + type: python + function: mymod:train + + active_learn: + type: python + function: mymod:select + + stop_criterion: + metric: mse + threshold: 0.1 + operator: "<" + evaluator: + type: python + function: mymod:eval_mse +""") + + +def test_sequential_spec_roundtrip(tmp_path): + p = tmp_path / "spec.yaml" + p.write_text(SEQ_YAML) + cfg = SpecConfig.from_yaml(p) + assert cfg.learner.type == "sequential_active_learner" + assert cfg.learner.max_iter == 5 + assert cfg.simulation.command == "python sim.py" + assert cfg.training.function == "mymod:train" + assert cfg.stop_criterion.metric == "mse" + assert cfg.tasks == { + "simulation": cfg.simulation, + "training": cfg.training, + "active_learn": cfg.active_learn, + } + + +def test_sequential_spec_missing_slot(tmp_path): + yaml = textwrap.dedent("""\ + learner: + type: sequential_active_learner + simulation: + type: shell + command: python sim.py + training: + type: python + function: mymod:train + stop_criterion: + metric: mse + threshold: 0.1 + evaluator: + type: python + function: mymod:eval + """) + p = tmp_path / "bad.yaml" + p.write_text(yaml) + with pytest.raises(ValueError, match="active_learn"): + SpecConfig.from_yaml(p) + + +def test_sequential_spec_extra_slot(tmp_path): + yaml = textwrap.dedent("""\ + learner: + type: sequential_active_learner + simulation: + type: shell + command: python sim.py + training: + type: python + function: mymod:train + active_learn: + type: python + function: mymod:select + environment: + type: shell + command: python env.py + stop_criterion: + metric: mse + threshold: 0.1 + evaluator: + type: python + function: mymod:eval + """) + p = tmp_path / "bad.yaml" + p.write_text(yaml) + with pytest.raises(ValueError, match="Unexpected"): + SpecConfig.from_yaml(p) + + +def test_unknown_learner_type(tmp_path): + yaml = textwrap.dedent("""\ + learner: + type: unknown_learner + simulation: + type: shell + command: x + stop_criterion: + metric: mse + threshold: 0.1 + evaluator: + type: shell + command: x + """) + p = tmp_path / "bad.yaml" + p.write_text(yaml) + with pytest.raises(ValueError, match="Unknown learner type"): + SpecConfig.from_yaml(p) + + +# ── SpecConfig — parallel learner with candidates ──────────────────────────── + +PAR_YAML = textwrap.dedent("""\ + learner: + type: parallel_active_learner + max_iter: 3 + + candidates: + - label: rf + simulation: + type: shell + command: python sim.py --model rf + training: + type: python + function: mymod:train_rf + active_learn: + type: python + function: mymod:select + - label: mlp + simulation: + type: shell + command: python sim.py --model mlp + training: + type: python + function: mymod:train_mlp + active_learn: + type: python + function: mymod:select + + stop_criterion: + metric: r2 + threshold: 0.9 + operator: ">" + evaluator: + type: python + function: mymod:eval_r2 +""") + + +def test_parallel_candidates_roundtrip(tmp_path): + p = tmp_path / "par.yaml" + p.write_text(PAR_YAML) + cfg = SpecConfig.from_yaml(p) + assert len(cfg.candidates) == 2 + assert cfg.candidates[0].label == "rf" + assert cfg.candidates[1].tasks["training"].function == "mymod:train_mlp" + + +def test_parallel_candidate_missing_slot(tmp_path): + yaml = textwrap.dedent("""\ + learner: + type: parallel_active_learner + candidates: + - label: rf + simulation: + type: shell + command: python sim.py + training: + type: python + function: mymod:train + stop_criterion: + metric: r2 + threshold: 0.9 + evaluator: + type: python + function: mymod:eval + """) + p = tmp_path / "bad.yaml" + p.write_text(yaml) + with pytest.raises(ValueError, match="active_learn"): + SpecConfig.from_yaml(p) + + +def test_parallel_candidates_mixed_types(tmp_path): + yaml = textwrap.dedent("""\ + learner: + type: parallel_active_learner + candidates: + - label: a + simulation: + type: shell + command: python sim.py + training: + type: python + function: mymod:train_a + active_learn: + type: python + function: mymod:select + - label: b + simulation: + type: python + function: mymod:sim_b + training: + type: python + function: mymod:train_b + active_learn: + type: python + function: mymod:select + stop_criterion: + metric: r2 + threshold: 0.9 + evaluator: + type: python + function: mymod:eval + """) + p = tmp_path / "bad.yaml" + p.write_text(yaml) + with pytest.raises(ValueError, match="mixed types"): + SpecConfig.from_yaml(p) From b7bd9ac2dfe3e960cf7b2662809b11e6d90fd757 Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Sat, 13 Jun 2026 00:41:13 +0200 Subject: [PATCH 02/13] support config learner/task --- rose/spec/adapters.py | 36 +++++++++++++++++----------- rose/spec/builder.py | 15 ++++++++---- rose/spec/schema.py | 3 ++- tests/integration/__init__.py | 0 tests/unit/spec/test_adapters.py | 35 ++++++++++++++++++++++++++++ tests/unit/spec/test_schema.py | 40 ++++++++++++++++++++++++++++++++ 6 files changed, 110 insertions(+), 19 deletions(-) create mode 100644 tests/integration/__init__.py diff --git a/rose/spec/adapters.py b/rose/spec/adapters.py index dee7e9c5..06fb6eaf 100644 --- a/rose/spec/adapters.py +++ b/rose/spec/adapters.py @@ -8,9 +8,10 @@ class TaskAdapterFactory: @staticmethod def make_closure(task_def: TaskDef, remote: RemoteConfig) -> Callable: + td = dict(task_def.task_description or {}) if task_def.type == "shell": - return _make_shell_closure(task_def.command) - return _make_python_closure(task_def.function, remote.pythonpath) + return _make_shell_closure(task_def.command, td) + return _make_python_closure(task_def.function, remote.pythonpath, td) @staticmethod def as_executable(task_def: TaskDef) -> bool: @@ -22,33 +23,38 @@ def make_dispatch_closure( task_defs: list[TaskDef], remote: RemoteConfig, ) -> Callable: - """Dispatch closure for parallel learners: routes per-candidate based on learner_id kwarg.""" + """Dispatch closure for parallel learners: routes per-candidate based on learner_id kwarg. + task_description uses the first candidate's value — asyncflow reads it once at registration.""" + td = dict(task_defs[0].task_description or {}) if task_defs[0].type == "shell": return _make_shell_dispatch( - {i: td.command for i, td in enumerate(task_defs)}, slot_name + {i: tdi.command for i, tdi in enumerate(task_defs)}, slot_name, td ) return _make_python_dispatch( - {i: td.function for i, td in enumerate(task_defs)}, + {i: tdi.function for i, tdi in enumerate(task_defs)}, list(remote.pythonpath), slot_name, + td, ) -def _make_shell_closure(command: str) -> Callable: +def _make_shell_closure(command: str, task_description: dict) -> Callable: _cmd = command + _td = task_description - async def _task(*args, **kwargs) -> str: + async def _task(*args, task_description=_td, **kwargs) -> str: return _cmd _task.__name__ = "shell_task" return _task -def _make_python_closure(spec: str, remote_paths: list[str]) -> Callable: - _spec = spec +def _make_python_closure(spec: str, remote_paths: list[str], task_description: dict) -> Callable: + _spec = spec _paths = list(remote_paths) + _td = task_description - async def _task(*args, **kwargs): + async def _task(*args, task_description=_td, **kwargs): import importlib as _il import inspect as _ins import sys as _sys @@ -65,10 +71,11 @@ async def _task(*args, **kwargs): return _task -def _make_shell_dispatch(cmds: dict[int, str], slot_name: str) -> Callable: +def _make_shell_dispatch(cmds: dict[int, str], slot_name: str, task_description: dict) -> Callable: _cmds = dict(cmds) + _td = task_description - async def _dispatch(*args, **kwargs) -> str: + async def _dispatch(*args, task_description=_td, **kwargs) -> str: return _cmds[kwargs.get("learner_id", 0)] _dispatch.__name__ = f"{slot_name}_dispatch" @@ -76,12 +83,13 @@ async def _dispatch(*args, **kwargs) -> str: def _make_python_dispatch( - specs: dict[int, str], remote_paths: list[str], slot_name: str + specs: dict[int, str], remote_paths: list[str], slot_name: str, task_description: dict ) -> Callable: _specs = dict(specs) _paths = list(remote_paths) + _td = task_description - async def _dispatch(*args, **kwargs): + async def _dispatch(*args, task_description=_td, **kwargs): import importlib as _il import inspect as _ins import sys as _sys diff --git a/rose/spec/builder.py b/rose/spec/builder.py index 316569fd..2f0375a1 100644 --- a/rose/spec/builder.py +++ b/rose/spec/builder.py @@ -75,11 +75,18 @@ def build_learner_configs(self): from rose.learner import LearnerConfig, TaskConfig required = _REQUIRED_SLOTS[cfg.learner.type] - configs = [] + configs = [] max_iter = cfg.learner.max_iter - for i, _candidate in enumerate(cfg.candidates): - schedule = {n: TaskConfig(kwargs={"learner_id": i}) for n in range(max_iter + 1)} - schedule[-1] = TaskConfig(kwargs={"learner_id": i}) + for i, candidate in enumerate(cfg.candidates): + # Build per-iteration TaskConfig kwargs, mirroring how ROSE expects LearnerConfig to work. + # learner_id → dispatch routing key (popped by dispatch closure, never reaches user fn) + # iteration → current iteration n, placed per-entry so get_task_config(slot, n) returns it + # learner_label → human-readable candidate name; only added when candidate has a label + base = {"learner_id": i} + if candidate.label: + base["learner_label"] = candidate.label + schedule = {n: TaskConfig(kwargs={**base, "iteration": n}) for n in range(max_iter + 1)} + schedule[-1] = TaskConfig(kwargs={**base, "iteration": max_iter}) configs.append( LearnerConfig(**{slot: schedule for slot in required}, criterion=schedule) ) diff --git a/rose/spec/schema.py b/rose/spec/schema.py index 3f21fb90..0e23105e 100644 --- a/rose/spec/schema.py +++ b/rose/spec/schema.py @@ -1,7 +1,7 @@ from __future__ import annotations from pathlib import Path -from typing import Literal +from typing import Any, Literal from pydantic import BaseModel, Field, model_validator @@ -11,6 +11,7 @@ class TaskDef(BaseModel): type: Literal["shell", "python"] command: str | None = None # required when type=="shell" function: str | None = None # required when type=="python"; "module:callable" + task_description: dict[str, Any] | None = None # resource hints forwarded to asyncflow backend model_config = {"extra": "forbid"} diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/spec/test_adapters.py b/tests/unit/spec/test_adapters.py index 3c05dee0..bd87fca8 100644 --- a/tests/unit/spec/test_adapters.py +++ b/tests/unit/spec/test_adapters.py @@ -1,5 +1,6 @@ """Unit tests for rose.spec.adapters — closure factories, no HPC needed.""" import asyncio +import inspect import pytest @@ -94,3 +95,37 @@ def test_python_dispatch_routes_by_learner_id(): fn = TaskAdapterFactory.make_dispatch_closure("training", tds, _remote()) assert asyncio.run(fn("/a/b/c", learner_id=0)) == "c" assert asyncio.run(fn("/a/b/c", learner_id=1)) == "/a/b" + + +# ── task_description default kwarg injection ────────────────────────────────── + +def test_shell_closure_task_description_in_signature(): + td = TaskDef(type="shell", command="echo hi", task_description={"cpu_count": 4}) + fn = TaskAdapterFactory.make_closure(td, _remote()) + sig = inspect.signature(fn) + assert "task_description" in sig.parameters + assert sig.parameters["task_description"].default == {"cpu_count": 4} + + +def test_shell_closure_no_task_description_defaults_empty(): + td = TaskDef(type="shell", command="echo hi") + fn = TaskAdapterFactory.make_closure(td, _remote()) + sig = inspect.signature(fn) + assert sig.parameters["task_description"].default == {} + + +def test_python_closure_task_description_in_signature(): + td = TaskDef(type="python", function="os:getcwd", task_description={"gpu_count": 2}) + fn = TaskAdapterFactory.make_closure(td, _remote()) + sig = inspect.signature(fn) + assert sig.parameters["task_description"].default == {"gpu_count": 2} + + +def test_dispatch_closure_task_description_uses_first_candidate(): + tds = [ + TaskDef(type="shell", command="cmd_0", task_description={"cpu_count": 8}), + TaskDef(type="shell", command="cmd_1", task_description={"cpu_count": 4}), + ] + fn = TaskAdapterFactory.make_dispatch_closure("simulation", tds, _remote()) + sig = inspect.signature(fn) + assert sig.parameters["task_description"].default == {"cpu_count": 8} diff --git a/tests/unit/spec/test_schema.py b/tests/unit/spec/test_schema.py index 9a4e3626..ca6d748f 100644 --- a/tests/unit/spec/test_schema.py +++ b/tests/unit/spec/test_schema.py @@ -39,6 +39,46 @@ def test_taskdef_extra_field_rejected(): TaskDef(type="shell", command="x", unknown_field="y") +def test_taskdef_task_description_roundtrip(): + t = TaskDef(type="shell", command="python sim.py", task_description={"cpu_count": 4, "gpu_count": 1}) + assert t.task_description == {"cpu_count": 4, "gpu_count": 1} + + +def test_taskdef_task_description_defaults_none(): + t = TaskDef(type="shell", command="x") + assert t.task_description is None + + +def test_sequential_spec_task_description_from_yaml(tmp_path): + yaml = textwrap.dedent("""\ + learner: + type: sequential_active_learner + max_iter: 1 + simulation: + type: shell + command: python sim.py + task_description: + cpu_count: 8 + training: + type: python + function: mymod:train + active_learn: + type: python + function: mymod:select + stop_criterion: + metric: mse + threshold: 0.1 + evaluator: + type: python + function: mymod:eval + """) + p = tmp_path / "td.yaml" + p.write_text(yaml) + cfg = SpecConfig.from_yaml(p) + assert cfg.simulation.task_description == {"cpu_count": 8} + assert cfg.training.task_description is None + + # ── SpecConfig — sequential learner ────────────────────────────────────────── SEQ_YAML = textwrap.dedent("""\ From ed0ad86035fe4d14c86544e347723966a7e00fbd Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Tue, 16 Jun 2026 22:51:50 +0200 Subject: [PATCH 03/13] docs, tests and schema expansion --- docs/user-guide/rose-aas.md | 82 +++ docs/user-guide/spec-api.md | 585 +++++++++++++++++++ mkdocs.yml | 4 + rose/spec/__init__.py | 104 +++- rose/spec/adapters.py | 9 +- rose/spec/builder.py | 56 +- rose/spec/schema.py | 57 +- tests/integration/spec/helpers.py | 22 + tests/integration/spec/test_yaml_workflow.py | 141 +++++ tests/unit/spec/test_adapters.py | 28 + tests/unit/spec/test_builder.py | 261 +++++++++ tests/unit/spec/test_schema.py | 306 +++++++++- tests/unit/spec/test_workflow_spec.py | 247 ++++++++ 13 files changed, 1853 insertions(+), 49 deletions(-) create mode 100644 docs/user-guide/rose-aas.md create mode 100644 docs/user-guide/spec-api.md create mode 100644 tests/unit/spec/test_builder.py create mode 100644 tests/unit/spec/test_workflow_spec.py diff --git a/docs/user-guide/rose-aas.md b/docs/user-guide/rose-aas.md new file mode 100644 index 00000000..009c1c55 --- /dev/null +++ b/docs/user-guide/rose-aas.md @@ -0,0 +1,82 @@ +# ROSE as a Service + +This page describes ROSE's service model: how the active-learning/RL loop you define stays under your control while the actual simulation and training work executes on HPC, and what that does and does not require from you today. + +--- + +## Core idea: BYOF — Bring Your Own Workflow + +ROSE does not own the science. It owns the **loop**: submit a task, wait for it, decide what runs next, check a stop criterion, repeat — until `max_iter` or the criterion is met. + +You bring: + +- A `simulate` / `train` / `active_learn` (or `environment` / `update`, or `prediction` / `uncertainty`) implementation, as a plain Python function or a shell command. +- The environment that implementation needs (packages, data, scripts). +- The decision of where it should run — a laptop, a login node, or a leadership-class HPC allocation. + +ROSE brings the orchestration: dependency-correct task submission, iteration bookkeeping, stop-criterion evaluation, checkpoint-safe state, and tracking integration. The [YAML Spec API](spec-api.md) is the declarative form of this contract — `function: tasks:simulate` is a pointer to *your* code, not a hook into ROSE's. + +This is the same boundary that makes ROSE usable across domains without ROSE having to understand any of them: it sees `*args, **kwargs` in, a return value out, nothing about what happened in between. + +--- + +## Two places work can happen + +A ROSE learner doesn't run anything itself — it submits tasks through whichever `asyncflow` execution backend you hand it. Two backends matter for the "as a Service" framing: + +| Backend | Where the orchestration loop runs | Where tasks execute | Documented at | +|---|---|---|---| +| `RadicalExecutionBackend` (RHAPSODY) | Wherever your script runs — it submits a pilot job and blocks inside it | Inside the pilot job it just submitted | [Target Resources](target-resources.md) | +| Edge backend (RHAPSODY, `bridge_url` + `edge_name`) | Wherever your script runs — does **not** need to be the HPC machine | Inside a separate, already-running edge agent — possibly on a different machine entirely | this page | + +The first model is "submit a job, then run my loop inside it." The second is "run my loop here; dispatch tasks to a job running somewhere else." That second model is what makes ROSE service-like: the loop's control plane (your `learner.start()` call, deciding when to stop, talking to MLflow/ClearML) is decoupled from the compute plane (the HPC allocation actually executing `simulate`/`train`). + +--- + +## How ROSE operates within a job + +The edge backend connects two things over a bridge: + +1. **An edge agent**, running inside an HPC job allocation. The job itself is requested through whatever your site normally uses to get an allocation — it doesn't have to be requested by ROSE. Once the job starts, the edge agent comes up inside it and stays alive for the allocation's lifetime, ready to accept task descriptions. +2. **Your orchestration process**, running anywhere — your laptop, a long-lived service host, a CI runner. It builds the learner from your spec (`LearnerBuilder(cfg, asyncflow).build()`), starts the loop (`learner.start(...)`), and for every `simulation`/`training`/`active_learn`/... task it submits, the asyncflow engine forwards that task description over the bridge to the edge, waits for the result, and resumes the loop. + +The practical effect: you submit a ROSE workflow to HPC without an interactive session on the cluster, and without your laptop needing to stay connected to the scheduler — only to the bridge. The job allocation is what's expensive and scheduler-queued; your control process is cheap and can be restarted independently of it. Stop-criterion checks, `set_next_config()` decisions, and tracking calls all happen on your side of the bridge, on every iteration, with no per-iteration job resubmission. + +This is additive to the model in [Target Resources](target-resources.md), not a replacement — both go through the same `WorkflowEngine`/`LearnerBuilder` plumbing. Which backend you choose only changes *where* the edge lives; the spec, the learner, and the loop semantics are identical either way. + +--- + +## What ROSE automates for you today + +- **The loop itself** — iterate, await, check criterion, repeat, with no boilerplate beyond defining your tasks and the threshold. +- **Preemption-safe state** — every completed iteration is durable before the next starts, so a killed job loses at most the in-flight iteration. +- **Tracking** — `tracking.backend: mlflow | clearml` in the spec wires a tracker once; every iteration's metrics, params, and lifecycle events are reported automatically, with no tracking code inside your task functions. +- **Heterogeneous dispatch** — the same loop drives CPU-only, GPU, MPI, or shell-executable tasks, and (via the `learners:` block) distinct task implementations per parallel learner. + +## What ROSE does **not** yet automate: data movement + +Today, getting a simulation's output into the training task's hands is entirely your task code's responsibility: + +- `type: python` slots pass the previous task's return value in-memory, as the first positional argument — this is implicit in the task type, not something you declare. +- `type: shell` slots pass nothing automatically; your command's stdout becomes the next task's input only if your scripts agree on a file path or convention outside the spec (this is exactly the pattern in the M3DC1 use case, where `simulate`/`train`/`active_learn` hand-roll a namespaced directory of files to communicate). + +**Planned:** a `DataExchangeProtocol` field in the YAML spec — `FileBased` or `MemoryBased` — that makes this an explicit, ROSE-managed choice instead of an implicit consequence of task type: + +- `MemoryBased` formalizes what `type: python` already does today — in-process object passing between tasks on the same worker. +- `FileBased` would let ROSE own staging a per-iteration (and per-learner, for parallel runs) working directory, instead of every use case reimplementing its own namespace/path convention by hand inside task code. + +This isn't implemented yet — it's the next piece of surface area on the spec, and the natural place to close the "no environment/output contract" gap without ROSE needing to know anything about the science moving through it. + +--- + +## What you bring, concretely + +| You supply | ROSE does not check this for you (today) | +|---|---| +| Task implementations (`simulate`, `train`, `active_learn`, criterion evaluator, ...) matching the `*args, **kwargs` convention | Whether return shapes match what the next task expects | +| The runtime environment those implementations need on the worker | No environment/dependency manifest in the spec | +| Access to wherever the edge agent's job runs (account, queue, allocation) | Out of scope for the spec layer entirely | +| Your data, and a convention for how tasks find it (today: `parameters:` + your own file paths) | No declared output/data contract yet — see `DataExchangeProtocol` above | +| Correct `remote.pythonpath` if task modules aren't already importable where the edge runs | `validate_imports=True` only checks importability in *your local* environment | + +None of this is unique to the service model — it's the same BYOF boundary as running ROSE locally. What the service model changes is *where* that boundary sits physically: your task code and its dependencies need to be reachable from the edge's job allocation, not from your laptop. diff --git a/docs/user-guide/spec-api.md b/docs/user-guide/spec-api.md new file mode 100644 index 00000000..7a98d766 --- /dev/null +++ b/docs/user-guide/spec-api.md @@ -0,0 +1,585 @@ +# YAML Spec API + +The YAML spec API lets you declare a ROSE workflow as a data file instead of Python code. The spec is schema-validated when loaded — missing slots, unknown keys, and type errors are caught before any infrastructure starts. Task functions are plain Python functions with no decorators required. The same spec file works with any ROSE learner type: Active Learning, Reinforcement Learning, or UQ. + +--- + +## Loading a Spec + +```python +from rose.spec import load_spec + +spec = load_spec("workflow.yaml") # validates schema on load +cfg = spec.config # typed SpecConfig object +``` + +`load_spec` raises `ValueError` with a precise message on any schema violation. It does **not** import task modules by default — see [Import Validation](#import-validation) to enable that check. + +--- + +## Sequential Active Learner + +
+
+ +**Python API** + +```python +from rose.al.active_learner import SequentialActiveLearner + +acl = SequentialActiveLearner(asyncflow) + +@acl.simulation_task(as_executable=False) +async def simulate(*args, **kwargs): + ... # return simulation result + +@acl.training_task(as_executable=False) +async def train(sim_result, **kwargs): + ... # return trained model + +@acl.active_learn_task(as_executable=False) +async def active_learn(sim_result, model, **kwargs): + ... # return updated dataset + +@acl.as_stop_criterion( + metric_name="mse", + threshold=0.01, + operator="<", +) +async def check_mse(*args, **kwargs): + ... # return float metric + +async for state in acl.start(max_iter=5): + print(f"iter {state.iteration} " + f"mse={state.metric_value:.4f}") + +await asyncflow.shutdown() +``` + +
+
+ +**YAML Spec** + +```yaml +learner: + type: sequential_active_learner + max_iter: 5 + +simulation: + type: python + function: tasks:simulate + +training: + type: python + function: tasks:train + +active_learn: + type: python + function: tasks:active_learn + +stop_criterion: + metric: mse + threshold: 0.01 + operator: "<" + evaluator: + type: python + function: tasks:check_mse +``` + +```python +# run it +from rose.spec import load_spec +from rose.spec.builder import LearnerBuilder + +cfg = load_spec("workflow.yaml").config +builder = LearnerBuilder(cfg, asyncflow) +learner = builder.build() + +async for state in learner.start( + max_iter=cfg.learner.max_iter +): + print(f"iter {state.iteration} " + f"mse={state.metric_value:.4f}") +``` + +
+
+ +Task functions referenced by `function: tasks:simulate` are ordinary Python callables in a `tasks.py` module — no ROSE decorators, no async required: + +```python +# tasks.py +def simulate(*args, **kwargs): + ... # return simulation result + +def train(sim_result, **kwargs): + ... # return trained model + +def active_learn(sim_result, model, **kwargs): + ... # return updated dataset + +def check_mse(*args, **kwargs): + ... # return float metric value +``` + +--- + +## Parallel Learner — Shared Tasks + +When all parallel learners run the same task implementations, declare the task slots at the top level exactly as in the sequential case. + +
+
+ +**Python API** + +```python +from rose.al.active_learner import ParallelActiveLearner + +pal = ParallelActiveLearner(asyncflow) + +@pal.simulation_task(as_executable=False) +async def simulate(*args, **kwargs): + ... + +@pal.training_task(as_executable=False) +async def train(sim_result, **kwargs): + ... + +@pal.active_learn_task(as_executable=False) +async def active_learn(sim_result, model, **kwargs): + ... + +@pal.as_stop_criterion( + metric_name="mse", + threshold=0.01, +) +async def check_mse(*args, **kwargs): + ... + +async for state in pal.start( + parallel_learners=2, + max_iter=5, +): + print(f"learner {state.learner_id} " + f"iter {state.iteration}") +``` + +
+
+ +**YAML Spec** + +```yaml +learner: + type: parallel_active_learner + max_iter: 5 + parallel_learners: 2 + +simulation: + type: python + function: tasks:simulate + +training: + type: python + function: tasks:train + +active_learn: + type: python + function: tasks:active_learn + +stop_criterion: + metric: mse + threshold: 0.01 + operator: "<" + evaluator: + type: python + function: tasks:check_mse +``` + +```python +# run it +cfg = load_spec("workflow.yaml").config +builder = LearnerBuilder(cfg, asyncflow) +learner = builder.build() + +async for state in learner.start( + max_iter=cfg.learner.max_iter, + parallel_learners=cfg.learner.parallel_learners, +): + print(f"learner {state.learner_id} " + f"iter {state.iteration}") +``` + +
+
+ +Each learner receives a unique `learner_id` kwarg (0, 1, …) in its task calls. Task functions can use it to write to isolated files or namespaces. + +--- + +## Parallel Learner — Distinct Tasks (`learners:` block) + +When each parallel learner needs its own task implementations — different models, different executables, different command-line flags — use the `learners:` block. Each entry gets a `label` that is injected as `learner_label` into every task kwarg. + +
+
+ +**Python API** + +```python +pal = ParallelActiveLearner(asyncflow) + +# Manual routing by learner_id +TRAIN_CMD = { + 0: "python train.py --model linear", + 1: "python train.py --model ridge", +} + +@pal.simulation_task +async def simulate(*args, **kwargs): + return "python sim.py" + +@pal.training_task +async def train(*args, **kwargs): + return TRAIN_CMD[kwargs["learner_id"]] + +@pal.active_learn_task +async def active_learn(*args, **kwargs): + return "python active_learn.py" + +@pal.as_stop_criterion( + metric_name="mse", threshold=0.01 +) +async def check_mse(*args, **kwargs): + ... + +# Build per-learner configs manually +from rose.learner import LearnerConfig, TaskConfig +lcs = [ + LearnerConfig( + simulation={i: TaskConfig(kwargs={"learner_id": i}) + for i in range(6)}, + ... + ) + for lid in range(2) +] + +async for state in pal.start( + parallel_learners=2, + max_iter=5, + learner_configs=lcs, +): + ... +``` + +
+
+ +**YAML Spec** + +```yaml +learner: + type: parallel_active_learner + max_iter: 5 + +learners: + - label: linear_regression + simulation: + type: shell + command: python sim.py --label a + training: + type: shell + command: python train.py --label a --model linear + active_learn: + type: shell + command: python active_learn.py --label a + + - label: ridge_regression + simulation: + type: shell + command: python sim.py --label b + training: + type: shell + command: python train.py --label b --model ridge + active_learn: + type: shell + command: python active_learn.py --label b + +stop_criterion: + metric: mse + threshold: 0.01 + operator: "<" + evaluator: + type: python + function: tasks:check_mse +``` + +```python +# run it — builder handles routing + LearnerConfig +cfg = load_spec("workflow.yaml").config +builder = LearnerBuilder(cfg, asyncflow) +learner = builder.build() +lcs = builder.build_learner_configs() + +async for state in learner.start( + max_iter=cfg.learner.max_iter, + parallel_learners=len(lcs), + learner_configs=lcs, +): + label = cfg.learners[state.learner_id].label + print(f"[{label}] iter {state.iteration}") +``` + +
+
+ +The builder creates a single dispatch closure per slot that routes to the correct command or function based on `learner_id`. The `learner_id` kwarg is consumed by the routing layer and never reaches the user function. + +!!! note + All learners in a `learners:` block must use the same task type (`python` or `shell`) for each slot. Mixed types within a slot are rejected at load time. + +--- + +## Task Types + +
+
+ +**`type: python`** + +Calls a Python function directly on the worker. Data flows in-memory between tasks. Function must be importable via `module:callable` syntax. + +```yaml +simulation: + type: python + function: my_package.tasks:simulate +``` + +```python +# tasks.py +def simulate(*args, **kwargs): + # receives positional results from + # prior tasks as *args + return {"X": ..., "y": ...} +``` + +The function may be sync or async. The return value is passed as the first positional argument to the next task in the chain. + +
+
+ +**`type: shell`** + +Runs a subprocess on the worker. Data flows through files on disk. The command string is the return value — ROSE submits it as an executable. + +```yaml +simulation: + type: shell + command: python sim.py +``` + +With `parameters:` placeholders: + +```yaml +simulation: + type: shell + command: python sim.py --dataset {dataset} +``` + +`{dataset}` is filled from `parameters.dataset` at runtime via `str.format_map(kwargs)`. + +The criterion evaluator must print a single float to stdout — ROSE captures stdout as the metric value. + +
+
+ +### Resource hints (`task_description`) + +Both task types accept an optional `task_description` dict forwarded to the execution backend as resource hints (CPU count, GPU count, etc.): + +```yaml +simulation: + type: python + function: tasks:simulate + task_description: + cpu_count: 4 + gpu_count: 1 +``` + +!!! note + For parallel learners using the `learners:` block, all entries must use the same `task_description` for each slot — the backend registers it once at task-registration time. + +--- + +## Stop Criterion + +
+
+ +**Python API** + +```python +@acl.as_stop_criterion( + metric_name="mse", + threshold=0.01, + operator="<", + as_executable=False, +) +async def check_mse(*args, **kwargs): + ... # return float +``` + +
+
+ +**YAML Spec** + +```yaml +stop_criterion: + metric: mse + threshold: 0.01 + operator: "<" # default; also: >, ==, <=, >= + evaluator: + type: python + function: tasks:check_mse +``` + +
+
+ +The evaluator function follows the same `*args, **kwargs` convention as task functions. ROSE stops the loop when `metric_value threshold` is satisfied, or when `max_iter` is reached — whichever comes first. + +--- + +## Parameters + +The `parameters:` block defines key-value pairs that are injected into every task's `**kwargs` at every iteration: + +```yaml +parameters: + dataset: my_dataset + batch_size: 32 + growing_pool: false +``` + +```python +def simulate(*args, **kwargs): + dataset = kwargs["dataset"] # "my_dataset" + batch_size = kwargs["batch_size"] # 32 + growing_pool = kwargs["growing_pool"] # False + iteration = kwargs["iteration"] # 0, 1, 2, ... + ... +``` + +### Reserved keys + +The following keys are injected automatically by the builder and must not appear in `parameters:`: + +| Key | Available in | Description | +|-----|-------------|-------------| +| `iteration` | all tasks | Current loop counter (0-based) | +| `pythonpath` | all tasks | Contents of `remote.pythonpath` as a list | +| `learner_id` | parallel tasks | Integer index of the learner (0, 1, …) | +| `learner_label` | parallel tasks (when `label` set) | Human-readable learner name from `learners[].label` | + +--- + +## Remote Config + +```yaml +remote: + pythonpath: + - /path/to/my/task/modules + - /path/to/shared/utilities +``` + +`remote.pythonpath` entries are added to `sys.path` on the remote worker before any task module is imported. They are also injected into every task's `kwargs["pythonpath"]` (as a list), so task functions can construct file paths without duplicating the value in `parameters:`: + +```python +def train(sim_result, **kwargs): + pythonpath = kwargs.get("pythonpath", []) + base = pythonpath[0] if pythonpath else "" + script = Path(base) / "scripts" / "train_model.py" + ... +``` + +`remote.pythonpath` is the single edit point for the remote worker path — changing it updates both the import path and the kwarg. + +--- + +## Tracking + +```yaml +tracking: + backend: mlflow # mlflow | clearml | none (default) + experiment: my-exp # experiment/project name + run_name: run-01 # optional run label +``` + +See [MLflow integration](../integrations/mlflow.md) and [ClearML integration](../integrations/clearml.md) for configuration details. + +--- + +## Spec Variants with `workflow_with()` + +`workflow_with()` returns a new `WorkflowSpec` with selective overrides applied. The original spec is never mutated: + +```python +base_spec = load_spec("workflow.yaml") +test_spec = base_spec.workflow_with(max_iter=2, parameters={"dataset": "test_ds"}) +large_spec = base_spec.workflow_with(max_iter=50, parameters={"batch_size": 128}) +``` + +Accepted override keys: + +- `parameters` — merged (not replaced) into the existing `parameters:` block +- Any `learner` field: `max_iter`, `parallel_learners` +- Any other top-level spec field + +Unknown keys raise `ValueError` immediately. + +--- + +## Import Validation + +By default, `load_spec` does not import task modules — this is intentional when `remote.pythonpath` points to paths that only exist on the remote worker. To catch `function:` typos locally during development: + +```python +spec = load_spec("workflow.yaml", validate_imports=True) +``` + +With `validate_imports=True`, every `module:callable` string is resolved in the current environment. All failures are collected and reported in a single `ValueError` before any infrastructure starts. Use this during development when task files are locally accessible. + +--- + +## Spec Reference + +### Top-level keys + +| Key | Type | Required | Default | Description | +|-----|------|----------|---------|-------------| +| `learner` | object | yes | — | Learner type and loop settings | +| `learner.type` | string | yes | — | `sequential_active_learner`, `parallel_active_learner`, `sequential_reinforcement_learner`, `uq_active_learner` | +| `learner.max_iter` | int | no | `0` | Maximum iterations | +| `learner.parallel_learners` | int | no | `2` | Parallel learner count when `learners:` is absent | +| `simulation` / `training` / `active_learn` | TaskDef | AL yes | — | Task slots for Active Learning | +| `environment` / `update` | TaskDef | RL yes | — | Task slots for Reinforcement Learning | +| `prediction` / `active_learn` / `uncertainty` | TaskDef | UQ yes | — | Task slots for UQ-based AL | +| `learners` | list | no | — | Per-learner task definitions for heterogeneous parallel | +| `stop_criterion` | object | yes | — | Stopping condition | +| `parameters` | dict | no | `{}` | User-defined kwargs injected into all tasks | +| `remote.pythonpath` | list[str] | no | `[]` | Paths added to `sys.path` on worker; injected as `pythonpath` kwarg | +| `tracking.backend` | string | no | `none` | `mlflow` / `clearml` / `none` | +| `tracking.experiment` | string | no | `ROSE-Spec` | Experiment name | +| `tracking.run_name` | string | no | `null` | Run label | + +### TaskDef fields + +| Key | Type | Required | Description | +|-----|------|----------|-------------| +| `type` | `python` \| `shell` | yes | Task execution mode | +| `function` | string | if `type: python` | `module:callable` — dotted module path and function name | +| `command` | string | if `type: shell` | Shell command; supports `{param}` placeholders filled from `parameters:` | +| `task_description` | dict | no | Resource hints forwarded to the execution backend | diff --git a/mkdocs.yml b/mkdocs.yml index c35fd6a4..f4873f74 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -75,6 +75,8 @@ markdown_extensions: emoji_generator: !!python/name:material.extensions.emoji.to_svg - pymdownx.snippets: check_paths: true + - pymdownx.tabbed: + alternate_style: true - attr_list - md_in_html @@ -95,6 +97,8 @@ nav: - 8.Experience Banks: user-guide/experience.md - 9.Advanced RL workflow: user-guide/advanced-rl-workflow.md - 10.Tracking & Observability: user-guide/tracking.md + - 11.YAML Spec API: user-guide/spec-api.md + - 12.ROSE as a Service: user-guide/rose-aas.md - Integrations: - MLflow: integrations/mlflow.md - ClearML: integrations/clearml.md diff --git a/rose/spec/__init__.py b/rose/spec/__init__.py index 6852be15..de2b44f5 100644 --- a/rose/spec/__init__.py +++ b/rose/spec/__init__.py @@ -9,9 +9,67 @@ __all__ = ["load_spec", "WorkflowSpec"] -def load_spec(path: str | Path) -> "WorkflowSpec": - """Load and validate a YAML workflow spec. Raises ValueError on schema errors.""" - return WorkflowSpec(SpecConfig.from_yaml(path)) +def load_spec(path: str | Path, validate_imports: bool = False) -> "WorkflowSpec": + """Load and validate a YAML workflow spec. Raises ValueError on schema errors. + + Args: + path: Path to the YAML file. + validate_imports: If True, verify that every python task's ``module:callable`` + is importable from the current environment (sys.path extended with + ``remote.pythonpath``). Use during development when task files are + locally accessible. Leave False (default) when ``remote.pythonpath`` + points to paths that only exist on the remote worker. + """ + spec = WorkflowSpec(SpecConfig.from_yaml(path)) + if validate_imports: + _validate_task_imports(spec.config) + return spec + + +def _collect_python_specs(cfg: "SpecConfig") -> list[str]: + """Return deduplicated 'module:callable' strings for all Python tasks in the spec.""" + specs: list[str] = [] + if cfg.learners is not None: + for l in cfg.learners: + for td in l.tasks.values(): + if td.type == "python": + specs.append(td.function) + else: + for td in cfg.tasks.values(): + if td.type == "python": + specs.append(td.function) + if cfg.stop_criterion.evaluator.type == "python": + specs.append(cfg.stop_criterion.evaluator.function) + return list(dict.fromkeys(specs)) # deduplicate, preserve order + + +def _validate_task_imports(cfg: "SpecConfig") -> None: + """Try to import every Python task callable. Raises ValueError listing all failures.""" + import importlib + import sys as _sys + + added = [p for p in cfg.remote.pythonpath if p not in _sys.path] + for p in added: + _sys.path.insert(0, p) + try: + failures: list[str] = [] + for spec_str in _collect_python_specs(cfg): + mod_path, fn_name = spec_str.rsplit(":", 1) + try: + mod = importlib.import_module(mod_path) + if not hasattr(mod, fn_name): + failures.append(f" {spec_str!r}: module has no attribute '{fn_name}'") + except ImportError as exc: + failures.append(f" {spec_str!r}: {exc}") + if failures: + raise ValueError( + "validate_imports failed — the following function specs are not importable:\n" + + "\n".join(failures) + ) + finally: + for p in added: + if p in _sys.path: + _sys.path.remove(p) class WorkflowSpec: @@ -20,6 +78,31 @@ class WorkflowSpec: def __init__(self, config: SpecConfig) -> None: self.config = config + def workflow_with(self, **overrides: Any) -> "WorkflowSpec": + """Return a new WorkflowSpec with selective overrides applied. + + Accepted keys: + - ``parameters``: dict merged (not replaced) into the existing parameters block + - Any ``LearnerSpec`` field: ``max_iter``, ``parallel_learners`` + - Any other top-level ``SpecConfig`` field + + Example:: + + spec.workflow_with(max_iter=3, parameters={"dataset": "test_ds"}) + """ + data = self.config.model_dump() + learner_fields = set(data["learner"].keys()) + for key, value in overrides.items(): + if key == "parameters" and isinstance(value, dict): + data["parameters"] = {**data["parameters"], **value} + elif key in learner_fields: + data["learner"][key] = value + elif key in data: + data[key] = value + else: + raise ValueError(f"workflow_with: unknown spec field '{key}'") + return WorkflowSpec(SpecConfig.model_validate(data)) + @property def workflow(self) -> Callable[..., Coroutine[Any, Any, None]]: cfg = self.config @@ -38,12 +121,19 @@ async def _workflow(bridge_url: str, edge_name: str) -> None: start_kwargs: dict[str, Any] = {"max_iter": cfg.learner.max_iter} if cfg.learner.type == "parallel_active_learner": - lc = builder.build_learner_configs() - if lc is not None: - start_kwargs["parallel_learners"] = len(lc) - start_kwargs["learner_configs"] = lc + lcs = builder.build_learner_configs() + if lcs is not None: + start_kwargs["parallel_learners"] = len(lcs) + start_kwargs["learner_configs"] = lcs else: start_kwargs["parallel_learners"] = cfg.learner.parallel_learners + else: + # Sequential learners accept initial_config — the same ROSE-native + # mechanism used by the parallel path, giving tasks access to + # parameters and iteration via kwargs at every iteration. + ic = builder.build_learner_config() + if ic is not None: + start_kwargs["initial_config"] = ic try: async for state in learner.start(**start_kwargs): diff --git a/rose/spec/adapters.py b/rose/spec/adapters.py index 06fb6eaf..e5480f1d 100644 --- a/rose/spec/adapters.py +++ b/rose/spec/adapters.py @@ -23,8 +23,8 @@ def make_dispatch_closure( task_defs: list[TaskDef], remote: RemoteConfig, ) -> Callable: - """Dispatch closure for parallel learners: routes per-candidate based on learner_id kwarg. - task_description uses the first candidate's value — asyncflow reads it once at registration.""" + """Dispatch closure for parallel learners: routes per-learner based on learner_id kwarg. + task_description uses the first learner's value — asyncflow reads it once at registration.""" td = dict(task_defs[0].task_description or {}) if task_defs[0].type == "shell": return _make_shell_dispatch( @@ -43,7 +43,7 @@ def _make_shell_closure(command: str, task_description: dict) -> Callable: _td = task_description async def _task(*args, task_description=_td, **kwargs) -> str: - return _cmd + return _cmd.format_map(kwargs) _task.__name__ = "shell_task" return _task @@ -76,7 +76,8 @@ def _make_shell_dispatch(cmds: dict[int, str], slot_name: str, task_description: _td = task_description async def _dispatch(*args, task_description=_td, **kwargs) -> str: - return _cmds[kwargs.get("learner_id", 0)] + cmd = _cmds[kwargs.get("learner_id", 0)] + return cmd.format_map(kwargs) _dispatch.__name__ = f"{slot_name}_dispatch" return _dispatch diff --git a/rose/spec/builder.py b/rose/spec/builder.py index 2f0375a1..f838a3f3 100644 --- a/rose/spec/builder.py +++ b/rose/spec/builder.py @@ -5,6 +5,12 @@ from .adapters import TaskAdapterFactory from .schema import SpecConfig, TrackingConfig, _REQUIRED_SLOTS +# Slots that exist as fields on LearnerConfig; uncertainty is required by uq_active_learner +# but is not a LearnerConfig field — filter it out when constructing LearnerConfig. +_LEARNER_CONFIG_SLOTS: frozenset[str] = frozenset( + {"simulation", "training", "prediction", "active_learn", "environment", "update"} +) + def _get_learner_class(learner_type: str): if learner_type == "sequential_active_learner": @@ -35,7 +41,7 @@ def build(self) -> Learner: cfg = self.config learner = _get_learner_class(cfg.learner.type)(self.asyncflow) - if cfg.candidates is not None: + if cfg.learners is not None: self._register_dispatched_tasks(learner, cfg) else: self._register_flat_tasks(learner, cfg) @@ -62,33 +68,57 @@ def _register_flat_tasks(self, learner: Learner, cfg: SpecConfig) -> None: def _register_dispatched_tasks(self, learner: Learner, cfg: SpecConfig) -> None: required = _REQUIRED_SLOTS[cfg.learner.type] for slot_name in required: - task_defs = [c.tasks[slot_name] for c in cfg.candidates] + task_defs = [l.tasks[slot_name] for l in cfg.learners] as_exec = task_defs[0].type == "shell" closure = TaskAdapterFactory.make_dispatch_closure(slot_name, task_defs, cfg.remote) getattr(learner, f"{slot_name}_task")(as_executable=as_exec)(closure) + def build_learner_config(self): + """Single LearnerConfig for sequential learners. + + Passed as initial_config to learner.start() so that parameters and + iteration reach task kwargs via the same ROSE-native mechanism used + by the parallel path. Returns None when no parameters are defined. + """ + cfg = self.config + if not cfg.parameters and not cfg.remote.pythonpath: + return None + from rose.learner import LearnerConfig, TaskConfig + + required = _REQUIRED_SLOTS[cfg.learner.type] + max_iter = cfg.learner.max_iter + params = dict(cfg.parameters) + params["pythonpath"] = list(cfg.remote.pythonpath) + schedule = {n: TaskConfig(kwargs={**params, "iteration": n}) for n in range(max_iter + 1)} + schedule[-1] = TaskConfig(kwargs={**params, "iteration": max_iter}) + lc_slots = required & _LEARNER_CONFIG_SLOTS + return LearnerConfig(**{slot: schedule for slot in lc_slots}, criterion=schedule) + def build_learner_configs(self): - """Return auto-generated LearnerConfig list for parallel candidates, or None.""" + """LearnerConfig list for parallel learners, or None if no learners defined.""" cfg = self.config - if cfg.candidates is None: + if cfg.learners is None: return None from rose.learner import LearnerConfig, TaskConfig required = _REQUIRED_SLOTS[cfg.learner.type] configs = [] max_iter = cfg.learner.max_iter - for i, candidate in enumerate(cfg.candidates): - # Build per-iteration TaskConfig kwargs, mirroring how ROSE expects LearnerConfig to work. - # learner_id → dispatch routing key (popped by dispatch closure, never reaches user fn) - # iteration → current iteration n, placed per-entry so get_task_config(slot, n) returns it - # learner_label → human-readable candidate name; only added when candidate has a label - base = {"learner_id": i} - if candidate.label: - base["learner_label"] = candidate.label + params = dict(cfg.parameters) + params["pythonpath"] = list(cfg.remote.pythonpath) + lc_slots = required & _LEARNER_CONFIG_SLOTS + for i, learner_def in enumerate(cfg.learners): + # learner_id → dispatch routing key (popped by closure, never reaches user fn) + # iteration → per-entry so get_task_config(slot, n) returns the right value + # learner_label → human-readable learner name; only injected when non-empty + # parameters → user-defined values from the YAML parameters: block + base = {"learner_id": i, **params} + if learner_def.label: + base["learner_label"] = learner_def.label schedule = {n: TaskConfig(kwargs={**base, "iteration": n}) for n in range(max_iter + 1)} schedule[-1] = TaskConfig(kwargs={**base, "iteration": max_iter}) configs.append( - LearnerConfig(**{slot: schedule for slot in required}, criterion=schedule) + LearnerConfig(**{slot: schedule for slot in lc_slots}, criterion=schedule) ) return configs diff --git a/rose/spec/schema.py b/rose/spec/schema.py index 0e23105e..57d93d46 100644 --- a/rose/spec/schema.py +++ b/rose/spec/schema.py @@ -6,6 +6,8 @@ from pydantic import BaseModel, Field, model_validator + + # ── Task definition ─────────────────────────────────────────────────────────── class TaskDef(BaseModel): type: Literal["shell", "python"] @@ -51,13 +53,13 @@ class StopCriterionDef(BaseModel): class LearnerSpec(BaseModel): type: str max_iter: int = 0 - parallel_learners: int = 2 # used only when candidates is absent for parallel types + parallel_learners: int = 2 # used only when learners: is absent for parallel types model_config = {"extra": "forbid"} -# ── Per-candidate definition (parallel learners only) ───────────────────────── -class CandidateDef(BaseModel): +# ── Per-learner definition (parallel learners only) ─────────────────────────── +class LearnerDef(BaseModel): label: str = "" simulation: TaskDef | None = None training: TaskDef | None = None @@ -89,6 +91,12 @@ class TrackingConfig(BaseModel): model_config = {"extra": "forbid"} +# Keys the builder always injects — users must not put these in parameters: +_RESERVED_PARAMETER_KEYS: frozenset[str] = frozenset( + {"learner_id", "learner_label", "iteration", "pythonpath"} +) + + # ── Top-level spec ──────────────────────────────────────────────────────────── class SpecConfig(BaseModel): learner: LearnerSpec @@ -101,10 +109,11 @@ class SpecConfig(BaseModel): update: TaskDef | None = None prediction: TaskDef | None = None uncertainty: TaskDef | None = None - candidates: list[CandidateDef] | None = None + learners: list[LearnerDef] | None = None stop_criterion: StopCriterionDef - remote: RemoteConfig = Field(default_factory=RemoteConfig) - tracking: TrackingConfig = Field(default_factory=TrackingConfig) + parameters: dict[str, Any] = Field(default_factory=dict) + remote: RemoteConfig = Field(default_factory=RemoteConfig) + tracking: TrackingConfig = Field(default_factory=TrackingConfig) model_config = {"extra": "forbid"} @@ -123,20 +132,28 @@ def _validate_task_slots(self) -> "SpecConfig": ) is_parallel = ltype == "parallel_active_learner" - if is_parallel and self.candidates is not None: - for c in self.candidates: - missing = required - set(c.tasks.keys()) + if is_parallel and self.learners is not None: + for l in self.learners: + missing = required - set(l.tasks.keys()) if missing: - raise ValueError(f"Candidate '{c.label}' missing: {sorted(missing)}") - extra = set(c.tasks.keys()) - required + raise ValueError(f"Learner '{l.label}' missing: {sorted(missing)}") + extra = set(l.tasks.keys()) - required if extra: - raise ValueError(f"Candidate '{c.label}' unexpected fields: {sorted(extra)}") + raise ValueError(f"Learner '{l.label}' unexpected fields: {sorted(extra)}") for slot in required: - types = {c.tasks[slot].type for c in self.candidates} + types = {l.tasks[slot].type for l in self.learners} if len(types) > 1: raise ValueError( - f"Slot '{slot}' has mixed types across candidates: {types}. " - "All candidates must use the same type for a given slot." + f"Slot '{slot}' has mixed types across learners: {types}. " + "All learners must use the same type for a given slot." + ) + descs = [dict(l.tasks[slot].task_description or {}) for l in self.learners] + if len({str(sorted(d.items())) for d in descs}) > 1: + raise ValueError( + f"Slot '{slot}' has different task_description values across learners. " + "asyncflow registers task_description once per slot at registration time — " + "all learners must use the same value. Use identical task_description " + "across all learners or omit it from all but the first." ) else: present = set(self.tasks.keys()) @@ -148,6 +165,16 @@ def _validate_task_slots(self) -> "SpecConfig": raise ValueError(f"Unexpected task fields for '{ltype}': {sorted(extra)}") return self + @model_validator(mode="after") + def _validate_parameters(self) -> "SpecConfig": + conflicts = _RESERVED_PARAMETER_KEYS & set(self.parameters.keys()) + if conflicts: + raise ValueError( + f"'parameters' must not use reserved keys: {sorted(conflicts)}. " + "These are injected automatically by the spec builder." + ) + return self + @classmethod def from_yaml(cls, path: str | Path) -> "SpecConfig": import yaml diff --git a/tests/integration/spec/helpers.py b/tests/integration/spec/helpers.py index d8ef3090..1d4496c4 100644 --- a/tests/integration/spec/helpers.py +++ b/tests/integration/spec/helpers.py @@ -1,5 +1,8 @@ """Lightweight task helpers used by the YAML spec integration tests.""" +# Shared capture list — cleared by tests that need to inspect received kwargs. +received_kwargs: list[dict] = [] + async def sim(*args, **kwargs): return [1.0, 2.0, 3.0] @@ -15,3 +18,22 @@ async def active_learn(sim_result, model, **kwargs): async def criterion(*args, **kwargs): return 0.05 + + +# ── Parameter-capturing variants ────────────────────────────────────────────── + +async def sim_capture(*args, **kwargs): + received_kwargs.append(dict(kwargs)) + return [1.0, 2.0, 3.0] + + +async def train_capture(data, **kwargs): + return {"mean": sum(data) / len(data)} + + +async def active_learn_capture(sim_result, model, **kwargs): + return abs(model["mean"] - 2.0) + + +async def criterion_capture(*args, **kwargs): + return 0.05 diff --git a/tests/integration/spec/test_yaml_workflow.py b/tests/integration/spec/test_yaml_workflow.py index 084de7c2..3c60ccc8 100644 --- a/tests/integration/spec/test_yaml_workflow.py +++ b/tests/integration/spec/test_yaml_workflow.py @@ -71,3 +71,144 @@ async def test_load_spec_returns_workflow_spec(tmp_path): spec = load_spec(p) assert hasattr(spec, "workflow") assert callable(spec.workflow) + + +# ── Sequential AL with parameters: block ───────────────────────────────────── + +SEQ_WITH_PARAMS_YAML = textwrap.dedent("""\ + learner: + type: sequential_active_learner + max_iter: 2 + + simulation: + type: python + function: tests.integration.spec.helpers:sim_capture + + training: + type: python + function: tests.integration.spec.helpers:train_capture + + active_learn: + type: python + function: tests.integration.spec.helpers:active_learn_capture + + stop_criterion: + metric: mse + threshold: 0.01 + operator: "<" + evaluator: + type: python + function: tests.integration.spec.helpers:criterion_capture + + parameters: + dataset: test_ds + scale: 1.5 +""") + + +@pytest.mark.asyncio +@pytest.mark.integration +async def test_sequential_workflow_parameters_reach_tasks(tmp_path): + import tests.integration.spec.helpers as helpers + helpers.received_kwargs.clear() + + p = tmp_path / "spec.yaml" + p.write_text(SEQ_WITH_PARAMS_YAML) + + engine = await ConcurrentExecutionBackend(ThreadPoolExecutor()) + asyncflow = await WorkflowEngine.create(engine) + + cfg = SpecConfig.from_yaml(p) + builder = LearnerBuilder(cfg, asyncflow) + learner = builder.build() + + ic = builder.build_learner_config() + start_kwargs = {"max_iter": 2} + if ic is not None: + start_kwargs["initial_config"] = ic + + async for _ in learner.start(**start_kwargs): + pass + + await asyncflow.shutdown() + + assert len(helpers.received_kwargs) > 0, "sim_capture was never called" + for kw in helpers.received_kwargs: + assert kw.get("dataset") == "test_ds", f"expected dataset in kwargs, got {kw}" + assert kw.get("scale") == 1.5, f"expected scale in kwargs, got {kw}" + + +# ── Parallel AL with parameters: block ─────────────────────────────────────── + +PAR_WITH_PARAMS_YAML = textwrap.dedent("""\ + learner: + type: parallel_active_learner + max_iter: 2 + + learners: + - label: rf + simulation: + type: python + function: tests.integration.spec.helpers:sim_capture + training: + type: python + function: tests.integration.spec.helpers:train_capture + active_learn: + type: python + function: tests.integration.spec.helpers:active_learn_capture + - label: mlp + simulation: + type: python + function: tests.integration.spec.helpers:sim_capture + training: + type: python + function: tests.integration.spec.helpers:train_capture + active_learn: + type: python + function: tests.integration.spec.helpers:active_learn_capture + + stop_criterion: + metric: mse + threshold: 0.01 + operator: "<" + evaluator: + type: python + function: tests.integration.spec.helpers:criterion_capture + + parameters: + dataset: par_ds + lr: 0.01 +""") + + +@pytest.mark.asyncio +@pytest.mark.integration +async def test_parallel_workflow_parameters_reach_tasks(tmp_path): + import tests.integration.spec.helpers as helpers + helpers.received_kwargs.clear() + + p = tmp_path / "spec.yaml" + p.write_text(PAR_WITH_PARAMS_YAML) + + engine = await ConcurrentExecutionBackend(ThreadPoolExecutor()) + asyncflow = await WorkflowEngine.create(engine) + + cfg = SpecConfig.from_yaml(p) + builder = LearnerBuilder(cfg, asyncflow) + learner = builder.build() + + lcs = builder.build_learner_configs() + start_kwargs = {"max_iter": 2, "parallel_learners": len(lcs), "learner_configs": lcs} + + async for _ in learner.start(**start_kwargs): + pass + + await asyncflow.shutdown() + + assert len(helpers.received_kwargs) > 0, "sim_capture was never called" + for kw in helpers.received_kwargs: + assert kw.get("dataset") == "par_ds", f"expected dataset in kwargs, got {kw}" + assert kw.get("lr") == 0.01, f"expected lr in kwargs, got {kw}" + # learner_id is stripped by the dispatch closure; learner_label passes through + assert "learner_id" not in kw, "learner_id should be stripped by dispatch" + assert kw.get("learner_label") in {"rf", "mlp"} diff --git a/tests/unit/spec/test_adapters.py b/tests/unit/spec/test_adapters.py index bd87fca8..4de803bb 100644 --- a/tests/unit/spec/test_adapters.py +++ b/tests/unit/spec/test_adapters.py @@ -121,6 +121,34 @@ def test_python_closure_task_description_in_signature(): assert sig.parameters["task_description"].default == {"gpu_count": 2} +# ── Shell command interpolation ─────────────────────────────────────────────── + +def test_shell_closure_interpolates_kwargs(): + td = TaskDef(type="shell", command="python sim.py --dataset {dataset} --n {n}") + fn = TaskAdapterFactory.make_closure(td, _remote()) + result = asyncio.run(fn(dataset="m3dc1", n=42)) + assert result == "python sim.py --dataset m3dc1 --n 42" + + +def test_shell_closure_no_placeholders_unchanged(): + td = TaskDef(type="shell", command="python sim.py") + fn = TaskAdapterFactory.make_closure(td, _remote()) + result = asyncio.run(fn(dataset="ignored")) + assert result == "python sim.py" + + +def test_shell_dispatch_interpolates_kwargs(): + tds = [ + TaskDef(type="shell", command="python sim.py --model rf --dataset {dataset}"), + TaskDef(type="shell", command="python sim.py --model mlp --dataset {dataset}"), + ] + fn = TaskAdapterFactory.make_dispatch_closure("simulation", tds, _remote()) + assert asyncio.run(fn(learner_id=0, dataset="m3dc1")) == \ + "python sim.py --model rf --dataset m3dc1" + assert asyncio.run(fn(learner_id=1, dataset="test_ds")) == \ + "python sim.py --model mlp --dataset test_ds" + + def test_dispatch_closure_task_description_uses_first_candidate(): tds = [ TaskDef(type="shell", command="cmd_0", task_description={"cpu_count": 8}), diff --git a/tests/unit/spec/test_builder.py b/tests/unit/spec/test_builder.py new file mode 100644 index 00000000..5fce0581 --- /dev/null +++ b/tests/unit/spec/test_builder.py @@ -0,0 +1,261 @@ +"""Unit tests for rose.spec.builder — LearnerConfig construction, no HPC needed.""" +import textwrap +from unittest.mock import MagicMock + +import pytest + +from rose.spec.schema import SpecConfig + + +SEQ_WITH_PARAMS = textwrap.dedent("""\ + learner: + type: sequential_active_learner + max_iter: 3 + simulation: + type: python + function: mymod:sim + training: + type: python + function: mymod:train + active_learn: + type: python + function: mymod:select + stop_criterion: + metric: mse + threshold: 0.1 + evaluator: + type: python + function: mymod:eval + parameters: + dataset: test_ds + scale: 2.5 +""") + +SEQ_NO_PARAMS = textwrap.dedent("""\ + learner: + type: sequential_active_learner + max_iter: 3 + simulation: + type: python + function: mymod:sim + training: + type: python + function: mymod:train + active_learn: + type: python + function: mymod:select + stop_criterion: + metric: mse + threshold: 0.1 + evaluator: + type: python + function: mymod:eval +""") + +PAR_WITH_PARAMS = textwrap.dedent("""\ + learner: + type: parallel_active_learner + max_iter: 2 + learners: + - label: rf + simulation: + type: python + function: mymod:sim + training: + type: python + function: mymod:train + active_learn: + type: python + function: mymod:select + - label: mlp + simulation: + type: python + function: mymod:sim + training: + type: python + function: mymod:train + active_learn: + type: python + function: mymod:select + stop_criterion: + metric: r2 + threshold: 0.9 + operator: ">" + evaluator: + type: python + function: mymod:eval + parameters: + dataset: prod_ds + lr: 0.001 +""") + + +def _make_builder(yaml_text, tmp_path): + p = tmp_path / "spec.yaml" + p.write_text(yaml_text) + cfg = SpecConfig.from_yaml(p) + from rose.spec.builder import LearnerBuilder + return LearnerBuilder(cfg, MagicMock()) + + +# ── build_learner_config (sequential path) ──────────────────────────────────── + +def test_build_learner_config_no_parameters_returns_none(tmp_path): + builder = _make_builder(SEQ_NO_PARAMS, tmp_path) + assert builder.build_learner_config() is None + + +def test_build_learner_config_returns_learner_config(tmp_path): + from rose.learner import LearnerConfig + builder = _make_builder(SEQ_WITH_PARAMS, tmp_path) + lc = builder.build_learner_config() + assert lc is not None + assert isinstance(lc, LearnerConfig) + + +def test_build_learner_config_per_iteration_kwargs(tmp_path): + builder = _make_builder(SEQ_WITH_PARAMS, tmp_path) + lc = builder.build_learner_config() + # iteration 0: parameters + iteration=0 + kw0 = lc.simulation[0].kwargs + assert kw0["dataset"] == "test_ds" + assert kw0["scale"] == 2.5 + assert kw0["iteration"] == 0 + # iteration 3 (max_iter): parameters + iteration=3 + kw3 = lc.simulation[3].kwargs + assert kw3["dataset"] == "test_ds" + assert kw3["iteration"] == 3 + + +def test_build_learner_config_criterion_schedule_included(tmp_path): + builder = _make_builder(SEQ_WITH_PARAMS, tmp_path) + lc = builder.build_learner_config() + assert lc.criterion is not None + assert lc.criterion[0].kwargs["dataset"] == "test_ds" + + +# ── build_learner_configs (parallel path) ──────────────────────────────────── + +def test_build_learner_configs_no_learners_returns_none(tmp_path): + builder = _make_builder(SEQ_WITH_PARAMS, tmp_path) + assert builder.build_learner_configs() is None + + +def test_build_learner_configs_with_parameters(tmp_path): + from rose.learner import LearnerConfig + builder = _make_builder(PAR_WITH_PARAMS, tmp_path) + lcs = builder.build_learner_configs() + assert lcs is not None + assert len(lcs) == 2 + assert all(isinstance(lc, LearnerConfig) for lc in lcs) + + +def test_build_learner_configs_learner_0_kwargs(tmp_path): + builder = _make_builder(PAR_WITH_PARAMS, tmp_path) + lcs = builder.build_learner_configs() + kw = lcs[0].simulation[0].kwargs + assert kw["learner_id"] == 0 + assert kw["learner_label"] == "rf" + assert kw["dataset"] == "prod_ds" + assert kw["lr"] == 0.001 + assert kw["iteration"] == 0 + + +def test_build_learner_configs_learner_1_kwargs(tmp_path): + builder = _make_builder(PAR_WITH_PARAMS, tmp_path) + lcs = builder.build_learner_configs() + kw = lcs[1].simulation[1].kwargs + assert kw["learner_id"] == 1 + assert kw["learner_label"] == "mlp" + assert kw["dataset"] == "prod_ds" + assert kw["iteration"] == 1 + + +# ── pythonpath auto-injection ───────────────────────────────────────────────── + +SEQ_WITH_PYTHONPATH = textwrap.dedent("""\ + learner: + type: sequential_active_learner + max_iter: 2 + simulation: + type: python + function: mymod:sim + training: + type: python + function: mymod:train + active_learn: + type: python + function: mymod:select + stop_criterion: + metric: mse + threshold: 0.1 + evaluator: + type: python + function: mymod:eval + remote: + pythonpath: + - /remote/path/a + - /remote/path/b +""") + +PAR_WITH_PYTHONPATH = textwrap.dedent("""\ + learner: + type: parallel_active_learner + max_iter: 2 + learners: + - label: rf + simulation: + type: python + function: mymod:sim + training: + type: python + function: mymod:train + active_learn: + type: python + function: mymod:select + - label: mlp + simulation: + type: python + function: mymod:sim + training: + type: python + function: mymod:train + active_learn: + type: python + function: mymod:select + stop_criterion: + metric: r2 + threshold: 0.9 + operator: ">" + evaluator: + type: python + function: mymod:eval + remote: + pythonpath: + - /remote/path/a +""") + + +def test_build_learner_config_injects_pythonpath(tmp_path): + builder = _make_builder(SEQ_WITH_PYTHONPATH, tmp_path) + lc = builder.build_learner_config() + assert lc is not None + assert lc.simulation[0].kwargs["pythonpath"] == ["/remote/path/a", "/remote/path/b"] + + +def test_build_learner_config_pythonpath_empty_when_no_remote(tmp_path): + builder = _make_builder(SEQ_WITH_PARAMS, tmp_path) + lc = builder.build_learner_config() + assert lc.simulation[0].kwargs["pythonpath"] == [] + + +def test_build_learner_configs_injects_pythonpath(tmp_path): + builder = _make_builder(PAR_WITH_PYTHONPATH, tmp_path) + lcs = builder.build_learner_configs() + assert lcs[0].simulation[0].kwargs["pythonpath"] == ["/remote/path/a"] + assert lcs[1].simulation[0].kwargs["pythonpath"] == ["/remote/path/a"] + + +def test_build_learner_config_no_params_no_pythonpath_returns_none(tmp_path): + builder = _make_builder(SEQ_NO_PARAMS, tmp_path) + assert builder.build_learner_config() is None diff --git a/tests/unit/spec/test_schema.py b/tests/unit/spec/test_schema.py index ca6d748f..efb77c66 100644 --- a/tests/unit/spec/test_schema.py +++ b/tests/unit/spec/test_schema.py @@ -196,14 +196,14 @@ def test_unknown_learner_type(tmp_path): SpecConfig.from_yaml(p) -# ── SpecConfig — parallel learner with candidates ──────────────────────────── +# ── SpecConfig — parallel learner with learners ────────────────────────────── PAR_YAML = textwrap.dedent("""\ learner: type: parallel_active_learner max_iter: 3 - candidates: + learners: - label: rf simulation: type: shell @@ -235,20 +235,20 @@ def test_unknown_learner_type(tmp_path): """) -def test_parallel_candidates_roundtrip(tmp_path): +def test_parallel_learners_roundtrip(tmp_path): p = tmp_path / "par.yaml" p.write_text(PAR_YAML) cfg = SpecConfig.from_yaml(p) - assert len(cfg.candidates) == 2 - assert cfg.candidates[0].label == "rf" - assert cfg.candidates[1].tasks["training"].function == "mymod:train_mlp" + assert len(cfg.learners) == 2 + assert cfg.learners[0].label == "rf" + assert cfg.learners[1].tasks["training"].function == "mymod:train_mlp" -def test_parallel_candidate_missing_slot(tmp_path): +def test_parallel_learner_missing_slot(tmp_path): yaml = textwrap.dedent("""\ learner: type: parallel_active_learner - candidates: + learners: - label: rf simulation: type: shell @@ -269,11 +269,297 @@ def test_parallel_candidate_missing_slot(tmp_path): SpecConfig.from_yaml(p) -def test_parallel_candidates_mixed_types(tmp_path): +# ── SpecConfig — parameters block ──────────────────────────────────────────── + +# ── SpecConfig — reserved parameter keys ────────────────────────────────────── + +# ── SpecConfig — task_description consistency across parallel learners ──────── + +def test_parallel_learners_identical_task_description_valid(tmp_path): + yaml = textwrap.dedent("""\ + learner: + type: parallel_active_learner + learners: + - label: a + simulation: + type: shell + command: python sim.py + task_description: + cpu_count: 4 + training: + type: python + function: mymod:train + active_learn: + type: python + function: mymod:select + - label: b + simulation: + type: shell + command: python sim.py --model b + task_description: + cpu_count: 4 + training: + type: python + function: mymod:train + active_learn: + type: python + function: mymod:select + stop_criterion: + metric: r2 + threshold: 0.9 + evaluator: + type: python + function: mymod:eval + """) + p = tmp_path / "ok.yaml" + p.write_text(yaml) + cfg = SpecConfig.from_yaml(p) + assert len(cfg.learners) == 2 + + +def test_parallel_learners_different_task_description_raises(tmp_path): + yaml = textwrap.dedent("""\ + learner: + type: parallel_active_learner + learners: + - label: a + simulation: + type: shell + command: python sim.py + task_description: + cpu_count: 4 + training: + type: python + function: mymod:train + active_learn: + type: python + function: mymod:select + - label: b + simulation: + type: shell + command: python sim.py + task_description: + cpu_count: 8 + training: + type: python + function: mymod:train + active_learn: + type: python + function: mymod:select + stop_criterion: + metric: r2 + threshold: 0.9 + evaluator: + type: python + function: mymod:eval + """) + p = tmp_path / "bad.yaml" + p.write_text(yaml) + with pytest.raises(ValueError, match="task_description"): + SpecConfig.from_yaml(p) + + +def test_parallel_learners_one_has_task_description_other_absent_raises(tmp_path): + yaml = textwrap.dedent("""\ + learner: + type: parallel_active_learner + learners: + - label: a + simulation: + type: shell + command: python sim.py + task_description: + cpu_count: 4 + training: + type: python + function: mymod:train + active_learn: + type: python + function: mymod:select + - label: b + simulation: + type: shell + command: python sim.py + training: + type: python + function: mymod:train + active_learn: + type: python + function: mymod:select + stop_criterion: + metric: r2 + threshold: 0.9 + evaluator: + type: python + function: mymod:eval + """) + p = tmp_path / "bad.yaml" + p.write_text(yaml) + with pytest.raises(ValueError, match="task_description"): + SpecConfig.from_yaml(p) + + +def test_parameters_reserved_key_pythonpath(tmp_path): + yaml = textwrap.dedent("""\ + learner: + type: sequential_active_learner + max_iter: 1 + simulation: + type: python + function: mymod:sim + training: + type: python + function: mymod:train + active_learn: + type: python + function: mymod:select + stop_criterion: + metric: mse + threshold: 0.1 + evaluator: + type: python + function: mymod:eval + parameters: + pythonpath: /some/path + """) + p = tmp_path / "bad.yaml" + p.write_text(yaml) + with pytest.raises(ValueError, match="reserved keys"): + SpecConfig.from_yaml(p) + + +def test_parameters_reserved_key_iteration(tmp_path): + yaml = textwrap.dedent("""\ + learner: + type: sequential_active_learner + max_iter: 1 + simulation: + type: python + function: mymod:sim + training: + type: python + function: mymod:train + active_learn: + type: python + function: mymod:select + stop_criterion: + metric: mse + threshold: 0.1 + evaluator: + type: python + function: mymod:eval + parameters: + iteration: 5 + """) + p = tmp_path / "bad.yaml" + p.write_text(yaml) + with pytest.raises(ValueError, match="reserved keys"): + SpecConfig.from_yaml(p) + + +def test_parameters_reserved_key_learner_id(tmp_path): + yaml = textwrap.dedent("""\ + learner: + type: sequential_active_learner + max_iter: 1 + simulation: + type: python + function: mymod:sim + training: + type: python + function: mymod:train + active_learn: + type: python + function: mymod:select + stop_criterion: + metric: mse + threshold: 0.1 + evaluator: + type: python + function: mymod:eval + parameters: + learner_id: 0 + """) + p = tmp_path / "bad.yaml" + p.write_text(yaml) + with pytest.raises(ValueError, match="reserved keys"): + SpecConfig.from_yaml(p) + + +def test_parameters_non_reserved_key_allowed(tmp_path): + yaml = textwrap.dedent("""\ + learner: + type: sequential_active_learner + max_iter: 1 + simulation: + type: python + function: mymod:sim + training: + type: python + function: mymod:train + active_learn: + type: python + function: mymod:select + stop_criterion: + metric: mse + threshold: 0.1 + evaluator: + type: python + function: mymod:eval + parameters: + dataset: my_ds + """) + p = tmp_path / "ok.yaml" + p.write_text(yaml) + cfg = SpecConfig.from_yaml(p) + assert cfg.parameters == {"dataset": "my_ds"} + + +def test_parameters_roundtrip(tmp_path): + yaml = textwrap.dedent("""\ + learner: + type: sequential_active_learner + max_iter: 5 + simulation: + type: python + function: mymod:sim + training: + type: python + function: mymod:train + active_learn: + type: python + function: mymod:select + stop_criterion: + metric: mse + threshold: 0.1 + evaluator: + type: python + function: mymod:eval + parameters: + dataset: my_ds + scale: 1.5 + growing_pool: false + """) + p = tmp_path / "params.yaml" + p.write_text(yaml) + cfg = SpecConfig.from_yaml(p) + assert cfg.parameters == {"dataset": "my_ds", "scale": 1.5, "growing_pool": False} + + +def test_parameters_defaults_empty(tmp_path): + p = tmp_path / "spec.yaml" + p.write_text(SEQ_YAML) + cfg = SpecConfig.from_yaml(p) + assert cfg.parameters == {} + + +# ── SpecConfig — parallel learners (mixed types) ───────────────────────────── + +def test_parallel_learners_mixed_types(tmp_path): yaml = textwrap.dedent("""\ learner: type: parallel_active_learner - candidates: + learners: - label: a simulation: type: shell diff --git a/tests/unit/spec/test_workflow_spec.py b/tests/unit/spec/test_workflow_spec.py new file mode 100644 index 00000000..76d9cafb --- /dev/null +++ b/tests/unit/spec/test_workflow_spec.py @@ -0,0 +1,247 @@ +"""Unit tests for WorkflowSpec — load_spec and workflow_with().""" +import textwrap + +import pytest + +from rose.spec import WorkflowSpec, load_spec +from rose.spec.schema import SpecConfig + +SEQ_YAML = textwrap.dedent("""\ + learner: + type: sequential_active_learner + max_iter: 5 + + simulation: + type: python + function: mymod:sim + + training: + type: python + function: mymod:train + + active_learn: + type: python + function: mymod:select + + stop_criterion: + metric: mse + threshold: 0.1 + evaluator: + type: python + function: mymod:eval + + parameters: + dataset: base_ds + scale: 1.0 +""") + + +def _make_spec(tmp_path, yaml=SEQ_YAML): + p = tmp_path / "spec.yaml" + p.write_text(yaml) + return load_spec(p) + + +# ── workflow_with: learner field override ───────────────────────────────────── + +def test_workflow_with_max_iter(tmp_path): + spec = _make_spec(tmp_path) + new = spec.workflow_with(max_iter=2) + assert new.config.learner.max_iter == 2 + assert spec.config.learner.max_iter == 5 # original unchanged + + +def test_workflow_with_does_not_mutate_original(tmp_path): + spec = _make_spec(tmp_path) + _ = spec.workflow_with(max_iter=1) + assert spec.config.learner.max_iter == 5 + + +# ── workflow_with: parameters merge ────────────────────────────────────────── + +def test_workflow_with_parameters_merges(tmp_path): + spec = _make_spec(tmp_path) + new = spec.workflow_with(parameters={"dataset": "test_ds"}) + assert new.config.parameters["dataset"] == "test_ds" + assert new.config.parameters["scale"] == 1.0 # existing key preserved + + +def test_workflow_with_parameters_adds_new_key(tmp_path): + spec = _make_spec(tmp_path) + new = spec.workflow_with(parameters={"lr": 0.001}) + assert new.config.parameters["lr"] == 0.001 + assert new.config.parameters["dataset"] == "base_ds" + + +def test_workflow_with_parameters_does_not_mutate_original(tmp_path): + spec = _make_spec(tmp_path) + _ = spec.workflow_with(parameters={"dataset": "other"}) + assert spec.config.parameters["dataset"] == "base_ds" + + +# ── workflow_with: combined overrides ───────────────────────────────────────── + +def test_workflow_with_combined(tmp_path): + spec = _make_spec(tmp_path) + new = spec.workflow_with(max_iter=3, parameters={"dataset": "test_ds", "scale": 2.0}) + assert new.config.learner.max_iter == 3 + assert new.config.parameters["dataset"] == "test_ds" + assert new.config.parameters["scale"] == 2.0 + + +# ── workflow_with: unknown field raises ─────────────────────────────────────── + +def test_workflow_with_unknown_field_raises(tmp_path): + spec = _make_spec(tmp_path) + with pytest.raises(ValueError, match="unknown spec field"): + spec.workflow_with(nonexistent_key=42) + + +# ── workflow_with: returns WorkflowSpec with callable .workflow ─────────────── + +def test_workflow_with_returns_workflow_spec(tmp_path): + spec = _make_spec(tmp_path) + new = spec.workflow_with(max_iter=1) + assert isinstance(new, WorkflowSpec) + assert callable(new.workflow) + + +# ── validate_imports ────────────────────────────────────────────────────────── + +IMPORTABLE_YAML = textwrap.dedent("""\ + learner: + type: sequential_active_learner + max_iter: 1 + + simulation: + type: python + function: os.path:join + + training: + type: python + function: os.path:dirname + + active_learn: + type: python + function: os.path:basename + + stop_criterion: + metric: mse + threshold: 0.1 + evaluator: + type: python + function: os.path:exists +""") + +BAD_MODULE_YAML = textwrap.dedent("""\ + learner: + type: sequential_active_learner + max_iter: 1 + + simulation: + type: python + function: no_such_module_xyz:fn + + training: + type: python + function: os.path:dirname + + active_learn: + type: python + function: os.path:basename + + stop_criterion: + metric: mse + threshold: 0.1 + evaluator: + type: python + function: os.path:exists +""") + +BAD_ATTR_YAML = textwrap.dedent("""\ + learner: + type: sequential_active_learner + max_iter: 1 + + simulation: + type: python + function: os.path:no_such_fn_xyz + + training: + type: python + function: os.path:dirname + + active_learn: + type: python + function: os.path:basename + + stop_criterion: + metric: mse + threshold: 0.1 + evaluator: + type: python + function: os.path:exists +""") + +MULTI_BAD_YAML = textwrap.dedent("""\ + learner: + type: sequential_active_learner + max_iter: 1 + + simulation: + type: python + function: no_such_module_xyz:fn + + training: + type: python + function: os.path:no_such_fn_xyz + + active_learn: + type: python + function: os.path:basename + + stop_criterion: + metric: mse + threshold: 0.1 + evaluator: + type: python + function: os.path:exists +""") + + +def test_validate_imports_succeeds_for_importable_functions(tmp_path): + p = tmp_path / "spec.yaml" + p.write_text(IMPORTABLE_YAML) + spec = load_spec(p, validate_imports=True) + assert spec is not None + + +def test_validate_imports_raises_on_bad_module(tmp_path): + p = tmp_path / "spec.yaml" + p.write_text(BAD_MODULE_YAML) + with pytest.raises(ValueError, match="no_such_module_xyz:fn"): + load_spec(p, validate_imports=True) + + +def test_validate_imports_raises_on_bad_attribute(tmp_path): + p = tmp_path / "spec.yaml" + p.write_text(BAD_ATTR_YAML) + with pytest.raises(ValueError, match="no_such_fn_xyz"): + load_spec(p, validate_imports=True) + + +def test_validate_imports_reports_all_failures(tmp_path): + p = tmp_path / "spec.yaml" + p.write_text(MULTI_BAD_YAML) + with pytest.raises(ValueError) as exc_info: + load_spec(p, validate_imports=True) + msg = str(exc_info.value) + assert "no_such_module_xyz:fn" in msg + assert "no_such_fn_xyz" in msg + + +def test_validate_imports_default_false_skips_check(tmp_path): + p = tmp_path / "spec.yaml" + p.write_text(BAD_MODULE_YAML) + spec = load_spec(p) # validate_imports=False by default — should not raise + assert spec is not None From 846d33f914f00d6f30e7609324b9f20d36f3c84c Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Thu, 18 Jun 2026 18:01:15 +0200 Subject: [PATCH 04/13] major docs refine --- docs/getting-started/faq.md | 3 + docs/getting-started/installation.md | 13 ++-- docs/index.md | 2 +- docs/integrations/clearml.md | 4 ++ docs/user-guide/acl-metrics.md | 7 +- docs/user-guide/advanced-acl-workflow.md | 47 ++++++++----- docs/user-guide/advanced-rl-workflow.md | 41 ++++-------- docs/user-guide/basic-acl-workflow.md | 6 +- docs/user-guide/basic-rl-workflow.md | 4 +- docs/user-guide/experience.md | 4 ++ docs/user-guide/parallel_learners_docs.md | 3 +- docs/user-guide/rose-aas.md | 2 +- docs/user-guide/spec-api.md | 81 ++++++++++++++++++----- docs/user-guide/target-resources.md | 12 ++-- docs/user-guide/uq_based-acl-workflow.md | 50 +++++++------- docs/user-guide/visualization.md | 72 ++++++++++++++++++++ mkdocs.yml | 31 +++++---- 17 files changed, 266 insertions(+), 116 deletions(-) diff --git a/docs/getting-started/faq.md b/docs/getting-started/faq.md index c0011c2e..3f2a63c5 100644 --- a/docs/getting-started/faq.md +++ b/docs/getting-started/faq.md @@ -9,6 +9,9 @@ Refer to the [Installation Guide](installation.md) for detailed instructions bas A dry run allows you to simulate the execution of your tasks locally without making any changes to the behavior of your ML workflow. It's a way to verify your setup is working on HPC before applying it to prevent time and resource waste. +!!! tip + See the [Dry Run](dry-run.md) guide for a full step-by-step walkthrough. + ## Q3: What if I have an issue? If you have any additional questions or an issue, feel free to reach out to our support team by opening a Github [ticket](https://github.com/radical-cybertools/ROSE/issues). diff --git a/docs/getting-started/installation.md b/docs/getting-started/installation.md index b3c2cad5..7000283f 100644 --- a/docs/getting-started/installation.md +++ b/docs/getting-started/installation.md @@ -7,7 +7,7 @@ machines, please refer to the following link: [RADICAL-Pilot Supported HPC Machi ## For Linux and macOS 🐧 1. Clone the latest version from the [official website](https://github.com/radical-cybertools/ROSE). - ``` + ```bash git clone https://github.com/radical-cybertools/ROSE.git ``` 2. Run the following commands to install ROSE and its dependencies: @@ -21,19 +21,19 @@ machines, please refer to the following link: [RADICAL-Pilot Supported HPC Machi 1. Download the Windows WSL installer from the [official website](https://apps.microsoft.com/detail/9pdxgncfsczv?hl=en-US&gl=US). 2. Setup you WSL user name and password. 3. Make sure you have Python 3.9 or higher in your WSL as follows: - ``` + ```bash python --version ``` 4. create new pip virtual env: - ``` + ```bash python3 -m venv rose_env ``` 5. Activate the env: - ``` + ```bash source rose_env/bin/activate ``` 6. Clone the latest version from the [official website](https://github.com/radical-cybertools/ROSE). - ``` + ```bash git clone https://github.com/radical-cybertools/ROSE.git ``` 7. Run the following commands to install ROSE and its dependencies: @@ -42,4 +42,5 @@ machines, please refer to the following link: [RADICAL-Pilot Supported HPC Machi pip install . ``` -If you encounter any issues, refer to the [Issues Section](https://github.com/radical-cybertools/ROSE/issues). +!!! note + If you encounter any issues, refer to the [Issues Section](https://github.com/radical-cybertools/ROSE/issues). diff --git a/docs/index.md b/docs/index.md index 76823901..4abd15d4 100644 --- a/docs/index.md +++ b/docs/index.md @@ -1,6 +1,6 @@ ## What is ROSE? 🌹 -ROSE: RADICAL Orchestrator for Surrogate Exploration (ROSE) toolkit is a framework designed to enable the concurrent and adaptive execution of simulation, surrogate training, and selection tasks on High-Performance Computing (HPC) resources. ROSE is a Python package that provides tools for developing active and reinforcement learning (AL and RL) methods for scientific applications. It enables users to define simulation and surrogate training tasks and automatically manage their execution on HPC resources via a predefined set of Learning Policies (Learners). +ROSE: RADICAL Orchestrator for Surrogate Exploration (ROSE) toolkit is a framework designed to enable the distribuited, concurrent and adaptive execution of simulation, surrogate training, and selection tasks on High-Performance Computing (HPC) resources at a scale. ROSE is a Python package that provides tools for developing active and reinforcement learning (AL and RL) methods for scientific applications. It enables users to define simulation and surrogate training tasks and automatically manage their execution on HPC resources via a predefined set of Learning Policies (Learners). ROSE also includes tools to facilitate the selection of the most effective surrogate model for a given simulation based on performance metrics. diff --git a/docs/integrations/clearml.md b/docs/integrations/clearml.md index d5b1ee79..0eb2a461 100644 --- a/docs/integrations/clearml.md +++ b/docs/integrations/clearml.md @@ -61,6 +61,10 @@ For **parallel learners**, `state.learner_id` is included automatically. The tra each state as a separate `series` inside the same scalar title, making per-learner curves directly comparable without any user code. +!!! note + No user code is required to get this overlay — `state.learner_id` is already populated + by the parallel learner framework before the tracker ever sees the state. + ### Task tags — logged in `on_stop` | ClearML tag | Value | diff --git a/docs/user-guide/acl-metrics.md b/docs/user-guide/acl-metrics.md index d8d71dc9..4a303f7d 100644 --- a/docs/user-guide/acl-metrics.md +++ b/docs/user-guide/acl-metrics.md @@ -4,7 +4,7 @@ ROSE supports different Machine Learning (ML) Metrics such as `RMSE`, `MAE`, and ## Standard Metrics -For a full list of the supported metrics please refer to the following link [ROSE Standard Metrics](https://github.com/radical-cybertools/ROSE/blob/feature/al_algo_selector/rose/metrics.py) +For a full list of the supported metrics please refer to the following link [ROSE Standard Metrics](https://github.com/radical-cybertools/ROSE/blob/main/rose/metrics.py) ## Custom Metrics @@ -26,3 +26,8 @@ async def check_metric(*args): ``` In this way, ROSE will understand the relation between the custom metric and the target threshold value. + +!!! note + `metric_name` is just a label used for logging and tracking — ROSE does not validate it + against a fixed list. Any string is accepted as long as the decorated function returns a + numerical value to compare against `threshold`. diff --git a/docs/user-guide/advanced-acl-workflow.md b/docs/user-guide/advanced-acl-workflow.md index 6c501bcc..64412395 100644 --- a/docs/user-guide/advanced-acl-workflow.md +++ b/docs/user-guide/advanced-acl-workflow.md @@ -1,26 +1,39 @@ +# Advanced Active Learning Workflow + To support the rapid advancement of AL techniques, ROSE offers an additional approach to building and executing complex AL workflows. In this example, we demonstrate how to express an AL workflow with different levels of parallelism. What does that mean? In some cases, AL workflows may require the execution of N simulation or training tasks **concurrently**. But not only that—additionally, they may also require the submission of M AL workflows concurrently. This introduces two levels of parallelism: one at the task level and another at the AL workflow level. Such an approach is possible and can be easily expressed and executed using ROSE's **custom AL policy**. -```sh - (N AL WFs in Parallel) - +-------------------+ +-------------------+ - | AL WF 1 | | AL WF 2 | - +-------------------+ +-------------------+ - │ │ - +----------------+-----------------+ +----------------+-----------------+ - | (N tasks Parallel) | | (N AL tasks Parallel) | - +---------------+ +---------------+ +---------------+ +---------------+ - | Simulation 1 | | Simulation 2 | | Simulation 1 | | Simulation 2 | - +---------------+ +---------------+ +---------------+ +---------------+ - | | | | - +---------------+ +---------------+ +---------------+ +---------------+ - | Training 1 | | Training 2 | | Training 1 | | Training 2 | - +---------------+ +---------------+ +---------------+ +---------------+ - | | | | - (...) (...) (...) (...) +!!! note + Task-level parallelism (N simulation/training tasks per iteration) and workflow-level + parallelism (M AL loops running side by side) are independent dimensions — you can scale + either one without changing the other. + +```mermaid +graph TD + N["N AL WFs in Parallel"] + N --> WF1["AL WF 1"] + N --> WF2["AL WF 2"] + + subgraph G1[" "] + WF1 --> S1a["Simulation 1"] + WF1 --> S1b["Simulation 2"] + S1a --> T1a["Training 1"] + S1b --> T1b["Training 2"] + T1a --> E1a["..."] + T1b --> E1b["..."] + end + + subgraph G2[" "] + WF2 --> S2a["Simulation 1"] + WF2 --> S2b["Simulation 2"] + S2a --> T2a["Training 1"] + S2b --> T2b["Training 2"] + T2a --> E2a["..."] + T2b --> E2b["..."] + end ``` Since we have already learned how to deploy and load ROSE, and how to instruct it to use different resources, we will skip this part and focus only on expressing the AL workflow. diff --git a/docs/user-guide/advanced-rl-workflow.md b/docs/user-guide/advanced-rl-workflow.md index 25323987..9091a418 100644 --- a/docs/user-guide/advanced-rl-workflow.md +++ b/docs/user-guide/advanced-rl-workflow.md @@ -1,33 +1,21 @@ +# Advanced Reinforcement Learning Workflow + In addition to basic reinforcement learning (RL) workflows, ROSE supports advanced RL workflows that can run multiple environment instances in parallel. The 'ParallelLearner' gives you the ability to run multiple environment tasks simultaneously, each with different parameters, and then merge their experiences for training. This is particularly useful for scenarios where you want to explore different configurations or hyperparameters in parallel, speeding up the learning process. -```sh - - +-------------------+ - | RL WF | - +-------------------+ - │ - +-------------------------+---------------------------+ - | (N Environment Tasks Parallel) | - +---------------+ +---------------+ +---------------+ - | Environment 1 | | Environment 2 | | Environment 3 | - +---------------+ +---------------+ +---------------+ - | | | - └────────────────┼────────────────────┘ - │ - +------v------+ - | Merge | - +------+------+ - │ - +------v------+ - | Update | - +------+------+ - │ - +------v------+ - | Test | - +-------------+ +```mermaid +graph TD + WF["RL WF"] + WF --> E1["Environment 1"] + WF --> E2["Environment 2"] + WF --> E3["Environment 3"] + E1 --> M["Merge"] + E2 --> M + E3 --> M + M --> U["Update"] + U --> T["Test"] ``` Import ROSE parallel RL modules: @@ -80,8 +68,7 @@ async def main(): Now that each environment task is defined, we define the rest of the workflow components: !!! note - -This snippet of code must be inside an async context or inside `main` function + This snippet of code must be inside an async context or inside `main` function ```python @pe.update_task diff --git a/docs/user-guide/basic-acl-workflow.md b/docs/user-guide/basic-acl-workflow.md index 4ff8d701..485f1626 100644 --- a/docs/user-guide/basic-acl-workflow.md +++ b/docs/user-guide/basic-acl-workflow.md @@ -38,8 +38,8 @@ async def active_learn(*args): ``` !!! tip -ROSE supports defining tasks with python code instead of executables (i.e., python scripts, shell scripts, etc.). To do that, the user have to -pass the `as_executable=False` argument to the decorator as follows: + ROSE supports defining tasks with python code instead of executables (i.e., python scripts, shell scripts, etc.). To do that, the user have to + pass the `as_executable=False` argument to the decorator as follows: ```python @acl.simulation_task(as_executable=False) @@ -92,7 +92,7 @@ async def check_mse(*args): return f'python3 check_mse.py' ``` -!!! Warning +!!! warning For any metric function like `@acl.as_stop_criterion` the invoked script like `check_mse.py` must return a numerical value. diff --git a/docs/user-guide/basic-rl-workflow.md b/docs/user-guide/basic-rl-workflow.md index 1eeb2711..a40152c7 100644 --- a/docs/user-guide/basic-rl-workflow.md +++ b/docs/user-guide/basic-rl-workflow.md @@ -1,3 +1,5 @@ +# Basic Reinforcement Learning Workflow + ## Define your target machine to run on Import ROSE main modules: @@ -33,7 +35,7 @@ async def check_reward(*args): return 'python3 check_reward.py' ``` -!!! Warning +!!! warning For any metric function like `@rl.as_stop_criterion` the invoked script like `check_reward.py` must return a numerical value. Finally invoke the tasks and register them with the reinforcement learner as a workflow: diff --git a/docs/user-guide/experience.md b/docs/user-guide/experience.md index 59155772..696e7d02 100644 --- a/docs/user-guide/experience.md +++ b/docs/user-guide/experience.md @@ -39,6 +39,10 @@ batch = bank.sample(batch_size=32, replace=True) ROSE assigns experience banks a unique session ID automatically, or session IDs can be assigned manually. The session ID is used to identify the bank and can be used to save/load the bank to/from disk: +!!! note + Automatic session IDs require no bookkeeping on your part — pass `session_id` explicitly + only when you need a stable, predictable filename across runs (e.g. resuming a bank by name). + ```python # Custom session ID bank = ExperienceBank(session_id="rose_session") diff --git a/docs/user-guide/parallel_learners_docs.md b/docs/user-guide/parallel_learners_docs.md index b5cfe5f0..2b458caf 100644 --- a/docs/user-guide/parallel_learners_docs.md +++ b/docs/user-guide/parallel_learners_docs.md @@ -10,8 +10,7 @@ This tutorial demonstrates how to configure and run multiple learning pipelines --- !!! note - -This approach can be applied for both Active and Reinforcement learners (Sequential and Parallel). + This approach can be applied for both Active and Reinforcement learners (Sequential and Parallel). ## Example Overview diff --git a/docs/user-guide/rose-aas.md b/docs/user-guide/rose-aas.md index 009c1c55..96f2fb9b 100644 --- a/docs/user-guide/rose-aas.md +++ b/docs/user-guide/rose-aas.md @@ -6,7 +6,7 @@ This page describes ROSE's service model: how the active-learning/RL loop you de ## Core idea: BYOF — Bring Your Own Workflow -ROSE does not own the science. It owns the **loop**: submit a task, wait for it, decide what runs next, check a stop criterion, repeat — until `max_iter` or the criterion is met. +ROSE does not own the science. It owns the **orchestration loop**: submit a task, wait for it, decide what runs next, check a stop criterion, repeat — until `max_iter` or the criterion is met. You bring: diff --git a/docs/user-guide/spec-api.md b/docs/user-guide/spec-api.md index 7a98d766..bc481137 100644 --- a/docs/user-guide/spec-api.md +++ b/docs/user-guide/spec-api.md @@ -4,6 +4,15 @@ The YAML spec API lets you declare a ROSE workflow as a data file instead of Pyt --- + +!!! note + The YAML spec currently supports only the four built-in `learner.type` values listed + above. Custom `Learner` subclasses are not yet expressible in YAML — `LearnerBuilder` + raises `ValueError` for any other `learner.type` value. Use the Python API + (decorator-based, e.g. `SequentialActiveLearner(asyncflow)`) if you need a custom + learner implementation. + + ## Loading a Spec ```python @@ -17,6 +26,28 @@ cfg = spec.config # typed SpecConfig object --- +## Resource Specification via `task_description` + +Every task — `simulation`, `training`, `active_learn`, `environment`, `update`, … — can +carry an optional `task_description`. In the Python API this is a parameter on the +decorated function itself, with a default value: + +```python +@acl.simulation_task(as_executable=False) +async def simulate(*args, task_description={"process_templates": [(4, {})]}, **kwargs): + ... +``` + +The accepted keys depend entirely on which execution backend your `asyncflow` session is +using — see asyncflow's +[Execution Backends guide](https://radical-cybertools.github.io/radical.asyncflow/exec_backends/?h=task_desc#assign-resources-for-your-application-task) +for the full per-backend reference. + +The YAML spec exposes the same dict as a `task_description:` key on the corresponding task +slot — see the side-by-side example in [Task Types](#task-types) below. + +--- + ## Sequential Active Learner
@@ -137,28 +168,28 @@ When all parallel learners run the same task implementations, declare the task s ```python from rose.al.active_learner import ParallelActiveLearner -pal = ParallelActiveLearner(asyncflow) +acl = ParallelActiveLearner(asyncflow) -@pal.simulation_task(as_executable=False) +@acl.simulation_task(as_executable=False) async def simulate(*args, **kwargs): ... -@pal.training_task(as_executable=False) +@acl.training_task(as_executable=False) async def train(sim_result, **kwargs): ... -@pal.active_learn_task(as_executable=False) +@acl.active_learn_task(as_executable=False) async def active_learn(sim_result, model, **kwargs): ... -@pal.as_stop_criterion( +@acl.as_stop_criterion( metric_name="mse", threshold=0.01, ) async def check_mse(*args, **kwargs): ... -async for state in pal.start( +async for state in acl.start( parallel_learners=2, max_iter=5, ): @@ -229,7 +260,7 @@ When each parallel learner needs its own task implementations — different mode **Python API** ```python -pal = ParallelActiveLearner(asyncflow) +acl = ParallelActiveLearner(asyncflow) # Manual routing by learner_id TRAIN_CMD = { @@ -237,19 +268,19 @@ TRAIN_CMD = { 1: "python train.py --model ridge", } -@pal.simulation_task +@acl.simulation_task async def simulate(*args, **kwargs): return "python sim.py" -@pal.training_task +@acl.training_task async def train(*args, **kwargs): return TRAIN_CMD[kwargs["learner_id"]] -@pal.active_learn_task +@acl.active_learn_task async def active_learn(*args, **kwargs): return "python active_learn.py" -@pal.as_stop_criterion( +@acl.as_stop_criterion( metric_name="mse", threshold=0.01 ) async def check_mse(*args, **kwargs): @@ -266,7 +297,7 @@ lcs = [ for lid in range(2) ] -async for state in pal.start( +async for state in acl.start( parallel_learners=2, max_iter=5, learner_configs=lcs, @@ -395,19 +426,36 @@ The criterion evaluator must print a single float to stdout — ROSE captures st
-### Resource hints (`task_description`) +### `task_description` example + +
+
+ +**Python API** -Both task types accept an optional `task_description` dict forwarded to the execution backend as resource hints (CPU count, GPU count, etc.): +```python +@acl.simulation_task(as_executable=False) +async def simulate(*args, task_description={"process_templates": [(4, {})]}, **kwargs): + ... +``` + +
+
+ +**YAML Spec** ```yaml simulation: type: python function: tasks:simulate task_description: - cpu_count: 4 - gpu_count: 1 + process_templates: + - [4, {}] ``` +
+
+ !!! note For parallel learners using the `learners:` block, all entries must use the same `task_description` for each slot — the backend registers it once at task-registration time. @@ -575,6 +623,7 @@ With `validate_imports=True`, every `module:callable` string is resolved in the | `tracking.experiment` | string | no | `ROSE-Spec` | Experiment name | | `tracking.run_name` | string | no | `null` | Run label | + ### TaskDef fields | Key | Type | Required | Description | diff --git a/docs/user-guide/target-resources.md b/docs/user-guide/target-resources.md index eb6add8c..086dd196 100644 --- a/docs/user-guide/target-resources.md +++ b/docs/user-guide/target-resources.md @@ -6,14 +6,14 @@ For local execution, user can use their desktops, laptops, and their own small c ```python import os +from concurrent.futures import ProcessPoolExecutor + from radical.asyncflow import WorkflowEngine -from rhapsody.backends import RadicalExecutionBackend +from rhapsody.backends import ConcurrentExecutionBackend from rose.al.active_learner import SequentialActiveLearner -engine = await RadicalExecutionBackend( - {'runtime': 30, - 'resource': 'local.localhost'}) +engine = await ConcurrentExecutionBackend(ProcessPoolExecutor()) asyncflow = await WorkflowEngine.create(engine) @@ -21,6 +21,10 @@ acl = SequentialActiveLearner(asyncflow) ``` ## HPC Resources + +!!! warning + `RadicalExecutionBackend` will be deprecated by the end of Q3 2026. + To execute AL workflows on HPC machines, users must have an active allocation on the target machine and specify their resource requirements, as well as the time needed to execute their workflows. Remember, ROSE uses `RadicalExecutionBackend` from [RHAPSODY](https://github.com/radical-cybertools/rhapsody) (`rhapsody-py`) which is an interface for RADICAL-Pilot runtime system. For more information on how to access, set up, and execute workflows on HPC machines, refer to the following link [RADICAL-Pilot Job Submission](https://radicalpilot.readthedocs.io/en/stable/tutorials/submission.html): ```python diff --git a/docs/user-guide/uq_based-acl-workflow.md b/docs/user-guide/uq_based-acl-workflow.md index 24cfc90f..449bbce4 100644 --- a/docs/user-guide/uq_based-acl-workflow.md +++ b/docs/user-guide/uq_based-acl-workflow.md @@ -1,3 +1,5 @@ +# UQ-Based Active Learning Workflow + To accelerate the development of UQ-driven active learning methods, ROSE provides a flexible approach for composing and executing complex AL workflows. In this example, we illustrate how to define a UQ–AL workflow that supports multiple levels of parallelism. @@ -10,25 +12,23 @@ This introduces **two levels of parallelism**: * **Workflow-level parallelism** for executing multiple UQ–AL loops side by side. Both levels can be naturally expressed and efficiently executed using ROSE’s **custom AL policy**, enabling scalable and adaptive uncertainty-aware learning. -```sh - (N AL WFs in Parallel) - +-------------------+ +-------------------+ - | UQ WF 1 | | UQ WF 2 | - +-------------------+ +-------------------+ - │ │ - +----------------+-----------------+ +----------------+-----------------+ - | (N tasks Parallel) | | (N AL tasks Parallel) | - +---------------+ +---------------+ +---------------+ +---------------+ - | Simulation 1,2,..............n | | Simulation 1,2,..............n | - +---------------+ +---------------+ +---------------+ +---------------+ - | | | | - +---------------+ +---------------+ +---------------+ +---------------+ - | Train Model 1,2,...........m | | Train Model 1,2,...........m | - +---------------+ +---------------+ +---------------+ +---------------+ - | | | | - +-----------------------------+ +-----------------------------+ - | AL based on UQ WF 1 | | AL based on UQ WF 1 | - +-----------------------------+ +-----------------------------+ +```mermaid +graph TD + N["N AL WFs in Parallel"] + N --> WF1["UQ WF 1"] + N --> WF2["UQ WF 2"] + + subgraph G1[" "] + WF1 --> S1["Simulation 1, 2, ..., n (parallel)"] + S1 --> T1["Train Model 1, 2, ..., m (parallel)"] + T1 --> A1["AL based on UQ WF 1"] + end + + subgraph G2[" "] + WF2 --> S2["Simulation 1, 2, ..., n (parallel)"] + S2 --> T2["Train Model 1, 2, ..., m (parallel)"] + T2 --> A2["AL based on UQ WF 2"] + end ``` @@ -56,7 +56,7 @@ To support this approach, two new task types have been introduced: code_path = f'{sys.executable} {os.getcwd()}' # Define and register the prediction task -@learner.prediction_task() +@uql.prediction_task() async def prediction(*args): return f'{code_path}/predict.py' @@ -71,11 +71,11 @@ async def prediction(*args): ```python # Defining the uncertainty quantification with a metric (PREDICTIVE_ENTROPY in this case) -@learner.uncertainty_quantification(uq_metric_name=PREDICTIVE_ENTROPY, +@uql.uncertainty_quantification(uq_metric_name=PREDICTIVE_ENTROPY, threshold=1.0, query_size=10) async def check_uq(*args): - return f'{code_path}/check_uq.py'xs + return f'{code_path}/check_uq.py' ``` @@ -95,7 +95,7 @@ Import and Initialize the UQ Learner ```python from rose.uq.uq_active_learner import ParallelUQLearner -learner = ParallelUQLearner(asyncflow) +uql = ParallelUQLearner(asyncflow) ``` Run the learning (Active Learning Loop) @@ -103,7 +103,7 @@ Run the learning (Active Learning Loop) ```python PIPELINES = ['UQ_learner1', 'UQ_learner2'] -results = await learner.start( +results = await uql.start( learner_names=PIPELINES, model_names=MODELS, learner_configs=learner_configs, @@ -122,7 +122,7 @@ print(results) with open(Path(os.getcwd(), 'UQ_training_results.json'), 'w') as f: json.dump(results, f, indent=4) -await learner.shutdown() +await uql.shutdown() ``` #### Defining costom UQ metric diff --git a/docs/user-guide/visualization.md b/docs/user-guide/visualization.md index e69de29b..881189e6 100644 --- a/docs/user-guide/visualization.md +++ b/docs/user-guide/visualization.md @@ -0,0 +1,72 @@ +# Visualization + +ROSE does not ship a plotting module — visualizing a run's progress is left to whichever +tool already fits your workflow: raw `IterationState` fields, the native file tracker, or +an MLflow/ClearML dashboard if `tracking.backend` is already wired up. + +!!! note + There is no `rose.plot` or similar utility. The three approaches below all read data + ROSE already produces — none requires changes to your task code. + +--- + +## Real-time, in-loop + +Every iteration yields an `IterationState` with `iteration`, `metric_value`, and +`metric_history` (the full list of past metric values so far): + +```python +import matplotlib.pyplot as plt + +async for state in learner.start(max_iter=20): + print(f"[iter {state.iteration}] mse={state.metric_value:.4f}") + +plt.plot(state.metric_history) +plt.xlabel("iteration") +plt.ylabel("mse") +plt.show() +``` + +This is the fastest path when you just want a quick look right after a local run. + +--- + +## Native file tracker + +For HPC runs where you want a durable, preemption-safe record, attach `HPC_FileTracker` +(see `examples/integrations/tracking/basic.py`) — it appends one JSON line per iteration: + +```python +learner.add_tracker(HPC_FileTracker("run.jsonl")) +``` + +Replay and plot after the run, or from a different machine entirely: + +```python +import pandas as pd + +df = pd.read_json("run.jsonl", lines=True) +iterations = df[df.event == "iteration"] +iterations.plot(x="iteration", y="mse", logy=True) +``` + +Any numeric value your tasks return as a `dict` (e.g. `n_labeled`, `train_mse`) is captured +automatically and available as its own column. + +--- + +## MLflow / ClearML dashboards + +If `tracking.backend: mlflow` or `tracking.backend: clearml` is set in your spec (or you +attach `MLflowTracker`/`ClearMLTracker` directly in the Python API), every iteration's +metrics are already logged with no extra code. The respective web UIs are the recommended +way to compare runs and overlay parallel-learner series: + +- [MLflow integration](../integrations/mlflow.md) — `mlflow ui --port 5000`, metric curves + per run, run comparison. +- [ClearML integration](../integrations/clearml.md) — **Scalars** tab, parallel learners + shown as separate series under the same title. + +!!! tip + Prefer the dashboards over building your own plots once you have more than a couple of + runs to compare — both tools already handle multi-run overlay and filtering. diff --git a/mkdocs.yml b/mkdocs.yml index f4873f74..5c82ca3e 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -66,7 +66,11 @@ markdown_extensions: - admonition - codehilite: linenums: true - - pymdownx.superfences + - pymdownx.superfences: + custom_fences: + - name: mermaid + class: mermaid + format: !!python/name:pymdownx.superfences.fence_code_format - pymdownx.details - pymdownx.tasklist: custom_checkbox: true @@ -88,17 +92,20 @@ nav: - 3.FAQ: getting-started/faq.md - User Guide: - 1.Target Resources: user-guide/target-resources.md - - 2.Basic AL workflow: user-guide/basic-acl-workflow.md - - 3.AL Metric: user-guide/acl-metrics.md - - 4.Building Parallel Learners: user-guide/parallel_learners_docs.md - - 5.Advanced AL workflow: user-guide/advanced-acl-workflow.md - - 6.Visualization: user-guide/visualization.md - - 7.Reinforcement Learning: user-guide/basic-rl-workflow.md - - 8.Experience Banks: user-guide/experience.md - - 9.Advanced RL workflow: user-guide/advanced-rl-workflow.md - - 10.Tracking & Observability: user-guide/tracking.md - - 11.YAML Spec API: user-guide/spec-api.md - - 12.ROSE as a Service: user-guide/rose-aas.md + - 2.Active Learners: + - 1.Sequential Active Learning Example: user-guide/basic-acl-workflow.md + - 2.Active Learning Metrics: user-guide/acl-metrics.md + - 3.Parallel Active Learning Example: user-guide/parallel_learners_docs.md + - 4.Custom Active Learning Example: user-guide/advanced-acl-workflow.md + - 5.UQ-Based Active Learning Example: user-guide/uq_based-acl-workflow.md + - 3.Reinforcement Learners: + - 1.Sequential Reinforcement Learning Example: user-guide/basic-rl-workflow.md + - 2.Experience Banks: user-guide/experience.md + - 3.Parallel Reinforcement Learning Example: user-guide/advanced-rl-workflow.md + - 4.Visualization: user-guide/visualization.md + - 5.Tracking & Observability: user-guide/tracking.md + - 6.YAML Spec API: user-guide/spec-api.md + - 7.ROSE as a Service: user-guide/rose-aas.md - Integrations: - MLflow: integrations/mlflow.md - ClearML: integrations/clearml.md From fae7af7f5ca2acd84412aecf22842137aa7eb238 Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Thu, 18 Jun 2026 22:20:14 +0200 Subject: [PATCH 05/13] fix Spec API and docs inconsistency --- docs/user-guide/spec-api.md | 30 ++++++++---- rose/spec/__init__.py | 14 +++--- rose/spec/builder.py | 16 +++++-- rose/spec/schema.py | 8 ++-- tests/integration/spec/test_yaml_workflow.py | 8 ++-- tests/unit/spec/test_builder.py | 4 +- tests/unit/spec/test_schema.py | 48 ++++++++++---------- tests/unit/spec/test_workflow_spec.py | 2 +- 8 files changed, 73 insertions(+), 57 deletions(-) diff --git a/docs/user-guide/spec-api.md b/docs/user-guide/spec-api.md index bc481137..564923db 100644 --- a/docs/user-guide/spec-api.md +++ b/docs/user-guide/spec-api.md @@ -15,14 +15,25 @@ The YAML spec API lets you declare a ROSE workflow as a data file instead of Pyt ## Loading a Spec +For the common case — build a learner and run it — `LearnerBuilder` loads and validates the YAML itself; you never need to touch a config object: + +```python +from rose.spec.builder import LearnerBuilder + +builder = LearnerBuilder("workflow.yaml", asyncflow) # validates schema on load +cfg = builder.config # typed WorkflowConfig, if you need it +``` + +Reach for `load_spec` instead when you need `WorkflowSpec`-level features: spec variants via [`workflow_with()`](#spec-variants-with-workflow_with), import validation, or the `.workflow` coroutine used by [ROSE as a Service](rose-aas.md): + ```python from rose.spec import load_spec spec = load_spec("workflow.yaml") # validates schema on load -cfg = spec.config # typed SpecConfig object +cfg = spec.config # typed WorkflowConfig object ``` -`load_spec` raises `ValueError` with a precise message on any schema violation. It does **not** import task modules by default — see [Import Validation](#import-validation) to enable that check. +Both raise `ValueError` with a precise message on any schema violation. Neither imports task modules by default — see [Import Validation](#import-validation) to enable that check. --- @@ -120,11 +131,10 @@ stop_criterion: ```python # run it -from rose.spec import load_spec from rose.spec.builder import LearnerBuilder -cfg = load_spec("workflow.yaml").config -builder = LearnerBuilder(cfg, asyncflow) +builder = LearnerBuilder("workflow.yaml", asyncflow) +cfg = builder.config learner = builder.build() async for state in learner.start( @@ -231,8 +241,8 @@ stop_criterion: ```python # run it -cfg = load_spec("workflow.yaml").config -builder = LearnerBuilder(cfg, asyncflow) +builder = LearnerBuilder("workflow.yaml", asyncflow) +cfg = builder.config learner = builder.build() async for state in learner.start( @@ -349,8 +359,8 @@ stop_criterion: ```python # run it — builder handles routing + LearnerConfig -cfg = load_spec("workflow.yaml").config -builder = LearnerBuilder(cfg, asyncflow) +builder = LearnerBuilder("workflow.yaml", asyncflow) +cfg = builder.config learner = builder.build() lcs = builder.build_learner_configs() @@ -403,7 +413,7 @@ The function may be sync or async. The return value is passed as the first posit **`type: shell`** -Runs a subprocess on the worker. Data flows through files on disk. The command string is the return value — ROSE submits it as an executable. +Runs in a single or multi-process on the worker. Data flows through files on disk. The command string is the return value — ROSE submits it as an executable. ```yaml simulation: diff --git a/rose/spec/__init__.py b/rose/spec/__init__.py index de2b44f5..8ee9acbd 100644 --- a/rose/spec/__init__.py +++ b/rose/spec/__init__.py @@ -4,7 +4,7 @@ from typing import Any, Callable, Coroutine from .builder import LearnerBuilder -from .schema import SpecConfig +from .schema import WorkflowConfig __all__ = ["load_spec", "WorkflowSpec"] @@ -20,13 +20,13 @@ def load_spec(path: str | Path, validate_imports: bool = False) -> "WorkflowSpec locally accessible. Leave False (default) when ``remote.pythonpath`` points to paths that only exist on the remote worker. """ - spec = WorkflowSpec(SpecConfig.from_yaml(path)) + spec = WorkflowSpec(WorkflowConfig.from_yaml(path)) if validate_imports: _validate_task_imports(spec.config) return spec -def _collect_python_specs(cfg: "SpecConfig") -> list[str]: +def _collect_python_specs(cfg: "WorkflowConfig") -> list[str]: """Return deduplicated 'module:callable' strings for all Python tasks in the spec.""" specs: list[str] = [] if cfg.learners is not None: @@ -43,7 +43,7 @@ def _collect_python_specs(cfg: "SpecConfig") -> list[str]: return list(dict.fromkeys(specs)) # deduplicate, preserve order -def _validate_task_imports(cfg: "SpecConfig") -> None: +def _validate_task_imports(cfg: "WorkflowConfig") -> None: """Try to import every Python task callable. Raises ValueError listing all failures.""" import importlib import sys as _sys @@ -75,7 +75,7 @@ def _validate_task_imports(cfg: "SpecConfig") -> None: class WorkflowSpec: """Validated workflow spec that produces a coroutine compatible with service_utils.run().""" - def __init__(self, config: SpecConfig) -> None: + def __init__(self, config: WorkflowConfig) -> None: self.config = config def workflow_with(self, **overrides: Any) -> "WorkflowSpec": @@ -84,7 +84,7 @@ def workflow_with(self, **overrides: Any) -> "WorkflowSpec": Accepted keys: - ``parameters``: dict merged (not replaced) into the existing parameters block - Any ``LearnerSpec`` field: ``max_iter``, ``parallel_learners`` - - Any other top-level ``SpecConfig`` field + - Any other top-level ``WorkflowConfig`` field Example:: @@ -101,7 +101,7 @@ def workflow_with(self, **overrides: Any) -> "WorkflowSpec": data[key] = value else: raise ValueError(f"workflow_with: unknown spec field '{key}'") - return WorkflowSpec(SpecConfig.model_validate(data)) + return WorkflowSpec(WorkflowConfig.model_validate(data)) @property def workflow(self) -> Callable[..., Coroutine[Any, Any, None]]: diff --git a/rose/spec/builder.py b/rose/spec/builder.py index f838a3f3..89dc217d 100644 --- a/rose/spec/builder.py +++ b/rose/spec/builder.py @@ -1,9 +1,11 @@ from __future__ import annotations +from pathlib import Path + from rose.learner import Learner from .adapters import TaskAdapterFactory -from .schema import SpecConfig, TrackingConfig, _REQUIRED_SLOTS +from .schema import WorkflowConfig, TrackingConfig, _REQUIRED_SLOTS # Slots that exist as fields on LearnerConfig; uncertainty is required by uq_active_learner # but is not a LearnerConfig field — filter it out when constructing LearnerConfig. @@ -33,8 +35,12 @@ def _get_learner_class(learner_type: str): class LearnerBuilder: - def __init__(self, config: SpecConfig, asyncflow) -> None: - self.config = config + def __init__(self, source: str | Path | WorkflowConfig, asyncflow) -> None: + if isinstance(source, (str, Path)): + source = WorkflowConfig.from_yaml(source) + elif hasattr(source, "config"): # duck-typed: also accepts a WorkflowSpec + source = source.config + self.config = source self.asyncflow = asyncflow def build(self) -> Learner: @@ -59,13 +65,13 @@ def build(self) -> Learner: _attach_tracker(learner, cfg.tracking) return learner - def _register_flat_tasks(self, learner: Learner, cfg: SpecConfig) -> None: + def _register_flat_tasks(self, learner: Learner, cfg: WorkflowConfig) -> None: for slot_name, task_def in cfg.tasks.items(): closure = TaskAdapterFactory.make_closure(task_def, cfg.remote) as_exec = TaskAdapterFactory.as_executable(task_def) getattr(learner, f"{slot_name}_task")(as_executable=as_exec)(closure) - def _register_dispatched_tasks(self, learner: Learner, cfg: SpecConfig) -> None: + def _register_dispatched_tasks(self, learner: Learner, cfg: WorkflowConfig) -> None: required = _REQUIRED_SLOTS[cfg.learner.type] for slot_name in required: task_defs = [l.tasks[slot_name] for l in cfg.learners] diff --git a/rose/spec/schema.py b/rose/spec/schema.py index 57d93d46..17cef4c9 100644 --- a/rose/spec/schema.py +++ b/rose/spec/schema.py @@ -98,7 +98,7 @@ class TrackingConfig(BaseModel): # ── Top-level spec ──────────────────────────────────────────────────────────── -class SpecConfig(BaseModel): +class WorkflowConfig(BaseModel): learner: LearnerSpec # Task slots — explicit fields keep extra="forbid" and enable IDE autocomplete. # Access non-None slots as a dict via the .tasks property. @@ -122,7 +122,7 @@ def tasks(self) -> dict[str, TaskDef]: return {s: getattr(self, s) for s in _ALL_SLOTS if getattr(self, s) is not None} @model_validator(mode="after") - def _validate_task_slots(self) -> "SpecConfig": + def _validate_task_slots(self) -> "WorkflowConfig": ltype = self.learner.type required = _REQUIRED_SLOTS.get(ltype) if required is None: @@ -166,7 +166,7 @@ def _validate_task_slots(self) -> "SpecConfig": return self @model_validator(mode="after") - def _validate_parameters(self) -> "SpecConfig": + def _validate_parameters(self) -> "WorkflowConfig": conflicts = _RESERVED_PARAMETER_KEYS & set(self.parameters.keys()) if conflicts: raise ValueError( @@ -176,7 +176,7 @@ def _validate_parameters(self) -> "SpecConfig": return self @classmethod - def from_yaml(cls, path: str | Path) -> "SpecConfig": + def from_yaml(cls, path: str | Path) -> "WorkflowConfig": import yaml raw = yaml.safe_load(Path(path).read_text()) diff --git a/tests/integration/spec/test_yaml_workflow.py b/tests/integration/spec/test_yaml_workflow.py index 3c60ccc8..d430ba11 100644 --- a/tests/integration/spec/test_yaml_workflow.py +++ b/tests/integration/spec/test_yaml_workflow.py @@ -8,7 +8,7 @@ from rose.spec import load_spec from rose.spec.builder import LearnerBuilder -from rose.spec.schema import SpecConfig +from rose.spec.schema import WorkflowConfig # ── Sequential AL via YAML ──────────────────────────────────────────────────── @@ -49,7 +49,7 @@ async def test_yaml_sequential_workflow(tmp_path): engine = await ConcurrentExecutionBackend(ThreadPoolExecutor()) asyncflow = await WorkflowEngine.create(engine) - cfg = SpecConfig.from_yaml(p) + cfg = WorkflowConfig.from_yaml(p) builder = LearnerBuilder(cfg, asyncflow) learner = builder.build() @@ -118,7 +118,7 @@ async def test_sequential_workflow_parameters_reach_tasks(tmp_path): engine = await ConcurrentExecutionBackend(ThreadPoolExecutor()) asyncflow = await WorkflowEngine.create(engine) - cfg = SpecConfig.from_yaml(p) + cfg = WorkflowConfig.from_yaml(p) builder = LearnerBuilder(cfg, asyncflow) learner = builder.build() @@ -193,7 +193,7 @@ async def test_parallel_workflow_parameters_reach_tasks(tmp_path): engine = await ConcurrentExecutionBackend(ThreadPoolExecutor()) asyncflow = await WorkflowEngine.create(engine) - cfg = SpecConfig.from_yaml(p) + cfg = WorkflowConfig.from_yaml(p) builder = LearnerBuilder(cfg, asyncflow) learner = builder.build() diff --git a/tests/unit/spec/test_builder.py b/tests/unit/spec/test_builder.py index 5fce0581..e3d23170 100644 --- a/tests/unit/spec/test_builder.py +++ b/tests/unit/spec/test_builder.py @@ -4,7 +4,7 @@ import pytest -from rose.spec.schema import SpecConfig +from rose.spec.schema import WorkflowConfig SEQ_WITH_PARAMS = textwrap.dedent("""\ @@ -93,7 +93,7 @@ def _make_builder(yaml_text, tmp_path): p = tmp_path / "spec.yaml" p.write_text(yaml_text) - cfg = SpecConfig.from_yaml(p) + cfg = WorkflowConfig.from_yaml(p) from rose.spec.builder import LearnerBuilder return LearnerBuilder(cfg, MagicMock()) diff --git a/tests/unit/spec/test_schema.py b/tests/unit/spec/test_schema.py index efb77c66..4105134a 100644 --- a/tests/unit/spec/test_schema.py +++ b/tests/unit/spec/test_schema.py @@ -3,7 +3,7 @@ import pytest -from rose.spec.schema import SpecConfig, TaskDef +from rose.spec.schema import WorkflowConfig, TaskDef # ── TaskDef ─────────────────────────────────────────────────────────────────── @@ -74,12 +74,12 @@ def test_sequential_spec_task_description_from_yaml(tmp_path): """) p = tmp_path / "td.yaml" p.write_text(yaml) - cfg = SpecConfig.from_yaml(p) + cfg = WorkflowConfig.from_yaml(p) assert cfg.simulation.task_description == {"cpu_count": 8} assert cfg.training.task_description is None -# ── SpecConfig — sequential learner ────────────────────────────────────────── +# ── WorkflowConfig — sequential learner ────────────────────────────────────────── SEQ_YAML = textwrap.dedent("""\ learner: @@ -111,7 +111,7 @@ def test_sequential_spec_task_description_from_yaml(tmp_path): def test_sequential_spec_roundtrip(tmp_path): p = tmp_path / "spec.yaml" p.write_text(SEQ_YAML) - cfg = SpecConfig.from_yaml(p) + cfg = WorkflowConfig.from_yaml(p) assert cfg.learner.type == "sequential_active_learner" assert cfg.learner.max_iter == 5 assert cfg.simulation.command == "python sim.py" @@ -144,7 +144,7 @@ def test_sequential_spec_missing_slot(tmp_path): p = tmp_path / "bad.yaml" p.write_text(yaml) with pytest.raises(ValueError, match="active_learn"): - SpecConfig.from_yaml(p) + WorkflowConfig.from_yaml(p) def test_sequential_spec_extra_slot(tmp_path): @@ -173,7 +173,7 @@ def test_sequential_spec_extra_slot(tmp_path): p = tmp_path / "bad.yaml" p.write_text(yaml) with pytest.raises(ValueError, match="Unexpected"): - SpecConfig.from_yaml(p) + WorkflowConfig.from_yaml(p) def test_unknown_learner_type(tmp_path): @@ -193,10 +193,10 @@ def test_unknown_learner_type(tmp_path): p = tmp_path / "bad.yaml" p.write_text(yaml) with pytest.raises(ValueError, match="Unknown learner type"): - SpecConfig.from_yaml(p) + WorkflowConfig.from_yaml(p) -# ── SpecConfig — parallel learner with learners ────────────────────────────── +# ── WorkflowConfig — parallel learner with learners ────────────────────────────── PAR_YAML = textwrap.dedent("""\ learner: @@ -238,7 +238,7 @@ def test_unknown_learner_type(tmp_path): def test_parallel_learners_roundtrip(tmp_path): p = tmp_path / "par.yaml" p.write_text(PAR_YAML) - cfg = SpecConfig.from_yaml(p) + cfg = WorkflowConfig.from_yaml(p) assert len(cfg.learners) == 2 assert cfg.learners[0].label == "rf" assert cfg.learners[1].tasks["training"].function == "mymod:train_mlp" @@ -266,14 +266,14 @@ def test_parallel_learner_missing_slot(tmp_path): p = tmp_path / "bad.yaml" p.write_text(yaml) with pytest.raises(ValueError, match="active_learn"): - SpecConfig.from_yaml(p) + WorkflowConfig.from_yaml(p) -# ── SpecConfig — parameters block ──────────────────────────────────────────── +# ── WorkflowConfig — parameters block ──────────────────────────────────────────── -# ── SpecConfig — reserved parameter keys ────────────────────────────────────── +# ── WorkflowConfig — reserved parameter keys ────────────────────────────────────── -# ── SpecConfig — task_description consistency across parallel learners ──────── +# ── WorkflowConfig — task_description consistency across parallel learners ──────── def test_parallel_learners_identical_task_description_valid(tmp_path): yaml = textwrap.dedent("""\ @@ -313,7 +313,7 @@ def test_parallel_learners_identical_task_description_valid(tmp_path): """) p = tmp_path / "ok.yaml" p.write_text(yaml) - cfg = SpecConfig.from_yaml(p) + cfg = WorkflowConfig.from_yaml(p) assert len(cfg.learners) == 2 @@ -356,7 +356,7 @@ def test_parallel_learners_different_task_description_raises(tmp_path): p = tmp_path / "bad.yaml" p.write_text(yaml) with pytest.raises(ValueError, match="task_description"): - SpecConfig.from_yaml(p) + WorkflowConfig.from_yaml(p) def test_parallel_learners_one_has_task_description_other_absent_raises(tmp_path): @@ -396,7 +396,7 @@ def test_parallel_learners_one_has_task_description_other_absent_raises(tmp_path p = tmp_path / "bad.yaml" p.write_text(yaml) with pytest.raises(ValueError, match="task_description"): - SpecConfig.from_yaml(p) + WorkflowConfig.from_yaml(p) def test_parameters_reserved_key_pythonpath(tmp_path): @@ -425,7 +425,7 @@ def test_parameters_reserved_key_pythonpath(tmp_path): p = tmp_path / "bad.yaml" p.write_text(yaml) with pytest.raises(ValueError, match="reserved keys"): - SpecConfig.from_yaml(p) + WorkflowConfig.from_yaml(p) def test_parameters_reserved_key_iteration(tmp_path): @@ -454,7 +454,7 @@ def test_parameters_reserved_key_iteration(tmp_path): p = tmp_path / "bad.yaml" p.write_text(yaml) with pytest.raises(ValueError, match="reserved keys"): - SpecConfig.from_yaml(p) + WorkflowConfig.from_yaml(p) def test_parameters_reserved_key_learner_id(tmp_path): @@ -483,7 +483,7 @@ def test_parameters_reserved_key_learner_id(tmp_path): p = tmp_path / "bad.yaml" p.write_text(yaml) with pytest.raises(ValueError, match="reserved keys"): - SpecConfig.from_yaml(p) + WorkflowConfig.from_yaml(p) def test_parameters_non_reserved_key_allowed(tmp_path): @@ -511,7 +511,7 @@ def test_parameters_non_reserved_key_allowed(tmp_path): """) p = tmp_path / "ok.yaml" p.write_text(yaml) - cfg = SpecConfig.from_yaml(p) + cfg = WorkflowConfig.from_yaml(p) assert cfg.parameters == {"dataset": "my_ds"} @@ -542,18 +542,18 @@ def test_parameters_roundtrip(tmp_path): """) p = tmp_path / "params.yaml" p.write_text(yaml) - cfg = SpecConfig.from_yaml(p) + cfg = WorkflowConfig.from_yaml(p) assert cfg.parameters == {"dataset": "my_ds", "scale": 1.5, "growing_pool": False} def test_parameters_defaults_empty(tmp_path): p = tmp_path / "spec.yaml" p.write_text(SEQ_YAML) - cfg = SpecConfig.from_yaml(p) + cfg = WorkflowConfig.from_yaml(p) assert cfg.parameters == {} -# ── SpecConfig — parallel learners (mixed types) ───────────────────────────── +# ── WorkflowConfig — parallel learners (mixed types) ───────────────────────────── def test_parallel_learners_mixed_types(tmp_path): yaml = textwrap.dedent("""\ @@ -590,4 +590,4 @@ def test_parallel_learners_mixed_types(tmp_path): p = tmp_path / "bad.yaml" p.write_text(yaml) with pytest.raises(ValueError, match="mixed types"): - SpecConfig.from_yaml(p) + WorkflowConfig.from_yaml(p) diff --git a/tests/unit/spec/test_workflow_spec.py b/tests/unit/spec/test_workflow_spec.py index 76d9cafb..fcc76a37 100644 --- a/tests/unit/spec/test_workflow_spec.py +++ b/tests/unit/spec/test_workflow_spec.py @@ -4,7 +4,7 @@ import pytest from rose.spec import WorkflowSpec, load_spec -from rose.spec.schema import SpecConfig +from rose.spec.schema import WorkflowConfig SEQ_YAML = textwrap.dedent("""\ learner: From 59990d3df5c6cca6df37b50af286775d6f0c2391 Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Thu, 18 Jun 2026 23:33:34 +0200 Subject: [PATCH 06/13] improve docs --- docs/index.md | 2 +- docs/scripts/custom.js | 38 +++++++ docs/styles/custom.css | 216 +++++++++++++++++++++++++++++++++--- docs/user-guide/spec-api.md | 4 +- mkdocs.yml | 9 +- 5 files changed, 246 insertions(+), 23 deletions(-) create mode 100644 docs/scripts/custom.js diff --git a/docs/index.md b/docs/index.md index 4abd15d4..40c14e0d 100644 --- a/docs/index.md +++ b/docs/index.md @@ -9,7 +9,7 @@ ROSE leverages [**RADICAL-Cybertools**](https://radical-cybertools.github.io), a ## Why ROSE? 🚀🚀🚀 ROSE allows you to enable, scale, and accelerate your learning workflows across thousands of CPU cores and GPUs effectively and efficiently with just a few lines of code. -ROSE is built on the [**RADICAL-AsyncFlow**](https://radical-cybertools.github.io/radical.asyncflow/) and [**RADICAL-Pilot**](https://github.com/radical-cybertools/radical.pilot) runtime system, a powerful execution engine that enables the distributed execution of millions of scientific tasks and applications such as executables, functions and containers effortlessly. +ROSE is built on the [**RADICAL-AsyncFlow**](https://radical-cybertools.github.io/radical.asyncflow/) and [**RHAPSODY**](https://radical-cybertools.github.io/rhapsody/landing.html) runtime system, a powerful execution engine that enables the distributed execution of millions of scientific tasks and applications such as executables, functions and containers effortlessly. **Preemption-safe on HPC.** Every completed iteration is written to disk before the next one starts. If the job is killed mid-run, all completed iterations are already on disk — inspect them, resume from the last checkpoint, and never rerun a finished iteration. diff --git a/docs/scripts/custom.js b/docs/scripts/custom.js new file mode 100644 index 00000000..7e385009 --- /dev/null +++ b/docs/scripts/custom.js @@ -0,0 +1,38 @@ +// ROSE docs — subtle scroll-reveal for a more fluid reading experience. +// Pairs with the .rose-reveal / .rose-revealed rules in styles/custom.css. +(() => { + const reduceMotion = window.matchMedia("(prefers-reduced-motion: reduce)").matches; + if (reduceMotion || !("IntersectionObserver" in window)) return; + + const SELECTOR = + ".md-content__inner > h2, .md-content__inner > h3, .md-content__inner > p, " + + ".md-content__inner > table, .md-content__inner > .admonition, " + + ".md-content__inner > pre, .md-content__inner > .highlight"; + + const observer = new IntersectionObserver( + (entries) => { + entries.forEach((entry) => { + if (entry.isIntersecting) { + entry.target.classList.add("rose-revealed"); + observer.unobserve(entry.target); + } + }); + }, + { threshold: 0.05, rootMargin: "0px 0px -40px 0px" } + ); + + const observe = () => { + document.querySelectorAll(SELECTOR).forEach((el) => { + el.classList.add("rose-reveal"); + observer.observe(el); + }); + }; + + // mkdocs-material's instant navigation swaps content without a full page + // load, so re-run on each page change via its document$ observable. + if (typeof document$ !== "undefined") { + document$.subscribe(observe); + } else { + document.addEventListener("DOMContentLoaded", observe); + } +})(); diff --git a/docs/styles/custom.css b/docs/styles/custom.css index c21edcfb..df3e7221 100644 --- a/docs/styles/custom.css +++ b/docs/styles/custom.css @@ -1,15 +1,186 @@ -/* Change the banner background color to dark rose */ +/* ============================================================ + ROSE docs — fluid theme refinement + Cohesive rose/pink brand palette, softened rounded surfaces, + smooth motion, and viewport-responsive (fluid) type & spacing. + ============================================================ */ + +/* ---- 1. Brand palette --------------------------------------- */ +:root { + --rose-50: #fdf2f6; + --rose-100: #fbe4ee; + --rose-300: #f06ea6; + --rose-500: #d81b60; + --rose-600: #c2185b; + --rose-700: #a3134c; + --rose-900: #6e0e35; + + --md-primary-fg-color: var(--rose-600); + --md-primary-fg-color--light: var(--rose-300); + --md-primary-fg-color--dark: var(--rose-900); + --md-accent-fg-color: var(--rose-500); + --md-accent-fg-color--transparent: rgba(216, 27, 96, 0.1); + + --rose-radius: 0.6rem; + --rose-radius-lg: 1rem; + --rose-shadow: 0 2px 10px rgba(194, 24, 91, 0.08), 0 1px 2px rgba(0, 0, 0, 0.04); + --rose-shadow-hover: 0 6px 20px rgba(194, 24, 91, 0.16), 0 2px 4px rgba(0, 0, 0, 0.06); + --rose-transition: 180ms cubic-bezier(0.4, 0, 0.2, 1); +} + +[data-md-color-scheme="slate"] { + --md-primary-fg-color: var(--rose-700); + --md-primary-fg-color--dark: #4a0a26; + --md-accent-fg-color: var(--rose-300); + --rose-shadow: 0 2px 12px rgba(0, 0, 0, 0.35); + --rose-shadow-hover: 0 8px 24px rgba(0, 0, 0, 0.45); +} + +/* ---- 2. Fluid typography & layout ----------------------------- */ +html { + font-size: clamp(18px, 0.6vw + 16px, 22px); + scroll-behavior: smooth; +} + +.md-typeset h1 { font-size: clamp(1.8rem, 1.2rem + 1.6vw, 2.6rem); } +.md-typeset h2 { font-size: clamp(1.4rem, 1.1rem + 1vw, 2rem); } +.md-typeset h3 { font-size: clamp(1.15rem, 1rem + 0.5vw, 1.5rem); } + +.md-typeset { + line-height: 1.65; +} + +.md-content__inner { + padding: clamp(0.8rem, 1vw + 0.5rem, 1.8rem) clamp(0.6rem, 2vw, 2rem); +} + +.md-grid { + max-width: clamp(960px, 90vw, 1440px); +} + +/* ---- 3. Header & tabs: subtle gradient instead of flat fill --- */ .md-header { - background-color: #c2185b !important; /* Dark rose color */ + background: linear-gradient(135deg, var(--rose-600), var(--rose-700)); + box-shadow: 0 2px 12px rgba(194, 24, 91, 0.25); + transition: box-shadow var(--rose-transition), transform var(--rose-transition); +} + +.md-tabs { + background: linear-gradient(135deg, var(--rose-700), var(--rose-900)); +} + +/* ---- 4. Motion: links, nav, tabs ------------------------------ */ +a, +.md-nav__link, +.md-tabs__link, +.md-typeset a { + transition: color var(--rose-transition), background-color var(--rose-transition), + opacity var(--rose-transition); } -/* Ensure the active navigation link is distinct */ .md-nav__link--active { - color: #c2185b !important; /* Slightly darker black for active link */ - font-weight: bold; + color: var(--rose-600) !important; + font-weight: 600; +} + +[data-md-color-scheme="slate"] .md-nav__link--active { + color: var(--rose-300) !important; +} + +.md-nav__link:hover { + opacity: 0.75; } +.md-typeset a:hover { + text-decoration: underline; +} + +/* ---- 5. Softened, elevated surfaces ---------------------------- */ +.md-typeset .admonition, +.md-typeset details { + border-radius: var(--rose-radius); + box-shadow: var(--rose-shadow); + border-width: 0; + border-left: 4px solid var(--md-accent-fg-color); + overflow: hidden; + transition: box-shadow var(--rose-transition); +} + +.md-typeset .admonition:hover, +.md-typeset details:hover { + box-shadow: var(--rose-shadow-hover); +} + +.md-typeset pre, +.md-typeset .highlight { + border-radius: var(--rose-radius); + box-shadow: var(--rose-shadow); +} +.md-typeset code { + border-radius: 0.3rem; +} + +.md-typeset table:not([class]) { + border-radius: var(--rose-radius); + overflow: hidden; + box-shadow: var(--rose-shadow); + border-collapse: separate; + border-spacing: 0; +} + +.md-typeset table:not([class]) th { + background-color: var(--md-accent-fg-color--transparent); +} + +.md-typeset table:not([class]) tr { + transition: background-color var(--rose-transition); +} + +.md-typeset table:not([class]) tr:hover { + background-color: var(--md-accent-fg-color--transparent); +} + +.md-typeset img, +.md-typeset figure { + border-radius: var(--rose-radius-lg); +} + +/* ---- 6. Search box: rounded, smooth focus ---------------------- */ +.md-search__form { + border-radius: 2rem; + transition: background-color var(--rose-transition), box-shadow var(--rose-transition); +} + +.md-search__input:focus ~ .md-search__form, +.md-search__form:hover { + box-shadow: 0 0 0 3px var(--md-accent-fg-color--transparent); +} + +/* ---- 7. Buttons / copy button ----------------------------------- */ +.md-typeset .md-button, +.md-clipboard { + border-radius: var(--rose-radius); + transition: transform var(--rose-transition), box-shadow var(--rose-transition); +} + +.md-typeset .md-button:hover { + transform: translateY(-1px); + box-shadow: var(--rose-shadow); +} + +/* ---- 8. Scroll-reveal (paired with scripts/custom.js) ----------- */ +.rose-reveal { + opacity: 0; + transform: translateY(8px); + transition: opacity 420ms ease, transform 420ms ease; +} + +.rose-reveal.rose-revealed { + opacity: 1; + transform: translateY(0); +} + +/* ---- 9. Config-reference table (existing component, restyled) --- */ table#config { font-family: monospace; font-size: 13px; @@ -22,20 +193,26 @@ table#config span.value { font-size: 0.9em; filter: opacity(70%); display: inline-block; - padding: 0px 4px; - border-radius: 5px; + padding: 1px 6px; + border-radius: 999px; } table#config span.example { - background-color: lightBlue; + background-color: var(--rose-100); + color: var(--rose-700); } table#config span.value { - background-color: lightGray; + background-color: rgba(0, 0, 0, 0.06); } table#config td.comment:nth-child(3) { - border-left: 2px solid whiteSmoke; + border-left: 2px solid rgba(0, 0, 0, 0.08); } table#config td.type { font-size: 0.8em; + filter: opacity(70%); + color: #999999; +} +table#config tr { + transition: background-color var(--rose-transition); } table#config tr:hover code { background-color: hsla(0, 0%, 100%, 1); @@ -53,10 +230,21 @@ table#config td .anchor:target::before { content: ""; display: block; } -table#config td.type { - filter: opacity(70%); - color: #999999; -} .md-typeset__scrollwrap { overflow-x: inherit; } + +/* ---- 10. Respect reduced-motion preference ----------------------- */ +@media (prefers-reduced-motion: reduce) { + html { + scroll-behavior: auto; + } + * { + transition-duration: 0.01ms !important; + animation-duration: 0.01ms !important; + } + .rose-reveal { + opacity: 1; + transform: none; + } +} diff --git a/docs/user-guide/spec-api.md b/docs/user-guide/spec-api.md index 564923db..69e20575 100644 --- a/docs/user-guide/spec-api.md +++ b/docs/user-guide/spec-api.md @@ -147,7 +147,7 @@ async for state in learner.start( -Task functions referenced by `function: tasks:simulate` are ordinary Python callables in a `tasks.py` module — no ROSE decorators, no async required: +Task functions referenced by `function: tasks:simulate` are ordinary Python callables in a `tasks.py` module — no ROSE decorators required: ```python # tasks.py @@ -406,7 +406,7 @@ def simulate(*args, **kwargs): return {"X": ..., "y": ...} ``` -The function may be sync or async. The return value is passed as the first positional argument to the next task in the chain. +The function must be async. The return value is passed as the first positional argument to the next task in the chain.
diff --git a/mkdocs.yml b/mkdocs.yml index 5c82ca3e..5efc9a7c 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -6,10 +6,6 @@ edit_uri: "" theme: name: material - features: - - navigation.tabs - - navigation.sections - - toc.integrate palette: - media: "(prefers-color-scheme: light)" scheme: default @@ -17,14 +13,14 @@ theme: icon: material/weather-night name: Switch to dark mode primary: black - accent: indigo + accent: pink - media: "(prefers-color-scheme: dark)" scheme: slate toggle: icon: material/weather-sunny name: Switch to light mode primary: black - accent: indigo + accent: pink highlightjs: true hljs_languages: @@ -37,6 +33,7 @@ theme: - navigation.top - navigation.tabs - navigation.tabs.sticky + - toc.integrate - search.highlight - search.share - search.suggest From 44512ff8779367593287b0ead3eff46ea81405c0 Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Sat, 20 Jun 2026 00:39:23 +0200 Subject: [PATCH 07/13] This commit: 1-Improve docs 2-Remove RadicalExecutionBackend from docs 3-Address Gemini comments --- docs/index.md | 2 +- docs/user-guide/advanced-rl-workflow.md | 7 +--- docs/user-guide/basic-acl-workflow.md | 6 ++- docs/user-guide/basic-rl-workflow.md | 9 +++-- docs/user-guide/parallel_learners_docs.md | 10 ++--- docs/user-guide/rose-aas.md | 2 +- docs/user-guide/spec-api.md | 14 +++++-- docs/user-guide/target-resources.md | 16 ++++---- rose/spec/builder.py | 30 +++++++++----- rose/spec/schema.py | 2 +- tests/unit/spec/test_builder.py | 45 +++++++++++++++++++++ tests/unit/spec/test_schema.py | 48 +++++++++++++++++++++++ 12 files changed, 151 insertions(+), 40 deletions(-) diff --git a/docs/index.md b/docs/index.md index 40c14e0d..800f1b89 100644 --- a/docs/index.md +++ b/docs/index.md @@ -1,6 +1,6 @@ ## What is ROSE? 🌹 -ROSE: RADICAL Orchestrator for Surrogate Exploration (ROSE) toolkit is a framework designed to enable the distribuited, concurrent and adaptive execution of simulation, surrogate training, and selection tasks on High-Performance Computing (HPC) resources at a scale. ROSE is a Python package that provides tools for developing active and reinforcement learning (AL and RL) methods for scientific applications. It enables users to define simulation and surrogate training tasks and automatically manage their execution on HPC resources via a predefined set of Learning Policies (Learners). +ROSE: RADICAL Orchestrator for Surrogate Exploration (ROSE) toolkit is a framework designed to enable the distributed, concurrent and adaptive execution of simulation, surrogate training, and selection tasks on High-Performance Computing (HPC) resources at a scale. ROSE is a Python package that provides tools for developing active and reinforcement learning (AL and RL) methods for scientific applications. It enables users to define simulation and surrogate training tasks and automatically manage their execution on HPC resources via a predefined set of Learning Policies (Learners). ROSE also includes tools to facilitate the selection of the most effective surrogate model for a given simulation based on performance metrics. diff --git a/docs/user-guide/advanced-rl-workflow.md b/docs/user-guide/advanced-rl-workflow.md index 9091a418..6c55920f 100644 --- a/docs/user-guide/advanced-rl-workflow.md +++ b/docs/user-guide/advanced-rl-workflow.md @@ -21,7 +21,7 @@ Import ROSE parallel RL modules: ```python from radical.asyncflow import WorkflowEngine -from rhapsody.backends import RadicalExecutionBackend +from rhapsody.backends import DragonExecutionBackendV3 from rose.rl.reinforcement_learner import SequentialReinforcementLearner ``` @@ -32,10 +32,7 @@ from rose.rl.reinforcement_learner import SequentialReinforcementLearner async def main(): - execution_engine = await RadicalExecutionBackend( - {'runtime': 30, - 'resource': 'local.localhost'} - ) + execution_engine = await DragonExecutionBackendV3() asyncflow = await WorkflowEngine.create(execution_engine) diff --git a/docs/user-guide/basic-acl-workflow.md b/docs/user-guide/basic-acl-workflow.md index 485f1626..91b88146 100644 --- a/docs/user-guide/basic-acl-workflow.md +++ b/docs/user-guide/basic-acl-workflow.md @@ -5,15 +5,17 @@ Import ROSE main modules: from rose.metrics import MEAN_SQUARED_ERROR_MSE from rose.al.active_learner import SequentialActiveLearner +from concurrent.futures import ProcessPoolExecutor + from radical.asyncflow import WorkflowEngine -from rhapsody.backends import RadicalExecutionBackend +from rhapsody.backends import ConcurrentExecutionBackend ``` Define your resource engine, as we described in our previous [Target Resources](target-resources.md) step: ```python -engine = await RadicalExecutionBackend({'resource': 'local.localhost'}) +engine = await ConcurrentExecutionBackend(ProcessPoolExecutor()) asyncflow = await WorkflowEngine.create(engine) acl = SequentialActiveLearner(asyncflow) diff --git a/docs/user-guide/basic-rl-workflow.md b/docs/user-guide/basic-rl-workflow.md index a40152c7..dbee4a4b 100644 --- a/docs/user-guide/basic-rl-workflow.md +++ b/docs/user-guide/basic-rl-workflow.md @@ -4,16 +4,19 @@ Import ROSE main modules: ```python +from concurrent.futures import ProcessPoolExecutor + from radical.asyncflow import WorkflowEngine -from rhapsody.backends import RadicalExecutionBackend +from rhapsody.backends import ConcurrentExecutionBackend from rose.metrics import GREATER_THAN_THRESHOLD from rose.rl.reinforcement_learner import SequentialReinforcementLearner ``` -Define your resource engine, as we described in our previous [Target Resources](target-resources.md) step: +You can use and setup an HPC engine as we described in our previous [Target Resources](target-resources.md) step: + ```python -engine = await RadicalExecutionBackend({'resource': 'local.localhost'}) +engine = await ConcurrentExecutionBackend(ProcessPoolExecutor()) asyncflow = await WorkflowEngine.create(engine) rl = SequentialReinforcementLearner(asyncflow) diff --git a/docs/user-guide/parallel_learners_docs.md b/docs/user-guide/parallel_learners_docs.md index 2b458caf..e70e1060 100644 --- a/docs/user-guide/parallel_learners_docs.md +++ b/docs/user-guide/parallel_learners_docs.md @@ -78,14 +78,12 @@ from rose import LearnerConfig from rose.al import ParallelActiveLearner from rose.metrics import MEAN_SQUARED_ERROR_MSE +from concurrent.futures import ProcessPoolExecutor + from radical.asyncflow import WorkflowEngine -from rhapsody.backends import RadicalExecutionBackend +from rhapsody.backends import ConcurrentExecutionBackend -engine = await RadicalExecutionBackend( - {'runtime': 30, - 'resource': 'local.localhost' - } - ) +engine = await ConcurrentExecutionBackend(ProcessPoolExecutor()) asyncflow = await WorkflowEngine.create(engine) acl = ParallelActiveLearner(asyncflow) code_path = f'{sys.executable} {os.getcwd()}' diff --git a/docs/user-guide/rose-aas.md b/docs/user-guide/rose-aas.md index 96f2fb9b..ff98cb3d 100644 --- a/docs/user-guide/rose-aas.md +++ b/docs/user-guide/rose-aas.md @@ -26,7 +26,7 @@ A ROSE learner doesn't run anything itself — it submits tasks through whicheve | Backend | Where the orchestration loop runs | Where tasks execute | Documented at | |---|---|---|---| -| `RadicalExecutionBackend` (RHAPSODY) | Wherever your script runs — it submits a pilot job and blocks inside it | Inside the pilot job it just submitted | [Target Resources](target-resources.md) | +| `DragonExecutionBackendV3` (RHAPSODY) | Wherever your script runs — it submits a pilot job and blocks inside it | Inside the pilot job it just submitted | [Target Resources](target-resources.md) | | Edge backend (RHAPSODY, `bridge_url` + `edge_name`) | Wherever your script runs — does **not** need to be the HPC machine | Inside a separate, already-running edge agent — possibly on a different machine entirely | this page | The first model is "submit a job, then run my loop inside it." The second is "run my loop here; dispatch tasks to a job running somewhere else." That second model is what makes ROSE service-like: the loop's control plane (your `learner.start()` call, deciding when to stop, talking to MLflow/ClearML) is decoupled from the compute plane (the HPC allocation actually executing `simulate`/`train`). diff --git a/docs/user-guide/spec-api.md b/docs/user-guide/spec-api.md index 69e20575..0928707f 100644 --- a/docs/user-guide/spec-api.md +++ b/docs/user-guide/spec-api.md @@ -6,10 +6,16 @@ The YAML spec API lets you declare a ROSE workflow as a data file instead of Pyt !!! note - The YAML spec currently supports only the four built-in `learner.type` values listed - above. Custom `Learner` subclasses are not yet expressible in YAML — `LearnerBuilder` - raises `ValueError` for any other `learner.type` value. Use the Python API - (decorator-based, e.g. `SequentialActiveLearner(asyncflow)`) if you need a custom + The YAML spec currently supports only the four built-in `learner.type` values listed below: + + - `SequentialActiveLearner` + - `ParallelActiveLearner` + - `SequentialReinforcementLearner` + - `ParallelReinforcementLearner` + - `SequentialUQLearner` + - `ParallelUQLearner` + + Custom `Learner` subclasses are not yet expressible in YAML — `LearnerBuilder` raises `ValueError` for any other `learner.type` value. Use the Python API (decorator-based, e.g. `SequentialActiveLearner(asyncflow)`) if you need a custom learner implementation. diff --git a/docs/user-guide/target-resources.md b/docs/user-guide/target-resources.md index 086dd196..7f4c523e 100644 --- a/docs/user-guide/target-resources.md +++ b/docs/user-guide/target-resources.md @@ -1,5 +1,5 @@ # Target Machines for Executing AL Workflows -ROSE enables the orchestration of ML Surrogate building workflows on diverse computing resources using [radical.asyncflow](https://github.com/radical-cybertools/radical.asyncflow). Below, we will show how you can specify your `local computer` and `remote HPC machine` as target resources using the `RadicalExecutionBackend` from [RHAPSODY](https://github.com/radical-cybertools/rhapsody). +ROSE enables the orchestration of ML Surrogate building workflows on diverse computing resources using [radical.asyncflow](https://github.com/radical-cybertools/radical.asyncflow). Below, we will show how you can specify your `local computer` and `remote HPC machine` as target resources using the `DragonExecutionBackendV3` from [RHAPSODY](https://radical-cybertools.github.io/rhapsody/getting-started/advanced-usage/?h=hpc#multiple-execution-backends). ## Local Computer For local execution, user can use their desktops, laptops, and their own small clusters to execute their AL workflows as follows: @@ -22,23 +22,23 @@ acl = SequentialActiveLearner(asyncflow) ## HPC Resources -!!! warning - `RadicalExecutionBackend` will be deprecated by the end of Q3 2026. +To execute AL workflows on HPC machines, users must have an active allocation on the target machine and specify their resource requirements to execute their workflows. Remember, ROSE uses `DragonExecutionBackendV3` from [RHAPSODY](https://github.com/radical-cybertools/rhapsody) (`rhapsody-py`) which is an interface for multip execution and AI runtime system. For more information on how to access, set up, and execute workflows on HPC machines, refer to the following link [ROSE with RHAPSODY on HPC](https://radical-cybertools.github.io/rhapsody/getting-started/advanced-usage/?h=hpc#hpc-workloads-with-dragon): -To execute AL workflows on HPC machines, users must have an active allocation on the target machine and specify their resource requirements, as well as the time needed to execute their workflows. Remember, ROSE uses `RadicalExecutionBackend` from [RHAPSODY](https://github.com/radical-cybertools/rhapsody) (`rhapsody-py`) which is an interface for RADICAL-Pilot runtime system. For more information on how to access, set up, and execute workflows on HPC machines, refer to the following link [RADICAL-Pilot Job Submission](https://radicalpilot.readthedocs.io/en/stable/tutorials/submission.html): + +!!! note + For any ROSE script that uses `DragonExecutionBackendV3`, user must run the + rose_script.py with `dragon` binary instead of `python`. Please refer to the following link for more information: [use ROSE with RHAPSODY-Dragon on HPC](https://radical-cybertools.github.io/rhapsody/hpc-machines/#backend-compatibility) ```python import os from radical.asyncflow import WorkflowEngine -from rhapsody.backends import RadicalExecutionBackend +from rhapsody.backends import DragonExecutionBackendV3 from rose.al.active_learner import SequentialActiveLearner -hpc_engine = await RadicalExecutionBackend( - {'runtime': 30, 'cores': 4096, - 'gpus' : 4, 'resource': 'tacc.frontera'}) +hpc_engine = await DragonExecutionBackendV3() asyncflow = await WorkflowEngine.create(hpc_engine) diff --git a/rose/spec/builder.py b/rose/spec/builder.py index 89dc217d..83964d14 100644 --- a/rose/spec/builder.py +++ b/rose/spec/builder.py @@ -101,25 +101,37 @@ def build_learner_config(self): return LearnerConfig(**{slot: schedule for slot in lc_slots}, criterion=schedule) def build_learner_configs(self): - """LearnerConfig list for parallel learners, or None if no learners defined.""" + """LearnerConfig list for parallel learners. + + Returns None for non-parallel learner types (use build_learner_config + instead) or when there's nothing to inject: no per-learner `learners:` + block, no `parameters`, and no `remote.pythonpath`. + + When `learners:` is absent (shared tasks across learners), the list + length falls back to `learner.parallel_learners` so that `parameters`/ + `pythonpath` still reach every learner's task kwargs. + """ cfg = self.config - if cfg.learners is None: + if cfg.learner.type != "parallel_active_learner": + return None + if cfg.learners is None and not cfg.parameters and not cfg.remote.pythonpath: return None from rose.learner import LearnerConfig, TaskConfig - required = _REQUIRED_SLOTS[cfg.learner.type] - configs = [] - max_iter = cfg.learner.max_iter - params = dict(cfg.parameters) + required = _REQUIRED_SLOTS[cfg.learner.type] + configs = [] + max_iter = cfg.learner.max_iter + params = dict(cfg.parameters) params["pythonpath"] = list(cfg.remote.pythonpath) - lc_slots = required & _LEARNER_CONFIG_SLOTS - for i, learner_def in enumerate(cfg.learners): + lc_slots = required & _LEARNER_CONFIG_SLOTS + learner_defs = cfg.learners or [None] * cfg.learner.parallel_learners + for i, learner_def in enumerate(learner_defs): # learner_id → dispatch routing key (popped by closure, never reaches user fn) # iteration → per-entry so get_task_config(slot, n) returns the right value # learner_label → human-readable learner name; only injected when non-empty # parameters → user-defined values from the YAML parameters: block base = {"learner_id": i, **params} - if learner_def.label: + if learner_def is not None and learner_def.label: base["learner_label"] = learner_def.label schedule = {n: TaskConfig(kwargs={**base, "iteration": n}) for n in range(max_iter + 1)} schedule[-1] = TaskConfig(kwargs={**base, "iteration": max_iter}) diff --git a/rose/spec/schema.py b/rose/spec/schema.py index 17cef4c9..8312a351 100644 --- a/rose/spec/schema.py +++ b/rose/spec/schema.py @@ -148,7 +148,7 @@ def _validate_task_slots(self) -> "WorkflowConfig": "All learners must use the same type for a given slot." ) descs = [dict(l.tasks[slot].task_description or {}) for l in self.learners] - if len({str(sorted(d.items())) for d in descs}) > 1: + if len(descs) > 1 and any(d != descs[0] for d in descs[1:]): raise ValueError( f"Slot '{slot}' has different task_description values across learners. " "asyncflow registers task_description once per slot at registration time — " diff --git a/tests/unit/spec/test_builder.py b/tests/unit/spec/test_builder.py index e3d23170..e71b78de 100644 --- a/tests/unit/spec/test_builder.py +++ b/tests/unit/spec/test_builder.py @@ -171,6 +171,51 @@ def test_build_learner_configs_learner_1_kwargs(tmp_path): assert kw["iteration"] == 1 +PAR_SHARED_WITH_PARAMS = textwrap.dedent("""\ + learner: + type: parallel_active_learner + max_iter: 2 + parallel_learners: 3 + simulation: + type: python + function: mymod:sim + training: + type: python + function: mymod:train + active_learn: + type: python + function: mymod:select + stop_criterion: + metric: r2 + threshold: 0.9 + operator: ">" + evaluator: + type: python + function: mymod:eval + parameters: + dataset: prod_ds + lr: 0.001 +""") + + +def test_build_learner_configs_shared_tasks_with_parameters_falls_back_to_parallel_learners( + tmp_path, +): + from rose.learner import LearnerConfig + builder = _make_builder(PAR_SHARED_WITH_PARAMS, tmp_path) + lcs = builder.build_learner_configs() + assert lcs is not None + assert len(lcs) == 3 # learner.parallel_learners, since learners: is absent + assert all(isinstance(lc, LearnerConfig) for lc in lcs) + kw0 = lcs[0].simulation[0].kwargs + assert kw0["learner_id"] == 0 + assert kw0["dataset"] == "prod_ds" + assert "learner_label" not in kw0 # no per-learner label when learners: is absent + kw2 = lcs[2].simulation[1].kwargs + assert kw2["learner_id"] == 2 + assert kw2["iteration"] == 1 + + # ── pythonpath auto-injection ───────────────────────────────────────────────── SEQ_WITH_PYTHONPATH = textwrap.dedent("""\ diff --git a/tests/unit/spec/test_schema.py b/tests/unit/spec/test_schema.py index 4105134a..71e7c7ff 100644 --- a/tests/unit/spec/test_schema.py +++ b/tests/unit/spec/test_schema.py @@ -317,6 +317,54 @@ def test_parallel_learners_identical_task_description_valid(tmp_path): assert len(cfg.learners) == 2 +def test_parallel_learners_nested_task_description_different_key_order_valid(tmp_path): + """Nested dicts with the same keys/values in a different insertion order + must compare equal — dict equality, not string-sorted serialization.""" + yaml = textwrap.dedent("""\ + learner: + type: parallel_active_learner + learners: + - label: a + simulation: + type: shell + command: python sim.py + task_description: + resources: + gpus: 1 + cpus: 4 + training: + type: python + function: mymod:train + active_learn: + type: python + function: mymod:select + - label: b + simulation: + type: shell + command: python sim.py --model b + task_description: + resources: + cpus: 4 + gpus: 1 + training: + type: python + function: mymod:train + active_learn: + type: python + function: mymod:select + stop_criterion: + metric: r2 + threshold: 0.9 + evaluator: + type: python + function: mymod:eval + """) + p = tmp_path / "ok_nested.yaml" + p.write_text(yaml) + cfg = WorkflowConfig.from_yaml(p) + assert len(cfg.learners) == 2 + + def test_parallel_learners_different_task_description_raises(tmp_path): yaml = textwrap.dedent("""\ learner: From fa090516898eaad6c6847d666ba4ab332e71e6bb Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Sat, 20 Jun 2026 00:47:32 +0200 Subject: [PATCH 08/13] fix pre-commit --- _typos.toml | 3 + docs/user-guide/spec-api.md | 2 +- rose/spec/__init__.py | 26 ++++--- rose/spec/adapters.py | 16 ++-- rose/spec/builder.py | 22 +++--- rose/spec/schema.py | 82 ++++++++++---------- tests/integration/spec/helpers.py | 1 + tests/integration/spec/test_yaml_workflow.py | 5 +- tests/unit/spec/test_adapters.py | 23 ++++-- tests/unit/spec/test_builder.py | 10 ++- tests/unit/spec/test_schema.py | 20 +++-- tests/unit/spec/test_workflow_spec.py | 21 +++-- 12 files changed, 132 insertions(+), 99 deletions(-) create mode 100644 _typos.toml diff --git a/_typos.toml b/_typos.toml new file mode 100644 index 00000000..82f6f264 --- /dev/null +++ b/_typos.toml @@ -0,0 +1,3 @@ +[default.extend-words] +# "aas" appears in rose-aas.md / "ROSE as a Service" links — not a typo of "as"/"ass". +aas = "aas" diff --git a/docs/user-guide/spec-api.md b/docs/user-guide/spec-api.md index 0928707f..0029ca7c 100644 --- a/docs/user-guide/spec-api.md +++ b/docs/user-guide/spec-api.md @@ -14,7 +14,7 @@ The YAML spec API lets you declare a ROSE workflow as a data file instead of Pyt - `ParallelReinforcementLearner` - `SequentialUQLearner` - `ParallelUQLearner` - + Custom `Learner` subclasses are not yet expressible in YAML — `LearnerBuilder` raises `ValueError` for any other `learner.type` value. Use the Python API (decorator-based, e.g. `SequentialActiveLearner(asyncflow)`) if you need a custom learner implementation. diff --git a/rose/spec/__init__.py b/rose/spec/__init__.py index 8ee9acbd..057f2972 100644 --- a/rose/spec/__init__.py +++ b/rose/spec/__init__.py @@ -1,7 +1,8 @@ from __future__ import annotations +from collections.abc import Callable, Coroutine from pathlib import Path -from typing import Any, Callable, Coroutine +from typing import Any from .builder import LearnerBuilder from .schema import WorkflowConfig @@ -9,7 +10,7 @@ __all__ = ["load_spec", "WorkflowSpec"] -def load_spec(path: str | Path, validate_imports: bool = False) -> "WorkflowSpec": +def load_spec(path: str | Path, validate_imports: bool = False) -> WorkflowSpec: """Load and validate a YAML workflow spec. Raises ValueError on schema errors. Args: @@ -26,12 +27,12 @@ def load_spec(path: str | Path, validate_imports: bool = False) -> "WorkflowSpec return spec -def _collect_python_specs(cfg: "WorkflowConfig") -> list[str]: +def _collect_python_specs(cfg: WorkflowConfig) -> list[str]: """Return deduplicated 'module:callable' strings for all Python tasks in the spec.""" specs: list[str] = [] if cfg.learners is not None: - for l in cfg.learners: - for td in l.tasks.values(): + for ld in cfg.learners: + for td in ld.tasks.values(): if td.type == "python": specs.append(td.function) else: @@ -43,8 +44,11 @@ def _collect_python_specs(cfg: "WorkflowConfig") -> list[str]: return list(dict.fromkeys(specs)) # deduplicate, preserve order -def _validate_task_imports(cfg: "WorkflowConfig") -> None: - """Try to import every Python task callable. Raises ValueError listing all failures.""" +def _validate_task_imports(cfg: WorkflowConfig) -> None: + """Try to import every Python task callable. + + Raises ValueError listing all failures. + """ import importlib import sys as _sys @@ -78,7 +82,7 @@ class WorkflowSpec: def __init__(self, config: WorkflowConfig) -> None: self.config = config - def workflow_with(self, **overrides: Any) -> "WorkflowSpec": + def workflow_with(self, **overrides: Any) -> WorkflowSpec: """Return a new WorkflowSpec with selective overrides applied. Accepted keys: @@ -111,9 +115,7 @@ async def _workflow(bridge_url: str, edge_name: str) -> None: import rhapsody from radical.asyncflow import WorkflowEngine - engine = await rhapsody.get_backend( - "edge", bridge_url=bridge_url, edge_name=edge_name - ) + engine = await rhapsody.get_backend("edge", bridge_url=bridge_url, edge_name=edge_name) asyncflow = await WorkflowEngine.create(engine) builder = LearnerBuilder(cfg, asyncflow) @@ -124,7 +126,7 @@ async def _workflow(bridge_url: str, edge_name: str) -> None: lcs = builder.build_learner_configs() if lcs is not None: start_kwargs["parallel_learners"] = len(lcs) - start_kwargs["learner_configs"] = lcs + start_kwargs["learner_configs"] = lcs else: start_kwargs["parallel_learners"] = cfg.learner.parallel_learners else: diff --git a/rose/spec/adapters.py b/rose/spec/adapters.py index e5480f1d..e8910108 100644 --- a/rose/spec/adapters.py +++ b/rose/spec/adapters.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Callable +from collections.abc import Callable from .schema import RemoteConfig, TaskDef @@ -24,7 +24,9 @@ def make_dispatch_closure( remote: RemoteConfig, ) -> Callable: """Dispatch closure for parallel learners: routes per-learner based on learner_id kwarg. - task_description uses the first learner's value — asyncflow reads it once at registration.""" + + task_description uses the first learner's value — asyncflow reads it once at registration. + """ td = dict(task_defs[0].task_description or {}) if task_defs[0].type == "shell": return _make_shell_dispatch( @@ -40,7 +42,7 @@ def make_dispatch_closure( def _make_shell_closure(command: str, task_description: dict) -> Callable: _cmd = command - _td = task_description + _td = task_description async def _task(*args, task_description=_td, **kwargs) -> str: return _cmd.format_map(kwargs) @@ -50,9 +52,9 @@ async def _task(*args, task_description=_td, **kwargs) -> str: def _make_python_closure(spec: str, remote_paths: list[str], task_description: dict) -> Callable: - _spec = spec + _spec = spec _paths = list(remote_paths) - _td = task_description + _td = task_description async def _task(*args, task_description=_td, **kwargs): import importlib as _il @@ -73,7 +75,7 @@ async def _task(*args, task_description=_td, **kwargs): def _make_shell_dispatch(cmds: dict[int, str], slot_name: str, task_description: dict) -> Callable: _cmds = dict(cmds) - _td = task_description + _td = task_description async def _dispatch(*args, task_description=_td, **kwargs) -> str: cmd = _cmds[kwargs.get("learner_id", 0)] @@ -88,7 +90,7 @@ def _make_python_dispatch( ) -> Callable: _specs = dict(specs) _paths = list(remote_paths) - _td = task_description + _td = task_description async def _dispatch(*args, task_description=_td, **kwargs): import importlib as _il diff --git a/rose/spec/builder.py b/rose/spec/builder.py index 83964d14..7c158faa 100644 --- a/rose/spec/builder.py +++ b/rose/spec/builder.py @@ -5,7 +5,7 @@ from rose.learner import Learner from .adapters import TaskAdapterFactory -from .schema import WorkflowConfig, TrackingConfig, _REQUIRED_SLOTS +from .schema import _REQUIRED_SLOTS, TrackingConfig, WorkflowConfig # Slots that exist as fields on LearnerConfig; uncertainty is required by uq_active_learner # but is not a LearnerConfig field — filter it out when constructing LearnerConfig. @@ -74,7 +74,7 @@ def _register_flat_tasks(self, learner: Learner, cfg: WorkflowConfig) -> None: def _register_dispatched_tasks(self, learner: Learner, cfg: WorkflowConfig) -> None: required = _REQUIRED_SLOTS[cfg.learner.type] for slot_name in required: - task_defs = [l.tasks[slot_name] for l in cfg.learners] + task_defs = [ld.tasks[slot_name] for ld in cfg.learners] as_exec = task_defs[0].type == "shell" closure = TaskAdapterFactory.make_dispatch_closure(slot_name, task_defs, cfg.remote) getattr(learner, f"{slot_name}_task")(as_executable=as_exec)(closure) @@ -82,9 +82,9 @@ def _register_dispatched_tasks(self, learner: Learner, cfg: WorkflowConfig) -> N def build_learner_config(self): """Single LearnerConfig for sequential learners. - Passed as initial_config to learner.start() so that parameters and - iteration reach task kwargs via the same ROSE-native mechanism used - by the parallel path. Returns None when no parameters are defined. + Passed as initial_config to learner.start() so that parameters and iteration reach task + kwargs via the same ROSE-native mechanism used by the parallel path. Returns None when no + parameters are defined. """ cfg = self.config if not cfg.parameters and not cfg.remote.pythonpath: @@ -93,7 +93,7 @@ def build_learner_config(self): required = _REQUIRED_SLOTS[cfg.learner.type] max_iter = cfg.learner.max_iter - params = dict(cfg.parameters) + params = dict(cfg.parameters) params["pythonpath"] = list(cfg.remote.pythonpath) schedule = {n: TaskConfig(kwargs={**params, "iteration": n}) for n in range(max_iter + 1)} schedule[-1] = TaskConfig(kwargs={**params, "iteration": max_iter}) @@ -118,12 +118,12 @@ def build_learner_configs(self): return None from rose.learner import LearnerConfig, TaskConfig - required = _REQUIRED_SLOTS[cfg.learner.type] - configs = [] - max_iter = cfg.learner.max_iter - params = dict(cfg.parameters) + required = _REQUIRED_SLOTS[cfg.learner.type] + configs = [] + max_iter = cfg.learner.max_iter + params = dict(cfg.parameters) params["pythonpath"] = list(cfg.remote.pythonpath) - lc_slots = required & _LEARNER_CONFIG_SLOTS + lc_slots = required & _LEARNER_CONFIG_SLOTS learner_defs = cfg.learners or [None] * cfg.learner.parallel_learners for i, learner_def in enumerate(learner_defs): # learner_id → dispatch routing key (popped by closure, never reaches user fn) diff --git a/rose/spec/schema.py b/rose/spec/schema.py index 8312a351..99641920 100644 --- a/rose/spec/schema.py +++ b/rose/spec/schema.py @@ -6,19 +6,17 @@ from pydantic import BaseModel, Field, model_validator - - # ── Task definition ─────────────────────────────────────────────────────────── class TaskDef(BaseModel): type: Literal["shell", "python"] - command: str | None = None # required when type=="shell" - function: str | None = None # required when type=="python"; "module:callable" + command: str | None = None # required when type=="shell" + function: str | None = None # required when type=="python"; "module:callable" task_description: dict[str, Any] | None = None # resource hints forwarded to asyncflow backend model_config = {"extra": "forbid"} @model_validator(mode="after") - def _check_fields(self) -> "TaskDef": + def _check_fields(self) -> TaskDef: if self.type == "shell" and not self.command: raise ValueError("shell task requires 'command'") if self.type == "python": @@ -41,11 +39,12 @@ class StopCriterionDef(BaseModel): # ── Learner spec (YAML key: "learner") ─────────────────────────────────────── _REQUIRED_SLOTS: dict[str, frozenset[str]] = { - "sequential_active_learner": frozenset({"simulation", "training", "active_learn"}), - "parallel_active_learner": frozenset({"simulation", "training", "active_learn"}), + "sequential_active_learner": frozenset({"simulation", "training", "active_learn"}), + "parallel_active_learner": frozenset({"simulation", "training", "active_learn"}), "sequential_reinforcement_learner": frozenset({"environment", "update"}), - "uq_active_learner": frozenset({"simulation", "training", "prediction", - "active_learn", "uncertainty"}), + "uq_active_learner": frozenset( + {"simulation", "training", "prediction", "active_learn", "uncertainty"} + ), } _ALL_SLOTS: frozenset[str] = frozenset().union(*_REQUIRED_SLOTS.values()) @@ -53,21 +52,21 @@ class StopCriterionDef(BaseModel): class LearnerSpec(BaseModel): type: str max_iter: int = 0 - parallel_learners: int = 2 # used only when learners: is absent for parallel types + parallel_learners: int = 2 # used only when learners: is absent for parallel types model_config = {"extra": "forbid"} # ── Per-learner definition (parallel learners only) ─────────────────────────── class LearnerDef(BaseModel): - label: str = "" - simulation: TaskDef | None = None - training: TaskDef | None = None + label: str = "" + simulation: TaskDef | None = None + training: TaskDef | None = None active_learn: TaskDef | None = None - environment: TaskDef | None = None - update: TaskDef | None = None - prediction: TaskDef | None = None - uncertainty: TaskDef | None = None + environment: TaskDef | None = None + update: TaskDef | None = None + prediction: TaskDef | None = None + uncertainty: TaskDef | None = None model_config = {"extra": "forbid"} @@ -99,21 +98,21 @@ class TrackingConfig(BaseModel): # ── Top-level spec ──────────────────────────────────────────────────────────── class WorkflowConfig(BaseModel): - learner: LearnerSpec + learner: LearnerSpec # Task slots — explicit fields keep extra="forbid" and enable IDE autocomplete. # Access non-None slots as a dict via the .tasks property. - simulation: TaskDef | None = None - training: TaskDef | None = None - active_learn: TaskDef | None = None - environment: TaskDef | None = None - update: TaskDef | None = None - prediction: TaskDef | None = None - uncertainty: TaskDef | None = None - learners: list[LearnerDef] | None = None + simulation: TaskDef | None = None + training: TaskDef | None = None + active_learn: TaskDef | None = None + environment: TaskDef | None = None + update: TaskDef | None = None + prediction: TaskDef | None = None + uncertainty: TaskDef | None = None + learners: list[LearnerDef] | None = None stop_criterion: StopCriterionDef - parameters: dict[str, Any] = Field(default_factory=dict) - remote: RemoteConfig = Field(default_factory=RemoteConfig) - tracking: TrackingConfig = Field(default_factory=TrackingConfig) + parameters: dict[str, Any] = Field(default_factory=dict) + remote: RemoteConfig = Field(default_factory=RemoteConfig) + tracking: TrackingConfig = Field(default_factory=TrackingConfig) model_config = {"extra": "forbid"} @@ -122,32 +121,31 @@ def tasks(self) -> dict[str, TaskDef]: return {s: getattr(self, s) for s in _ALL_SLOTS if getattr(self, s) is not None} @model_validator(mode="after") - def _validate_task_slots(self) -> "WorkflowConfig": - ltype = self.learner.type + def _validate_task_slots(self) -> WorkflowConfig: + ltype = self.learner.type required = _REQUIRED_SLOTS.get(ltype) if required is None: raise ValueError( - f"Unknown learner type '{ltype}'. " - f"Supported: {sorted(_REQUIRED_SLOTS.keys())}" + f"Unknown learner type '{ltype}'. Supported: {sorted(_REQUIRED_SLOTS.keys())}" ) is_parallel = ltype == "parallel_active_learner" if is_parallel and self.learners is not None: - for l in self.learners: - missing = required - set(l.tasks.keys()) + for ld in self.learners: + missing = required - set(ld.tasks.keys()) if missing: - raise ValueError(f"Learner '{l.label}' missing: {sorted(missing)}") - extra = set(l.tasks.keys()) - required + raise ValueError(f"Learner '{ld.label}' missing: {sorted(missing)}") + extra = set(ld.tasks.keys()) - required if extra: - raise ValueError(f"Learner '{l.label}' unexpected fields: {sorted(extra)}") + raise ValueError(f"Learner '{ld.label}' unexpected fields: {sorted(extra)}") for slot in required: - types = {l.tasks[slot].type for l in self.learners} + types = {ld.tasks[slot].type for ld in self.learners} if len(types) > 1: raise ValueError( f"Slot '{slot}' has mixed types across learners: {types}. " "All learners must use the same type for a given slot." ) - descs = [dict(l.tasks[slot].task_description or {}) for l in self.learners] + descs = [dict(ld.tasks[slot].task_description or {}) for ld in self.learners] if len(descs) > 1 and any(d != descs[0] for d in descs[1:]): raise ValueError( f"Slot '{slot}' has different task_description values across learners. " @@ -166,7 +164,7 @@ def _validate_task_slots(self) -> "WorkflowConfig": return self @model_validator(mode="after") - def _validate_parameters(self) -> "WorkflowConfig": + def _validate_parameters(self) -> WorkflowConfig: conflicts = _RESERVED_PARAMETER_KEYS & set(self.parameters.keys()) if conflicts: raise ValueError( @@ -176,7 +174,7 @@ def _validate_parameters(self) -> "WorkflowConfig": return self @classmethod - def from_yaml(cls, path: str | Path) -> "WorkflowConfig": + def from_yaml(cls, path: str | Path) -> WorkflowConfig: import yaml raw = yaml.safe_load(Path(path).read_text()) diff --git a/tests/integration/spec/helpers.py b/tests/integration/spec/helpers.py index 1d4496c4..45c60d70 100644 --- a/tests/integration/spec/helpers.py +++ b/tests/integration/spec/helpers.py @@ -22,6 +22,7 @@ async def criterion(*args, **kwargs): # ── Parameter-capturing variants ────────────────────────────────────────────── + async def sim_capture(*args, **kwargs): received_kwargs.append(dict(kwargs)) return [1.0, 2.0, 3.0] diff --git a/tests/integration/spec/test_yaml_workflow.py b/tests/integration/spec/test_yaml_workflow.py index d430ba11..7ce139b3 100644 --- a/tests/integration/spec/test_yaml_workflow.py +++ b/tests/integration/spec/test_yaml_workflow.py @@ -1,4 +1,5 @@ """Integration tests for the YAML spec layer — full AL loop, no HPC required.""" + import textwrap from concurrent.futures import ThreadPoolExecutor @@ -10,7 +11,6 @@ from rose.spec.builder import LearnerBuilder from rose.spec.schema import WorkflowConfig - # ── Sequential AL via YAML ──────────────────────────────────────────────────── SEQ_YAML = textwrap.dedent("""\ @@ -63,6 +63,7 @@ async def test_yaml_sequential_workflow(tmp_path): # ── load_spec convenience wrapper ───────────────────────────────────────────── + @pytest.mark.asyncio @pytest.mark.integration async def test_load_spec_returns_workflow_spec(tmp_path): @@ -110,6 +111,7 @@ async def test_load_spec_returns_workflow_spec(tmp_path): @pytest.mark.integration async def test_sequential_workflow_parameters_reach_tasks(tmp_path): import tests.integration.spec.helpers as helpers + helpers.received_kwargs.clear() p = tmp_path / "spec.yaml" @@ -185,6 +187,7 @@ async def test_sequential_workflow_parameters_reach_tasks(tmp_path): @pytest.mark.integration async def test_parallel_workflow_parameters_reach_tasks(tmp_path): import tests.integration.spec.helpers as helpers + helpers.received_kwargs.clear() p = tmp_path / "spec.yaml" diff --git a/tests/unit/spec/test_adapters.py b/tests/unit/spec/test_adapters.py index 4de803bb..c6608c27 100644 --- a/tests/unit/spec/test_adapters.py +++ b/tests/unit/spec/test_adapters.py @@ -1,9 +1,8 @@ """Unit tests for rose.spec.adapters — closure factories, no HPC needed.""" + import asyncio import inspect -import pytest - from rose.spec.adapters import TaskAdapterFactory from rose.spec.schema import RemoteConfig, TaskDef @@ -14,6 +13,7 @@ def _remote() -> RemoteConfig: # ── Shell closure ───────────────────────────────────────────────────────────── + def test_shell_closure_returns_command(): td = TaskDef(type="shell", command="echo hello") fn = TaskAdapterFactory.make_closure(td, _remote()) @@ -21,7 +21,6 @@ def test_shell_closure_returns_command(): assert result == "echo hello" - def test_shell_as_executable_true(): td = TaskDef(type="shell", command="x") assert TaskAdapterFactory.as_executable(td) is True @@ -29,6 +28,7 @@ def test_shell_as_executable_true(): # ── Python closure ──────────────────────────────────────────────────────────── + def test_python_closure_calls_sync_function(): td = TaskDef(type="python", function="os.path:join") fn = TaskAdapterFactory.make_closure(td, _remote()) @@ -41,9 +41,11 @@ async def _async_fn(x): return x * 2 import types + mod = types.ModuleType("_test_async_mod") mod.double = _async_fn import sys + sys.modules["_test_async_mod"] = mod td = TaskDef(type="python", function="_test_async_mod:double") @@ -69,6 +71,7 @@ def test_python_closure_injects_remote_path(tmp_path): # ── Shell dispatch ──────────────────────────────────────────────────────────── + def test_shell_dispatch_routes_by_learner_id(): tds = [ TaskDef(type="shell", command="cmd_0"), @@ -87,6 +90,7 @@ def test_shell_dispatch_defaults_to_zero(): # ── Python dispatch ─────────────────────────────────────────────────────────── + def test_python_dispatch_routes_by_learner_id(): tds = [ TaskDef(type="python", function="os.path:basename"), @@ -99,6 +103,7 @@ def test_python_dispatch_routes_by_learner_id(): # ── task_description default kwarg injection ────────────────────────────────── + def test_shell_closure_task_description_in_signature(): td = TaskDef(type="shell", command="echo hi", task_description={"cpu_count": 4}) fn = TaskAdapterFactory.make_closure(td, _remote()) @@ -123,6 +128,7 @@ def test_python_closure_task_description_in_signature(): # ── Shell command interpolation ─────────────────────────────────────────────── + def test_shell_closure_interpolates_kwargs(): td = TaskDef(type="shell", command="python sim.py --dataset {dataset} --n {n}") fn = TaskAdapterFactory.make_closure(td, _remote()) @@ -143,10 +149,13 @@ def test_shell_dispatch_interpolates_kwargs(): TaskDef(type="shell", command="python sim.py --model mlp --dataset {dataset}"), ] fn = TaskAdapterFactory.make_dispatch_closure("simulation", tds, _remote()) - assert asyncio.run(fn(learner_id=0, dataset="m3dc1")) == \ - "python sim.py --model rf --dataset m3dc1" - assert asyncio.run(fn(learner_id=1, dataset="test_ds")) == \ - "python sim.py --model mlp --dataset test_ds" + assert ( + asyncio.run(fn(learner_id=0, dataset="m3dc1")) == "python sim.py --model rf --dataset m3dc1" + ) + assert ( + asyncio.run(fn(learner_id=1, dataset="test_ds")) + == "python sim.py --model mlp --dataset test_ds" + ) def test_dispatch_closure_task_description_uses_first_candidate(): diff --git a/tests/unit/spec/test_builder.py b/tests/unit/spec/test_builder.py index e71b78de..9a893293 100644 --- a/tests/unit/spec/test_builder.py +++ b/tests/unit/spec/test_builder.py @@ -1,12 +1,10 @@ """Unit tests for rose.spec.builder — LearnerConfig construction, no HPC needed.""" + import textwrap from unittest.mock import MagicMock -import pytest - from rose.spec.schema import WorkflowConfig - SEQ_WITH_PARAMS = textwrap.dedent("""\ learner: type: sequential_active_learner @@ -95,11 +93,13 @@ def _make_builder(yaml_text, tmp_path): p.write_text(yaml_text) cfg = WorkflowConfig.from_yaml(p) from rose.spec.builder import LearnerBuilder + return LearnerBuilder(cfg, MagicMock()) # ── build_learner_config (sequential path) ──────────────────────────────────── + def test_build_learner_config_no_parameters_returns_none(tmp_path): builder = _make_builder(SEQ_NO_PARAMS, tmp_path) assert builder.build_learner_config() is None @@ -107,6 +107,7 @@ def test_build_learner_config_no_parameters_returns_none(tmp_path): def test_build_learner_config_returns_learner_config(tmp_path): from rose.learner import LearnerConfig + builder = _make_builder(SEQ_WITH_PARAMS, tmp_path) lc = builder.build_learner_config() assert lc is not None @@ -136,6 +137,7 @@ def test_build_learner_config_criterion_schedule_included(tmp_path): # ── build_learner_configs (parallel path) ──────────────────────────────────── + def test_build_learner_configs_no_learners_returns_none(tmp_path): builder = _make_builder(SEQ_WITH_PARAMS, tmp_path) assert builder.build_learner_configs() is None @@ -143,6 +145,7 @@ def test_build_learner_configs_no_learners_returns_none(tmp_path): def test_build_learner_configs_with_parameters(tmp_path): from rose.learner import LearnerConfig + builder = _make_builder(PAR_WITH_PARAMS, tmp_path) lcs = builder.build_learner_configs() assert lcs is not None @@ -202,6 +205,7 @@ def test_build_learner_configs_shared_tasks_with_parameters_falls_back_to_parall tmp_path, ): from rose.learner import LearnerConfig + builder = _make_builder(PAR_SHARED_WITH_PARAMS, tmp_path) lcs = builder.build_learner_configs() assert lcs is not None diff --git a/tests/unit/spec/test_schema.py b/tests/unit/spec/test_schema.py index 71e7c7ff..690adbf3 100644 --- a/tests/unit/spec/test_schema.py +++ b/tests/unit/spec/test_schema.py @@ -1,13 +1,15 @@ """Unit tests for rose.spec.schema — YAML validation without any HPC machinery.""" + import textwrap import pytest +from pydantic import ValidationError -from rose.spec.schema import WorkflowConfig, TaskDef - +from rose.spec.schema import TaskDef, WorkflowConfig # ── TaskDef ─────────────────────────────────────────────────────────────────── + def test_taskdef_shell_valid(): t = TaskDef(type="shell", command="python sim.py") assert t.type == "shell" @@ -35,12 +37,14 @@ def test_taskdef_python_bad_syntax(): def test_taskdef_extra_field_rejected(): - with pytest.raises(Exception): + with pytest.raises(ValidationError): TaskDef(type="shell", command="x", unknown_field="y") def test_taskdef_task_description_roundtrip(): - t = TaskDef(type="shell", command="python sim.py", task_description={"cpu_count": 4, "gpu_count": 1}) + t = TaskDef( + type="shell", command="python sim.py", task_description={"cpu_count": 4, "gpu_count": 1} + ) assert t.task_description == {"cpu_count": 4, "gpu_count": 1} @@ -119,7 +123,7 @@ def test_sequential_spec_roundtrip(tmp_path): assert cfg.stop_criterion.metric == "mse" assert cfg.tasks == { "simulation": cfg.simulation, - "training": cfg.training, + "training": cfg.training, "active_learn": cfg.active_learn, } @@ -275,6 +279,7 @@ def test_parallel_learner_missing_slot(tmp_path): # ── WorkflowConfig — task_description consistency across parallel learners ──────── + def test_parallel_learners_identical_task_description_valid(tmp_path): yaml = textwrap.dedent("""\ learner: @@ -318,8 +323,8 @@ def test_parallel_learners_identical_task_description_valid(tmp_path): def test_parallel_learners_nested_task_description_different_key_order_valid(tmp_path): - """Nested dicts with the same keys/values in a different insertion order - must compare equal — dict equality, not string-sorted serialization.""" + """Nested dicts with the same keys/values in a different insertion order must compare equal — + dict equality, not string-sorted serialization.""" yaml = textwrap.dedent("""\ learner: type: parallel_active_learner @@ -603,6 +608,7 @@ def test_parameters_defaults_empty(tmp_path): # ── WorkflowConfig — parallel learners (mixed types) ───────────────────────────── + def test_parallel_learners_mixed_types(tmp_path): yaml = textwrap.dedent("""\ learner: diff --git a/tests/unit/spec/test_workflow_spec.py b/tests/unit/spec/test_workflow_spec.py index fcc76a37..a9e7fab8 100644 --- a/tests/unit/spec/test_workflow_spec.py +++ b/tests/unit/spec/test_workflow_spec.py @@ -1,10 +1,10 @@ """Unit tests for WorkflowSpec — load_spec and workflow_with().""" + import textwrap import pytest from rose.spec import WorkflowSpec, load_spec -from rose.spec.schema import WorkflowConfig SEQ_YAML = textwrap.dedent("""\ learner: @@ -44,46 +44,49 @@ def _make_spec(tmp_path, yaml=SEQ_YAML): # ── workflow_with: learner field override ───────────────────────────────────── + def test_workflow_with_max_iter(tmp_path): spec = _make_spec(tmp_path) - new = spec.workflow_with(max_iter=2) + new = spec.workflow_with(max_iter=2) assert new.config.learner.max_iter == 2 assert spec.config.learner.max_iter == 5 # original unchanged def test_workflow_with_does_not_mutate_original(tmp_path): spec = _make_spec(tmp_path) - _ = spec.workflow_with(max_iter=1) + _ = spec.workflow_with(max_iter=1) assert spec.config.learner.max_iter == 5 # ── workflow_with: parameters merge ────────────────────────────────────────── + def test_workflow_with_parameters_merges(tmp_path): spec = _make_spec(tmp_path) - new = spec.workflow_with(parameters={"dataset": "test_ds"}) + new = spec.workflow_with(parameters={"dataset": "test_ds"}) assert new.config.parameters["dataset"] == "test_ds" assert new.config.parameters["scale"] == 1.0 # existing key preserved def test_workflow_with_parameters_adds_new_key(tmp_path): spec = _make_spec(tmp_path) - new = spec.workflow_with(parameters={"lr": 0.001}) + new = spec.workflow_with(parameters={"lr": 0.001}) assert new.config.parameters["lr"] == 0.001 assert new.config.parameters["dataset"] == "base_ds" def test_workflow_with_parameters_does_not_mutate_original(tmp_path): spec = _make_spec(tmp_path) - _ = spec.workflow_with(parameters={"dataset": "other"}) + _ = spec.workflow_with(parameters={"dataset": "other"}) assert spec.config.parameters["dataset"] == "base_ds" # ── workflow_with: combined overrides ───────────────────────────────────────── + def test_workflow_with_combined(tmp_path): spec = _make_spec(tmp_path) - new = spec.workflow_with(max_iter=3, parameters={"dataset": "test_ds", "scale": 2.0}) + new = spec.workflow_with(max_iter=3, parameters={"dataset": "test_ds", "scale": 2.0}) assert new.config.learner.max_iter == 3 assert new.config.parameters["dataset"] == "test_ds" assert new.config.parameters["scale"] == 2.0 @@ -91,6 +94,7 @@ def test_workflow_with_combined(tmp_path): # ── workflow_with: unknown field raises ─────────────────────────────────────── + def test_workflow_with_unknown_field_raises(tmp_path): spec = _make_spec(tmp_path) with pytest.raises(ValueError, match="unknown spec field"): @@ -99,9 +103,10 @@ def test_workflow_with_unknown_field_raises(tmp_path): # ── workflow_with: returns WorkflowSpec with callable .workflow ─────────────── + def test_workflow_with_returns_workflow_spec(tmp_path): spec = _make_spec(tmp_path) - new = spec.workflow_with(max_iter=1) + new = spec.workflow_with(max_iter=1) assert isinstance(new, WorkflowSpec) assert callable(new.workflow) From 8c30b1ee34c3abb2a5ca3637ff7c996b11a373b7 Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Sun, 21 Jun 2026 01:14:11 +0200 Subject: [PATCH 09/13] adding missing init for tests --- tests/__init__.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/__init__.py diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 00000000..e69de29b From 7cce508da946df80844f25f08ce0996c375b808a Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Mon, 6 Jul 2026 16:35:30 +0200 Subject: [PATCH 10/13] update rose with new orbit --- docs/user-guide/rose-aas.md | 8 ++++---- rose/spec/__init__.py | 4 ++-- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/docs/user-guide/rose-aas.md b/docs/user-guide/rose-aas.md index ff98cb3d..3d24f575 100644 --- a/docs/user-guide/rose-aas.md +++ b/docs/user-guide/rose-aas.md @@ -27,7 +27,7 @@ A ROSE learner doesn't run anything itself — it submits tasks through whicheve | Backend | Where the orchestration loop runs | Where tasks execute | Documented at | |---|---|---|---| | `DragonExecutionBackendV3` (RHAPSODY) | Wherever your script runs — it submits a pilot job and blocks inside it | Inside the pilot job it just submitted | [Target Resources](target-resources.md) | -| Edge backend (RHAPSODY, `bridge_url` + `edge_name`) | Wherever your script runs — does **not** need to be the HPC machine | Inside a separate, already-running edge agent — possibly on a different machine entirely | this page | +| Orbit backend (RHAPSODY, `bridge_url` + `endpoint_name`) | Wherever your script runs — does **not** need to be the HPC machine | Inside a separate, already-running orbit endpoint — possibly on a different machine entirely | this page | The first model is "submit a job, then run my loop inside it." The second is "run my loop here; dispatch tasks to a job running somewhere else." That second model is what makes ROSE service-like: the loop's control plane (your `learner.start()` call, deciding when to stop, talking to MLflow/ClearML) is decoupled from the compute plane (the HPC allocation actually executing `simulate`/`train`). @@ -35,10 +35,10 @@ The first model is "submit a job, then run my loop inside it." The second is "ru ## How ROSE operates within a job -The edge backend connects two things over a bridge: +The orbit backend connects two things over a bridge: -1. **An edge agent**, running inside an HPC job allocation. The job itself is requested through whatever your site normally uses to get an allocation — it doesn't have to be requested by ROSE. Once the job starts, the edge agent comes up inside it and stays alive for the allocation's lifetime, ready to accept task descriptions. -2. **Your orchestration process**, running anywhere — your laptop, a long-lived service host, a CI runner. It builds the learner from your spec (`LearnerBuilder(cfg, asyncflow).build()`), starts the loop (`learner.start(...)`), and for every `simulation`/`training`/`active_learn`/... task it submits, the asyncflow engine forwards that task description over the bridge to the edge, waits for the result, and resumes the loop. +1. **An orbit endpoint**, running inside an HPC job allocation. The job itself is requested through whatever your site normally uses to get an allocation — it doesn't have to be requested by ROSE. Once the job starts, the endpoint comes up inside it and stays alive for the allocation's lifetime, ready to accept task descriptions. +2. **Your orchestration process**, running anywhere — your laptop, a long-lived service host, a CI runner. It builds the learner from your spec (`LearnerBuilder(cfg, asyncflow).build()`), starts the loop (`learner.start(...)`), and for every `simulation`/`training`/`active_learn`/... task it submits, the asyncflow engine forwards that task description over the bridge to the endpoint, waits for the result, and resumes the loop. The practical effect: you submit a ROSE workflow to HPC without an interactive session on the cluster, and without your laptop needing to stay connected to the scheduler — only to the bridge. The job allocation is what's expensive and scheduler-queued; your control process is cheap and can be restarted independently of it. Stop-criterion checks, `set_next_config()` decisions, and tracking calls all happen on your side of the bridge, on every iteration, with no per-iteration job resubmission. diff --git a/rose/spec/__init__.py b/rose/spec/__init__.py index 057f2972..75a5ee72 100644 --- a/rose/spec/__init__.py +++ b/rose/spec/__init__.py @@ -111,11 +111,11 @@ def workflow_with(self, **overrides: Any) -> WorkflowSpec: def workflow(self) -> Callable[..., Coroutine[Any, Any, None]]: cfg = self.config - async def _workflow(bridge_url: str, edge_name: str) -> None: + async def _workflow(bridge_url: str, endpoint_name: str) -> None: import rhapsody from radical.asyncflow import WorkflowEngine - engine = await rhapsody.get_backend("edge", bridge_url=bridge_url, edge_name=edge_name) + engine = await rhapsody.get_backend("orbit", bridge_url=bridge_url, endpoint_name=endpoint_name) asyncflow = await WorkflowEngine.create(engine) builder = LearnerBuilder(cfg, asyncflow) From 85f2f17bf76b11cf20d7bef5f8d2cfc1d5f14bde Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Mon, 13 Jul 2026 18:56:21 +0200 Subject: [PATCH 11/13] leftovers --- docs/user-guide/spec-api.md | 26 ++++++++++++++ rose/spec/__init__.py | 4 ++- rose/spec/schema.py | 1 + tests/unit/spec/test_schema.py | 63 ++++++++++++++++++++++++++++++++++ 4 files changed, 93 insertions(+), 1 deletion(-) diff --git a/docs/user-guide/spec-api.md b/docs/user-guide/spec-api.md index 0029ca7c..1fd9dd4f 100644 --- a/docs/user-guide/spec-api.md +++ b/docs/user-guide/spec-api.md @@ -557,8 +557,11 @@ remote: pythonpath: - /path/to/my/task/modules - /path/to/shared/utilities + backends: [dragon_v3] # optional — default shown; use [concurrent] for CPU-only runs ``` +### `remote.pythonpath` + `remote.pythonpath` entries are added to `sys.path` on the remote worker before any task module is imported. They are also injected into every task's `kwargs["pythonpath"]` (as a list), so task functions can construct file paths without duplicating the value in `parameters:`: ```python @@ -571,6 +574,28 @@ def train(sim_result, **kwargs): `remote.pythonpath` is the single edit point for the remote worker path — changing it updates both the import path and the kwarg. +### `remote.backends` + +`remote.backends` selects which [Rhapsody](https://github.com/radical-cybertools/rhapsody) backends are instantiated on the remote orbit endpoint session. The value is a list of backend name strings: + +| Value | Backend class | When to use | +|---|---|---| +| `dragon_v3` | `DragonExecutionBackendV3` | Default; requires Dragon installed on the compute side | +| `concurrent` | `ConcurrentExecutionBackend` | CPU-only runs, no Dragon needed | +| `dragon_v1`, `dragon_v2` | older Dragon backends | Legacy Dragon versions | +| `dask` | `DaskExecutionBackend` | Dask-based execution | +| `radical_pilot` | `RadicalExecutionBackend` | RADICAL-Pilot | + +Example — switch to `concurrent` for a CPU-only endpoint: + +```yaml +remote: + pythonpath: [...] + backends: [concurrent] +``` + +`backends` is consumed only at engine-creation time and is **not** injected into task `kwargs`. + --- ## Tracking @@ -635,6 +660,7 @@ With `validate_imports=True`, every `module:callable` string is resolved in the | `stop_criterion` | object | yes | — | Stopping condition | | `parameters` | dict | no | `{}` | User-defined kwargs injected into all tasks | | `remote.pythonpath` | list[str] | no | `[]` | Paths added to `sys.path` on worker; injected as `pythonpath` kwarg | +| `remote.backends` | list[str] | no | `["dragon_v3"]` | Rhapsody backends requested on the remote orbit session | | `tracking.backend` | string | no | `none` | `mlflow` / `clearml` / `none` | | `tracking.experiment` | string | no | `ROSE-Spec` | Experiment name | | `tracking.run_name` | string | no | `null` | Run label | diff --git a/rose/spec/__init__.py b/rose/spec/__init__.py index 75a5ee72..4f80858f 100644 --- a/rose/spec/__init__.py +++ b/rose/spec/__init__.py @@ -115,7 +115,9 @@ async def _workflow(bridge_url: str, endpoint_name: str) -> None: import rhapsody from radical.asyncflow import WorkflowEngine - engine = await rhapsody.get_backend("orbit", bridge_url=bridge_url, endpoint_name=endpoint_name) + engine = await rhapsody.get_backend("orbit", bridge_url=bridge_url, + endpoint_name=endpoint_name, + backends=cfg.remote.backends) asyncflow = await WorkflowEngine.create(engine) builder = LearnerBuilder(cfg, asyncflow) diff --git a/rose/spec/schema.py b/rose/spec/schema.py index 99641920..492df021 100644 --- a/rose/spec/schema.py +++ b/rose/spec/schema.py @@ -78,6 +78,7 @@ def tasks(self) -> dict[str, TaskDef]: # ── Remote / tracking ──────────────────────────────────────────────────────── class RemoteConfig(BaseModel): pythonpath: list[str] = [] + backends: list[str] = ["dragon_v3"] model_config = {"extra": "forbid"} diff --git a/tests/unit/spec/test_schema.py b/tests/unit/spec/test_schema.py index 690adbf3..596aeac9 100644 --- a/tests/unit/spec/test_schema.py +++ b/tests/unit/spec/test_schema.py @@ -452,6 +452,69 @@ def test_parallel_learners_one_has_task_description_other_absent_raises(tmp_path WorkflowConfig.from_yaml(p) +def test_remote_backends_custom(tmp_path): + yaml = textwrap.dedent("""\ + learner: + type: sequential_active_learner + max_iter: 1 + simulation: + type: python + function: mymod:sim + training: + type: python + function: mymod:train + active_learn: + type: python + function: mymod:select + stop_criterion: + metric: mse + threshold: 0.1 + evaluator: + type: python + function: mymod:eval + remote: + backends: [concurrent] + """) + p = tmp_path / "cfg.yaml" + p.write_text(yaml) + cfg = WorkflowConfig.from_yaml(p) + assert cfg.remote.backends == ["concurrent"] + + +def test_remote_backends_default(): + from rose.spec.schema import RemoteConfig + assert RemoteConfig().backends == ["dragon_v3"] + + +def test_remote_extra_field_rejected(tmp_path): + yaml = textwrap.dedent("""\ + learner: + type: sequential_active_learner + max_iter: 1 + simulation: + type: python + function: mymod:sim + training: + type: python + function: mymod:train + active_learn: + type: python + function: mymod:select + stop_criterion: + metric: mse + threshold: 0.1 + evaluator: + type: python + function: mymod:eval + remote: + foo: bar + """) + p = tmp_path / "bad.yaml" + p.write_text(yaml) + with pytest.raises(ValidationError): + WorkflowConfig.from_yaml(p) + + def test_parameters_reserved_key_pythonpath(tmp_path): yaml = textwrap.dedent("""\ learner: From a4b3f0a59b854801e6904019f2cf03792dd71201 Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Mon, 13 Jul 2026 22:04:07 +0200 Subject: [PATCH 12/13] rename rose/al,rl,uq to descriptive package names; add rose run CLI MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Rename rose/al → rose/active_learning, rose/rl → rose/reinforcement_learning, rose/uq → rose/uncertainty_quantification so directory names match their content. No breaking changes: rose/__init__.py binds `from . import active_learning as al` and populates sys.modules with every rose.al.* short path, so all existing `from rose.al.active_learner import X` imports continue to work unchanged. - Add rose/cli.py and register the `rose` console script in pyproject.toml. `rose run workflow.yaml --local` runs any ROSE workflow locally using the rhapsody concurrent backend, eliminating the local.py boilerplate file. - Update rose/spec/builder.py deferred imports to the new package paths. - Update docs and README to use the short package-level aliases (from rose.al import X, from rose.rl import X, from rose.uq import X) instead of submodule paths. --- README.md | 2 +- docs/getting-started/dry-run.md | 2 +- docs/user-guide/advanced-rl-workflow.md | 2 +- docs/user-guide/basic-acl-workflow.md | 2 +- docs/user-guide/basic-rl-workflow.md | 2 +- docs/user-guide/experience.md | 4 +- docs/user-guide/spec-api.md | 12 ++-- docs/user-guide/target-resources.md | 4 +- docs/user-guide/uq_based-acl-workflow.md | 2 +- pyproject.toml | 3 + rose/__init__.py | 27 +++++++- rose/active_learning/__init__.py | 8 +++ .../{al => active_learning}/active_learner.py | 0 rose/{al => active_learning}/selector.py | 2 +- rose/al/__init__.py | 8 --- rose/cli.py | 64 +++++++++++++++++++ rose/reinforcement_learning/__init__.py | 8 +++ .../experience.py | 0 .../reinforcement_learner.py | 0 rose/rl/__init__.py | 8 --- rose/spec/builder.py | 8 +-- rose/uncertainty_quantification/__init__.py | 13 ++++ .../uq_active_learner.py | 2 +- .../uq_learner.py | 0 .../uq_scorer.py | 0 rose/uq/__init__.py | 13 ---- 26 files changed, 142 insertions(+), 54 deletions(-) create mode 100644 rose/active_learning/__init__.py rename rose/{al => active_learning}/active_learner.py (100%) rename rose/{al => active_learning}/selector.py (99%) delete mode 100644 rose/al/__init__.py create mode 100644 rose/cli.py create mode 100644 rose/reinforcement_learning/__init__.py rename rose/{rl => reinforcement_learning}/experience.py (100%) rename rose/{rl => reinforcement_learning}/reinforcement_learner.py (100%) delete mode 100644 rose/rl/__init__.py create mode 100644 rose/uncertainty_quantification/__init__.py rename rose/{uq => uncertainty_quantification}/uq_active_learner.py (99%) rename rose/{uq => uncertainty_quantification}/uq_learner.py (100%) rename rose/{uq => uncertainty_quantification}/uq_scorer.py (100%) delete mode 100644 rose/uq/__init__.py diff --git a/README.md b/README.md index 7f0f4b8c..22f95d73 100644 --- a/README.md +++ b/README.md @@ -27,7 +27,7 @@ For tutorials and walkthrough notebooks please check [here](examples) import asyncio from rose.metrics import MEAN_SQUARED_ERROR_MSE -from rose.al.active_learner import SequentialActiveLearner +from rose.al import SequentialActiveLearner from radical.asyncflow import WorkflowEngine from rhapsody.backends import RadicalExecutionBackend diff --git a/docs/getting-started/dry-run.md b/docs/getting-started/dry-run.md index 660061bd..e44e5775 100644 --- a/docs/getting-started/dry-run.md +++ b/docs/getting-started/dry-run.md @@ -12,7 +12,7 @@ import os import sys import asyncio -from rose.al.active_learner import SequentialActiveLearner +from rose.al import SequentialActiveLearner from radical.asyncflow import WorkflowEngine diff --git a/docs/user-guide/advanced-rl-workflow.md b/docs/user-guide/advanced-rl-workflow.md index 6c55920f..66c46fb5 100644 --- a/docs/user-guide/advanced-rl-workflow.md +++ b/docs/user-guide/advanced-rl-workflow.md @@ -22,7 +22,7 @@ Import ROSE parallel RL modules: ```python from radical.asyncflow import WorkflowEngine from rhapsody.backends import DragonExecutionBackendV3 -from rose.rl.reinforcement_learner import SequentialReinforcementLearner +from rose.rl import SequentialReinforcementLearner ``` diff --git a/docs/user-guide/basic-acl-workflow.md b/docs/user-guide/basic-acl-workflow.md index 91b88146..18db108e 100644 --- a/docs/user-guide/basic-acl-workflow.md +++ b/docs/user-guide/basic-acl-workflow.md @@ -3,7 +3,7 @@ Import ROSE main modules: ```python from rose.metrics import MEAN_SQUARED_ERROR_MSE -from rose.al.active_learner import SequentialActiveLearner +from rose.al import SequentialActiveLearner from concurrent.futures import ProcessPoolExecutor diff --git a/docs/user-guide/basic-rl-workflow.md b/docs/user-guide/basic-rl-workflow.md index dbee4a4b..7e1329ca 100644 --- a/docs/user-guide/basic-rl-workflow.md +++ b/docs/user-guide/basic-rl-workflow.md @@ -10,7 +10,7 @@ from radical.asyncflow import WorkflowEngine from rhapsody.backends import ConcurrentExecutionBackend from rose.metrics import GREATER_THAN_THRESHOLD -from rose.rl.reinforcement_learner import SequentialReinforcementLearner +from rose.rl import SequentialReinforcementLearner ``` You can use and setup an HPC engine as we described in our previous [Target Resources](target-resources.md) step: diff --git a/docs/user-guide/experience.md b/docs/user-guide/experience.md index 696e7d02..0dfafece 100644 --- a/docs/user-guide/experience.md +++ b/docs/user-guide/experience.md @@ -14,7 +14,7 @@ The `Experience` class captures a single environment interaction containing: Use the helper function to create experiences during environment interactions: ```python -from rose.rl.experience import Experience, create_experience +from rose.rl import Experience, create_experience experience = create_experience(state, action, reward, next_state, done, info={"step": 42}) ``` @@ -23,7 +23,7 @@ experience = create_experience(state, action, reward, next_state, done, info={"s Experience Banks provide persistent storage and efficient sampling for RL experiences: ```python -from rose.rl.experience import ExperienceBank +from rose.rl import ExperienceBank # Create a bank with size limit bank = ExperienceBank(max_size=10000) diff --git a/docs/user-guide/spec-api.md b/docs/user-guide/spec-api.md index 1fd9dd4f..97b0dca0 100644 --- a/docs/user-guide/spec-api.md +++ b/docs/user-guide/spec-api.md @@ -73,7 +73,7 @@ slot — see the side-by-side example in [Task Types](#task-types) below. **Python API** ```python -from rose.al.active_learner import SequentialActiveLearner +from rose.al import SequentialActiveLearner acl = SequentialActiveLearner(asyncflow) @@ -182,7 +182,7 @@ When all parallel learners run the same task implementations, declare the task s **Python API** ```python -from rose.al.active_learner import ParallelActiveLearner +from rose.al import ParallelActiveLearner acl = ParallelActiveLearner(asyncflow) @@ -580,11 +580,11 @@ def train(sim_result, **kwargs): | Value | Backend class | When to use | |---|---|---| -| `dragon_v3` | `DragonExecutionBackendV3` | Default; requires Dragon installed on the compute side | -| `concurrent` | `ConcurrentExecutionBackend` | CPU-only runs, no Dragon needed | -| `dragon_v1`, `dragon_v2` | older Dragon backends | Legacy Dragon versions | | `dask` | `DaskExecutionBackend` | Dask-based execution | -| `radical_pilot` | `RadicalExecutionBackend` | RADICAL-Pilot | +| `concurrent` | `ConcurrentExecutionBackend` | CPU-only runs, no Dragon needed | +| `dragon` | `DragonExecutionBackendV3` | For highly compute intensive surrogates | + + Example — switch to `concurrent` for a CPU-only endpoint: diff --git a/docs/user-guide/target-resources.md b/docs/user-guide/target-resources.md index 7f4c523e..84e3decd 100644 --- a/docs/user-guide/target-resources.md +++ b/docs/user-guide/target-resources.md @@ -11,7 +11,7 @@ from concurrent.futures import ProcessPoolExecutor from radical.asyncflow import WorkflowEngine from rhapsody.backends import ConcurrentExecutionBackend -from rose.al.active_learner import SequentialActiveLearner +from rose.al import SequentialActiveLearner engine = await ConcurrentExecutionBackend(ProcessPoolExecutor()) @@ -35,7 +35,7 @@ import os from radical.asyncflow import WorkflowEngine from rhapsody.backends import DragonExecutionBackendV3 -from rose.al.active_learner import SequentialActiveLearner +from rose.al import SequentialActiveLearner hpc_engine = await DragonExecutionBackendV3() diff --git a/docs/user-guide/uq_based-acl-workflow.md b/docs/user-guide/uq_based-acl-workflow.md index 449bbce4..e9a8e5e7 100644 --- a/docs/user-guide/uq_based-acl-workflow.md +++ b/docs/user-guide/uq_based-acl-workflow.md @@ -94,7 +94,7 @@ This design allows **parallel training and uncertainty-aware sampling** within A Import and Initialize the UQ Learner ```python -from rose.uq.uq_active_learner import ParallelUQLearner +from rose.uq import ParallelUQLearner uql = ParallelUQLearner(asyncflow) ``` diff --git a/pyproject.toml b/pyproject.toml index f57c4226..f404a4b7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,6 +24,9 @@ dependencies = [ 'rhapsody-py[dragon]; python_version >= "3.10" and python_version <= "3.12"', ] +[project.scripts] +rose = "rose.cli:main" + [project.urls] Homepage = "https://github.com/radical-cybertools/ROSE" Issues = "https://github.com/radical-cybertools/ROSE/issues" diff --git a/rose/__init__.py b/rose/__init__.py index db05a14b..b2d49597 100644 --- a/rose/__init__.py +++ b/rose/__init__.py @@ -1,12 +1,33 @@ -from rose.al import active_learner, selector +import sys as _sys + +from . import active_learning as al +from . import reinforcement_learning as rl +from . import uncertainty_quantification as uq +from .active_learning import active_learner, selector +from .reinforcement_learning import reinforcement_learner +from .uncertainty_quantification import uq_active_learner, uq_learner, uq_scorer from rose.learner import IterationState, Learner, LearnerConfig, TaskConfig from rose.metrics import * # noqa: F403 -from rose.rl import reinforcement_learner from rose.spec import WorkflowSpec, load_spec from rose.tracking import PipelineManifest, TrackerBase -from rose.uq import uq_active_learner, uq_learner, uq_scorer + +_sys.modules['rose.al'] = al +_sys.modules['rose.rl'] = rl +_sys.modules['rose.uq'] = uq +_sys.modules['rose.al.active_learner'] = _sys.modules['rose.active_learning.active_learner'] +_sys.modules['rose.al.selector'] = _sys.modules['rose.active_learning.selector'] +_sys.modules['rose.rl.experience'] = _sys.modules['rose.reinforcement_learning.experience'] +_sys.modules['rose.rl.reinforcement_learner'] = _sys.modules['rose.reinforcement_learning.reinforcement_learner'] +_sys.modules['rose.uq.uq_active_learner'] = _sys.modules['rose.uncertainty_quantification.uq_active_learner'] +_sys.modules['rose.uq.uq_learner'] = _sys.modules['rose.uncertainty_quantification.uq_learner'] +_sys.modules['rose.uq.uq_scorer'] = _sys.modules['rose.uncertainty_quantification.uq_scorer'] +del _sys __all__ = [ + # Short-name subpackage aliases + "al", + "rl", + "uq", # Submodules "active_learner", "selector", diff --git a/rose/active_learning/__init__.py b/rose/active_learning/__init__.py new file mode 100644 index 00000000..47fc07be --- /dev/null +++ b/rose/active_learning/__init__.py @@ -0,0 +1,8 @@ +from rose.active_learning.active_learner import ParallelActiveLearner, SequentialActiveLearner +from rose.active_learning.selector import AlgorithmSelector + +__all__ = [ + "ParallelActiveLearner", + "SequentialActiveLearner", + "AlgorithmSelector", +] diff --git a/rose/al/active_learner.py b/rose/active_learning/active_learner.py similarity index 100% rename from rose/al/active_learner.py rename to rose/active_learning/active_learner.py diff --git a/rose/al/selector.py b/rose/active_learning/selector.py similarity index 99% rename from rose/al/selector.py rename to rose/active_learning/selector.py index bed0073f..7b2c70c3 100644 --- a/rose/al/selector.py +++ b/rose/active_learning/selector.py @@ -97,7 +97,7 @@ def _create_algorithm_learner( Configured SequentialActiveLearner instance. """ # Import here to avoid circular imports - from rose.al import SequentialActiveLearner + from rose.active_learning import SequentialActiveLearner # Create a new sequential learner with the same asyncflow sequential_learner: SequentialActiveLearner = SequentialActiveLearner(self.asyncflow) diff --git a/rose/al/__init__.py b/rose/al/__init__.py deleted file mode 100644 index fd2e905d..00000000 --- a/rose/al/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -from rose.al.active_learner import ParallelActiveLearner, SequentialActiveLearner -from rose.al.selector import AlgorithmSelector - -__all__ = [ - "ParallelActiveLearner", - "SequentialActiveLearner", - "AlgorithmSelector", -] diff --git a/rose/cli.py b/rose/cli.py new file mode 100644 index 00000000..13ace782 --- /dev/null +++ b/rose/cli.py @@ -0,0 +1,64 @@ +#!/usr/bin/env python3 +import argparse +import asyncio +import sys +from pathlib import Path + + +def _run_local(yaml_path: Path, backend: str) -> None: + sys.path.insert(0, str(yaml_path.resolve().parent)) + import rhapsody + from radical.asyncflow import WorkflowEngine + from rose.spec import load_spec + from rose.spec.builder import LearnerBuilder + + async def _main(): + spec = load_spec(yaml_path) + cfg = spec.config + engine = await rhapsody.get_backend(backend) + asyncflow = await WorkflowEngine.create(engine) + builder = LearnerBuilder(cfg, asyncflow) + learner = builder.build() + + start_kwargs: dict = {"max_iter": cfg.learner.max_iter} + if cfg.learner.type == "parallel_active_learner": + lcs = builder.build_learner_configs() + if lcs is not None: + start_kwargs["parallel_learners"] = len(lcs) + start_kwargs["learner_configs"] = lcs + else: + start_kwargs["parallel_learners"] = cfg.learner.parallel_learners + else: + ic = builder.build_learner_config() + if ic is not None: + start_kwargs["initial_config"] = ic + + try: + async for state in learner.start(**start_kwargs): + print(f"[iter {state.iteration}] metric={state.metric_value}", flush=True) + finally: + await asyncflow.shutdown() + + asyncio.run(_main()) + + +def main(): + parser = argparse.ArgumentParser(prog="rose") + sub = parser.add_subparsers(dest="command") + + run_p = sub.add_parser("run", help="Execute a ROSE workflow YAML") + run_p.add_argument("yaml", type=Path, help="Path to workflow YAML") + run_p.add_argument("--local", action="store_true", + help="Run locally using rhapsody concurrent backend") + run_p.add_argument("--backend", default="concurrent", + help="Rhapsody backend name for --local (default: concurrent)") + + args = parser.parse_args() + + if args.command == "run": + if not args.local: + parser.error("specify a mode: --local (--orbit coming soon)") + _run_local(args.yaml, args.backend) + else: + parser.print_help() + sys.exit(1) diff --git a/rose/reinforcement_learning/__init__.py b/rose/reinforcement_learning/__init__.py new file mode 100644 index 00000000..db812ed1 --- /dev/null +++ b/rose/reinforcement_learning/__init__.py @@ -0,0 +1,8 @@ +from rose.reinforcement_learning.experience import Experience, ExperienceBank +from rose.reinforcement_learning.reinforcement_learner import ReinforcementLearner + +__all__ = [ + "Experience", + "ExperienceBank", + "ReinforcementLearner", +] diff --git a/rose/rl/experience.py b/rose/reinforcement_learning/experience.py similarity index 100% rename from rose/rl/experience.py rename to rose/reinforcement_learning/experience.py diff --git a/rose/rl/reinforcement_learner.py b/rose/reinforcement_learning/reinforcement_learner.py similarity index 100% rename from rose/rl/reinforcement_learner.py rename to rose/reinforcement_learning/reinforcement_learner.py diff --git a/rose/rl/__init__.py b/rose/rl/__init__.py deleted file mode 100644 index fc7b80ad..00000000 --- a/rose/rl/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -from rose.rl.experience import Experience, ExperienceBank -from rose.rl.reinforcement_learner import ReinforcementLearner - -__all__ = [ - "Experience", - "ExperienceBank", - "ReinforcementLearner", -] diff --git a/rose/spec/builder.py b/rose/spec/builder.py index 7c158faa..471b5fae 100644 --- a/rose/spec/builder.py +++ b/rose/spec/builder.py @@ -16,19 +16,19 @@ def _get_learner_class(learner_type: str): if learner_type == "sequential_active_learner": - from rose.al.active_learner import SequentialActiveLearner + from rose.active_learning.active_learner import SequentialActiveLearner return SequentialActiveLearner if learner_type == "parallel_active_learner": - from rose.al.active_learner import ParallelActiveLearner + from rose.active_learning.active_learner import ParallelActiveLearner return ParallelActiveLearner if learner_type == "sequential_reinforcement_learner": - from rose.rl.reinforcement_learner import SequentialReinforcementLearner + from rose.reinforcement_learning.reinforcement_learner import SequentialReinforcementLearner return SequentialReinforcementLearner if learner_type == "uq_active_learner": - from rose.uq.uq_active_learner import SeqUQLearner + from rose.uncertainty_quantification.uq_active_learner import SeqUQLearner return SeqUQLearner raise ValueError(f"No learner class registered for type '{learner_type}'") diff --git a/rose/uncertainty_quantification/__init__.py b/rose/uncertainty_quantification/__init__.py new file mode 100644 index 00000000..79dada40 --- /dev/null +++ b/rose/uncertainty_quantification/__init__.py @@ -0,0 +1,13 @@ +from rose.uncertainty_quantification.uq_active_learner import ParallelUQLearner, SeqUQLearner +from rose.uncertainty_quantification.uq_learner import UQLearner, UQLearnerConfig +from rose.uncertainty_quantification.uq_scorer import UQ_REGISTRY, UQScorer, register_uq + +__all__ = [ + "UQLearner", + "ParallelUQLearner", + "SeqUQLearner", + "UQScorer", + "register_uq", + "UQ_REGISTRY", + "UQLearnerConfig", +] diff --git a/rose/uq/uq_active_learner.py b/rose/uncertainty_quantification/uq_active_learner.py similarity index 99% rename from rose/uq/uq_active_learner.py rename to rose/uncertainty_quantification/uq_active_learner.py index fb9f5e93..8e4e78d1 100644 --- a/rose/uq/uq_active_learner.py +++ b/rose/uncertainty_quantification/uq_active_learner.py @@ -8,7 +8,7 @@ from radical.asyncflow import WorkflowEngine -from rose.uq.uq_learner import UQLearner, UQLearnerConfig +from rose.uncertainty_quantification.uq_learner import UQLearner, UQLearnerConfig from ..learner import IterationState, TaskConfig, _stream_parallel diff --git a/rose/uq/uq_learner.py b/rose/uncertainty_quantification/uq_learner.py similarity index 100% rename from rose/uq/uq_learner.py rename to rose/uncertainty_quantification/uq_learner.py diff --git a/rose/uq/uq_scorer.py b/rose/uncertainty_quantification/uq_scorer.py similarity index 100% rename from rose/uq/uq_scorer.py rename to rose/uncertainty_quantification/uq_scorer.py diff --git a/rose/uq/__init__.py b/rose/uq/__init__.py deleted file mode 100644 index e78cc8fa..00000000 --- a/rose/uq/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -from rose.uq.uq_active_learner import ParallelUQLearner, SeqUQLearner -from rose.uq.uq_learner import UQLearner, UQLearnerConfig -from rose.uq.uq_scorer import UQ_REGISTRY, UQScorer, register_uq - -__all__ = [ - "UQLearner", - "ParallelUQLearner", - "SeqUQLearner", - "UQScorer", - "register_uq", - "UQ_REGISTRY", - "UQLearnerConfig", -] From 2556dbd78bcec4ec04397a70df601f5bb3cf1f6d Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Mon, 13 Jul 2026 23:33:45 +0200 Subject: [PATCH 13/13] pre-commit --- rose/__init__.py | 33 +++++++++++-------- rose/cli.py | 27 ++++++++------- rose/spec/__init__.py | 9 +++-- rose/spec/schema.py | 2 +- .../integration/test_run_parallel_learner.py | 2 +- tests/integration/test_run_rl_par_learner.py | 2 +- tests/integration/test_run_rl_seq_learner.py | 2 +- .../test_run_sequential_learner.py | 2 +- tests/integration/test_run_uq_learner.py | 4 +-- tests/unit/spec/test_schema.py | 1 + 10 files changed, 49 insertions(+), 35 deletions(-) diff --git a/rose/__init__.py b/rose/__init__.py index b2d49597..8a48e1d6 100644 --- a/rose/__init__.py +++ b/rose/__init__.py @@ -1,26 +1,31 @@ import sys as _sys +from rose.learner import IterationState, Learner, LearnerConfig, TaskConfig +from rose.metrics import * # noqa: F403 +from rose.spec import WorkflowSpec, load_spec +from rose.tracking import PipelineManifest, TrackerBase + from . import active_learning as al from . import reinforcement_learning as rl from . import uncertainty_quantification as uq from .active_learning import active_learner, selector from .reinforcement_learning import reinforcement_learner from .uncertainty_quantification import uq_active_learner, uq_learner, uq_scorer -from rose.learner import IterationState, Learner, LearnerConfig, TaskConfig -from rose.metrics import * # noqa: F403 -from rose.spec import WorkflowSpec, load_spec -from rose.tracking import PipelineManifest, TrackerBase -_sys.modules['rose.al'] = al -_sys.modules['rose.rl'] = rl -_sys.modules['rose.uq'] = uq -_sys.modules['rose.al.active_learner'] = _sys.modules['rose.active_learning.active_learner'] -_sys.modules['rose.al.selector'] = _sys.modules['rose.active_learning.selector'] -_sys.modules['rose.rl.experience'] = _sys.modules['rose.reinforcement_learning.experience'] -_sys.modules['rose.rl.reinforcement_learner'] = _sys.modules['rose.reinforcement_learning.reinforcement_learner'] -_sys.modules['rose.uq.uq_active_learner'] = _sys.modules['rose.uncertainty_quantification.uq_active_learner'] -_sys.modules['rose.uq.uq_learner'] = _sys.modules['rose.uncertainty_quantification.uq_learner'] -_sys.modules['rose.uq.uq_scorer'] = _sys.modules['rose.uncertainty_quantification.uq_scorer'] +_sys.modules["rose.al"] = al +_sys.modules["rose.rl"] = rl +_sys.modules["rose.uq"] = uq +_sys.modules["rose.al.active_learner"] = _sys.modules["rose.active_learning.active_learner"] +_sys.modules["rose.al.selector"] = _sys.modules["rose.active_learning.selector"] +_sys.modules["rose.rl.experience"] = _sys.modules["rose.reinforcement_learning.experience"] +_sys.modules["rose.rl.reinforcement_learner"] = _sys.modules[ + "rose.reinforcement_learning.reinforcement_learner" +] +_sys.modules["rose.uq.uq_active_learner"] = _sys.modules[ + "rose.uncertainty_quantification.uq_active_learner" +] +_sys.modules["rose.uq.uq_learner"] = _sys.modules["rose.uncertainty_quantification.uq_learner"] +_sys.modules["rose.uq.uq_scorer"] = _sys.modules["rose.uncertainty_quantification.uq_scorer"] del _sys __all__ = [ diff --git a/rose/cli.py b/rose/cli.py index 13ace782..412aa094 100644 --- a/rose/cli.py +++ b/rose/cli.py @@ -9,23 +9,24 @@ def _run_local(yaml_path: Path, backend: str) -> None: sys.path.insert(0, str(yaml_path.resolve().parent)) import rhapsody from radical.asyncflow import WorkflowEngine + from rose.spec import load_spec from rose.spec.builder import LearnerBuilder async def _main(): - spec = load_spec(yaml_path) - cfg = spec.config - engine = await rhapsody.get_backend(backend) + spec = load_spec(yaml_path) + cfg = spec.config + engine = await rhapsody.get_backend(backend) asyncflow = await WorkflowEngine.create(engine) - builder = LearnerBuilder(cfg, asyncflow) - learner = builder.build() + builder = LearnerBuilder(cfg, asyncflow) + learner = builder.build() start_kwargs: dict = {"max_iter": cfg.learner.max_iter} if cfg.learner.type == "parallel_active_learner": lcs = builder.build_learner_configs() if lcs is not None: start_kwargs["parallel_learners"] = len(lcs) - start_kwargs["learner_configs"] = lcs + start_kwargs["learner_configs"] = lcs else: start_kwargs["parallel_learners"] = cfg.learner.parallel_learners else: @@ -44,14 +45,18 @@ async def _main(): def main(): parser = argparse.ArgumentParser(prog="rose") - sub = parser.add_subparsers(dest="command") + sub = parser.add_subparsers(dest="command") run_p = sub.add_parser("run", help="Execute a ROSE workflow YAML") run_p.add_argument("yaml", type=Path, help="Path to workflow YAML") - run_p.add_argument("--local", action="store_true", - help="Run locally using rhapsody concurrent backend") - run_p.add_argument("--backend", default="concurrent", - help="Rhapsody backend name for --local (default: concurrent)") + run_p.add_argument( + "--local", action="store_true", help="Run locally using rhapsody concurrent backend" + ) + run_p.add_argument( + "--backend", + default="concurrent", + help="Rhapsody backend name for --local (default: concurrent)", + ) args = parser.parse_args() diff --git a/rose/spec/__init__.py b/rose/spec/__init__.py index 4f80858f..39146f91 100644 --- a/rose/spec/__init__.py +++ b/rose/spec/__init__.py @@ -115,9 +115,12 @@ async def _workflow(bridge_url: str, endpoint_name: str) -> None: import rhapsody from radical.asyncflow import WorkflowEngine - engine = await rhapsody.get_backend("orbit", bridge_url=bridge_url, - endpoint_name=endpoint_name, - backends=cfg.remote.backends) + engine = await rhapsody.get_backend( + "orbit", + bridge_url=bridge_url, + endpoint_name=endpoint_name, + backends=cfg.remote.backends, + ) asyncflow = await WorkflowEngine.create(engine) builder = LearnerBuilder(cfg, asyncflow) diff --git a/rose/spec/schema.py b/rose/spec/schema.py index 492df021..d8da9cb2 100644 --- a/rose/spec/schema.py +++ b/rose/spec/schema.py @@ -78,7 +78,7 @@ def tasks(self) -> dict[str, TaskDef]: # ── Remote / tracking ──────────────────────────────────────────────────────── class RemoteConfig(BaseModel): pythonpath: list[str] = [] - backends: list[str] = ["dragon_v3"] + backends: list[str] = ["dragon_v3"] model_config = {"extra": "forbid"} diff --git a/tests/integration/test_run_parallel_learner.py b/tests/integration/test_run_parallel_learner.py index c180ccc9..8e7c76d1 100644 --- a/tests/integration/test_run_parallel_learner.py +++ b/tests/integration/test_run_parallel_learner.py @@ -3,8 +3,8 @@ import pytest from radical.asyncflow import WorkflowEngine from rhapsody.backends import ConcurrentExecutionBackend - from rose.al.active_learner import ParallelActiveLearner + from rose.metrics import MEAN_SQUARED_ERROR_MSE diff --git a/tests/integration/test_run_rl_par_learner.py b/tests/integration/test_run_rl_par_learner.py index ad244013..0abc4bc6 100644 --- a/tests/integration/test_run_rl_par_learner.py +++ b/tests/integration/test_run_rl_par_learner.py @@ -3,9 +3,9 @@ import pytest from radical.asyncflow import WorkflowEngine from rhapsody.backends import ConcurrentExecutionBackend +from rose.rl.reinforcement_learner import ParallelReinforcementLearner from rose.metrics import GREATER_THAN_THRESHOLD -from rose.rl.reinforcement_learner import ParallelReinforcementLearner @pytest.mark.asyncio diff --git a/tests/integration/test_run_rl_seq_learner.py b/tests/integration/test_run_rl_seq_learner.py index f61f4495..16372a63 100644 --- a/tests/integration/test_run_rl_seq_learner.py +++ b/tests/integration/test_run_rl_seq_learner.py @@ -3,9 +3,9 @@ import pytest from radical.asyncflow import WorkflowEngine from rhapsody.backends import ConcurrentExecutionBackend +from rose.rl.reinforcement_learner import SequentialReinforcementLearner from rose.metrics import GREATER_THAN_THRESHOLD -from rose.rl.reinforcement_learner import SequentialReinforcementLearner @pytest.mark.asyncio diff --git a/tests/integration/test_run_sequential_learner.py b/tests/integration/test_run_sequential_learner.py index d7a75efd..62c2029e 100644 --- a/tests/integration/test_run_sequential_learner.py +++ b/tests/integration/test_run_sequential_learner.py @@ -3,8 +3,8 @@ import pytest from radical.asyncflow import WorkflowEngine from rhapsody.backends import ConcurrentExecutionBackend - from rose.al.active_learner import SequentialActiveLearner + from rose.metrics import MEAN_SQUARED_ERROR_MSE diff --git a/tests/integration/test_run_uq_learner.py b/tests/integration/test_run_uq_learner.py index 86533e3f..8bcf63d4 100644 --- a/tests/integration/test_run_uq_learner.py +++ b/tests/integration/test_run_uq_learner.py @@ -5,11 +5,11 @@ import pytest from radical.asyncflow import WorkflowEngine from rhapsody.backends import ConcurrentExecutionBackend - -from rose.metrics import MEAN_SQUARED_ERROR_MSE, PREDICTIVE_ENTROPY from rose.uq import UQ_REGISTRY, register_uq from rose.uq.uq_active_learner import ParallelUQLearner, SeqUQLearner +from rose.metrics import MEAN_SQUARED_ERROR_MSE, PREDICTIVE_ENTROPY + @pytest.mark.asyncio async def test_active_learning_pipeline_functions(): diff --git a/tests/unit/spec/test_schema.py b/tests/unit/spec/test_schema.py index 596aeac9..beb3e9fb 100644 --- a/tests/unit/spec/test_schema.py +++ b/tests/unit/spec/test_schema.py @@ -483,6 +483,7 @@ def test_remote_backends_custom(tmp_path): def test_remote_backends_default(): from rose.spec.schema import RemoteConfig + assert RemoteConfig().backends == ["dragon_v3"]