From cae8e47da5b8f3bc99a023ce95cc6905c1068f4e Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Fri, 12 Jun 2026 19:59:59 +0200 Subject: [PATCH 01/19] 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/19] 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/19] 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/19] 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/19] 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/19] 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/19] 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/19] 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/19] 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/19] 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/19] 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/19] 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/19] 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"] From 3c60d65ca77cd9b1d7a394050333e24832909f59 Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Wed, 19 Aug 2026 00:12:27 +0200 Subject: [PATCH 14/19] implement `rose run --remote` (spec-driven remote execution via ORBIT) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ports amsc/use-cases/service_utils.py's discover/launch/teardown flow into ROSE proper, updated for radical.orbit's current API (BridgeClient/ EndpointClient no longer exist — replaced by EndpointRuntime) and made spec-driven instead of interactive, matching --local's design. - rose/remote.py: new module — bootstrap a remote endpoint via IRI (OLCF), SFAPI (NERSC — IRI's /compute/* is broken server-side there), or PsiJ on an already-connected login node; wait for registration; run the workflow; tear down. - rose/spec/schema.py: RemoteConfig.target (TargetConfig) describes the bootstrap job (account/queue/walltime/tunnel/etc.). - rose/spec/__init__.py: fix stale bridge_url= kwarg passed to rhapsody's OrbitExecutionBackend (renamed to broker_url upstream) — was a silent TypeError waiting to happen. - rose/cli.py: --remote flag, mutually exclusive with --local. Also adds remote.embedded: hosts the ORBIT broker in-process (EmbeddedBroker, new upstream) instead of requiring an already-running one. Combined with target.kind: psij, drops target.edge_name entirely — PsiJ loads directly on the embedded broker, so `rose run --remote` on a login node needs no separate bootstrap endpoint at all. Tests: 30 new/extended unit tests (schema validation, mocked EndpointRuntime/EmbeddedBroker orchestration). Verified against real, unmocked radical.orbit (embedded broker + live PSIJClient resolution) in addition to the CLI entry point end-to-end with mocks. Docs: spec-api.md (remote.target, remote.embedded), rose-aas.md (bootstrap automation, no-broker-required callout). --- docs/user-guide/rose-aas.md | 12 +- docs/user-guide/spec-api.md | 78 +++++ rose/cli.py | 32 +- rose/remote.py | 313 ++++++++++++++++++++ rose/spec/__init__.py | 7 +- rose/spec/schema.py | 85 ++++++ tests/unit/spec/test_schema.py | 206 ++++++++++++- tests/unit/spec/test_workflow_spec.py | 45 +++ tests/unit/test_remote.py | 411 ++++++++++++++++++++++++++ 9 files changed, 1176 insertions(+), 13 deletions(-) create mode 100644 rose/remote.py create mode 100644 tests/unit/test_remote.py diff --git a/docs/user-guide/rose-aas.md b/docs/user-guide/rose-aas.md index 3d24f575..c12c6981 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) | -| 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 | +| Orbit backend (RHAPSODY, `broker_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,12 +35,14 @@ The first model is "submit a job, then run my loop inside it." The second is "ru ## How ROSE operates within a job -The orbit backend connects two things over a bridge: +The orbit backend connects two things over the ORBIT broker: -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. +1. **An orbit endpoint**, running inside an HPC job allocation. Historically this job had to be requested through whatever your site normally uses to get an allocation, outside of ROSE. `rose run --remote` now automates this step too: given a `remote.target` block in the spec (see the [YAML Spec API](spec-api.md#remote-config)), it submits the bootstrap job itself — via IRI, NERSC's Superfacility API (SFAPI), or PsiJ on an already-connected login node — and waits for the endpoint to register before handing off to step 2. 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 broker 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. +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 broker. 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 broker, on every iteration, with no per-iteration job resubmission. + +**No standalone broker required either.** `remote.embedded: true` (see the [YAML Spec API](spec-api.md#remoteembedded--run-without-a-standalone-broker)) hosts the broker inside your own orchestration process — nothing to deploy or keep running separately. Combined with `target.kind: psij`, running `rose run --remote` directly on a login node needs no bootstrap connection at all: PsiJ submits the compute-node job from the same process that's coordinating the loop. This is the fullest "as a Service" story so far — one command, one process, no infrastructure to stand up first (you still need an operator-placed broker cert/token under `~/.radical/orbit`, same as any broker). 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. diff --git a/docs/user-guide/spec-api.md b/docs/user-guide/spec-api.md index 97b0dca0..ebf9b664 100644 --- a/docs/user-guide/spec-api.md +++ b/docs/user-guide/spec-api.md @@ -596,6 +596,81 @@ remote: `backends` is consumed only at engine-creation time and is **not** injected into task `kwargs`. +### `remote.target` — bootstrapping the endpoint for `rose run --remote` + +`remote.target` tells `rose run --remote` how to launch the remote +orbit endpoint itself, instead of assuming one is already running. Omit it +to keep today's manual-bootstrap behavior (start the endpoint yourself, use +`--local`/the orbit backend directly). + +```yaml +remote: + backends: [dragon_v3] + target: + kind: sfapi # iri | sfapi | psij + endpoint: nersc # iri/sfapi only: nersc | olcf + resource_id: perlmutter + account: amsc007 + queue_name: debug + walltime_min: 30 + n_nodes: 1 + constraint: cpu + home_dir: /global/u2/m/merzky + tunnel: none # none | forward | reverse +``` + +| Key | Type | Required | Description | +|-----|------|----------|-------------| +| `remote.target.kind` | `iri` \| `sfapi` \| `psij` | yes | Bootstrap mechanism. Use `sfapi` for NERSC (IRI's `/compute/*` routes are broken server-side there); `iri` for OLCF; `psij` to submit via an already-connected login-node endpoint instead of IRI/SFAPI. | +| `remote.target.endpoint` | string | `iri`/`sfapi` only | `nersc` or `olcf`. | +| `remote.target.resource_id` | string | `iri`/`sfapi` only | Target compute resource, e.g. `perlmutter`, `odo`. | +| `remote.target.home_dir` | string | `iri`/`sfapi` only | User `$HOME` on the target — resolves the endpoint wrapper script path. | +| `remote.target.login_host` | string | when `tunnel: forward` | Login host to tunnel through. | +| `remote.target.edge_name` | string | `psij` only, unless `remote.embedded: true` | Name of the already-connected login-node endpoint to submit through. Not needed when `remote.embedded: true` — PsiJ runs on the embedded broker itself. | +| `remote.target.executor` | string | no | PsiJ executor name (default: `local`). | +| `remote.target.account` | string | yes | Allocation/project account. | +| `remote.target.queue_name` | string | no | Queue/partition. | +| `remote.target.walltime_min` | int | no | Default `30`. | +| `remote.target.n_nodes` | int | no | Default `1`. | +| `remote.target.constraint` | string | no | Scheduler constraint (e.g. `cpu`). | +| `remote.target.reservation` | string | no | Reservation name. | +| `remote.target.workdir` | string | no | Job working directory. | +| `remote.target.environment` | dict | no | Extra environment variables for the bootstrap job. | +| `remote.target.setup` | list[str] | no | Shell setup lines run before the endpoint starts. | +| `remote.target.tunnel` | `none` \| `forward` \| `reverse` | no | Default `none`. SSH tunnel mode for the endpoint's broker connection. | + +Credentials are read from disk/env, never stored in the spec: `iri` reads a +bearer token from `~/.amsc/token_`; `sfapi` reads a client ID from +`$SFAPI_CLIENT_ID` and a private key from `~/.amsc/sfapi_key_.pem`. + +### `remote.embedded` — run without a standalone broker + +`remote.embedded: true` hosts the ORBIT broker inside the `rose run --remote` +process itself (`EndpointRuntime`/`EmbeddedBroker`) instead of connecting to +one already running elsewhere. No separate `radical-orbit-broker.py` +deployment is needed. Mutually exclusive with `remote.broker_url`. + +```yaml +remote: + embedded: true + target: + kind: psij + account: amsc007 + queue_name: debug + # edge_name omitted — PsiJ runs on the embedded broker itself +``` + +The embedded broker is still a **real** broker: it needs the same +operator-placed cert/key/token under `~/.radical/orbit` (or their env-var +redirects) as a standalone one — see `radical.orbit`'s `CLAUDE.md` +`TOKEN_RECIPE` if those aren't set up yet. + +Combined with `target.kind: psij`, `target.edge_name` is not required — PsiJ +loads directly on the embedded broker (this only works when the process +running `rose run --remote` itself has batch-scheduler access, e.g. you're +on a login node). `target.kind: iri`/`sfapi` work the same as without +`embedded` — only where the broker lives changes. + --- ## Tracking @@ -661,6 +736,9 @@ With `validate_imports=True`, every `module:callable` string is resolved in the | `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 | +| `remote.broker_url` | string | no | `$RADICAL_ORBIT_BROKER_URL` | Broker URL override for `--remote` (no on-disk fallback — orbit resolves URL from CLI/API/env only) | +| `remote.target` | object | no | `null` | Bootstrap config for `--remote` — see [`remote.target`](#remotetarget--bootstrapping-the-endpoint-for-rose-run---remote) | +| `remote.embedded` | bool | no | `false` | Host the broker in-process instead of connecting to one — see [`remote.embedded`](#remoteembedded--run-without-a-standalone-broker) | | `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/cli.py b/rose/cli.py index 412aa094..878c1da6 100644 --- a/rose/cli.py +++ b/rose/cli.py @@ -43,6 +43,19 @@ async def _main(): asyncio.run(_main()) +def _run_remote(yaml_path: Path) -> None: + from rose.remote import run_remote + from rose.spec import load_spec + + spec = load_spec(yaml_path) + if spec.config.remote.target is None: + sys.exit( + "--remote requires a 'remote.target' block in the spec " + f"({yaml_path}) — see docs/user-guide/spec-api.md#remote-config" + ) + asyncio.run(run_remote(spec)) + + def main(): parser = argparse.ArgumentParser(prog="rose") sub = parser.add_subparsers(dest="command") @@ -52,18 +65,29 @@ def main(): run_p.add_argument( "--local", action="store_true", help="Run locally using rhapsody concurrent backend" ) + run_p.add_argument( + "--remote", + action="store_true", + help="Run on a remote HPC target via the ORBIT broker " + "(spec-driven; see 'remote.target' in the YAML)", + ) run_p.add_argument( "--backend", - default="concurrent", + default="dragon_v3", 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) + if args.local and args.remote: + parser.error("--local and --remote are mutually exclusive") + elif args.local: + _run_local(args.yaml, args.backend) + elif args.remote: + _run_remote(args.yaml) + else: + parser.error("specify a mode: --local | --remote") else: parser.print_help() sys.exit(1) diff --git a/rose/remote.py b/rose/remote.py new file mode 100644 index 00000000..336bb396 --- /dev/null +++ b/rose/remote.py @@ -0,0 +1,313 @@ +""" +rose.remote — bootstrap and run a ROSE workflow on a remote ORBIT endpoint. + +Three phases, driven by the ``remote.target`` block of the workflow spec +(see ``rose.spec.schema.TargetConfig``) instead of interactive prompts: + +* **Bootstrap** — submit a job that starts a ``radical-orbit-endpoint`` + process on the configured target, via the broker's ``iri_connect`` / + ``sfapi_connect`` plugin (NERSC should use ``sfapi`` — IRI's + ``/compute/*`` routes are broken server-side there; OLCF stays on + ``iri``) or via PsiJ on an already-connected login-node endpoint. +* **Execute** — wait for the endpoint to register, then run the workflow's + task graph against it (``spec.workflow``, built in ``rose.spec``). +* **Teardown** — cancel the bootstrap job and disconnect. + +Ported from ``amsc/use-cases/service_utils.py``'s interactive discover/ +launch/teardown flow. That module's ``BridgeClient`` (``radical.edge.client``) +no longer exists — every participant, including this one, is now a +``radical.orbit.EndpointRuntime``. + +``remote.embedded: true`` hosts the ORBIT broker in this process +(``radical.orbit.embedded.EmbeddedBroker``) instead of connecting to an +already-running one — no separate ``radical-orbit-broker.py`` deployment +needed. Combined with ``target.kind: psij``, PsiJ runs on the embedded +broker itself (no separate login-node endpoint either). +""" + +from __future__ import annotations + +import os +import time +import uuid +from pathlib import Path +from typing import TYPE_CHECKING, Any, Dict, Optional + +if TYPE_CHECKING: + from .spec import WorkflowSpec + from .spec.schema import TargetConfig + +_ENDPOINT_WRAPPER = 'radical-orbit-endpoint-wrapper.sh' +_BROKER = 'broker' # the broker's own participant name (src='broker'), embedded or not + +_PLUGIN_FOR_KIND = {'iri': 'iri_connect', 'sfapi': 'sfapi_connect', 'psij': 'psij'} + +AMSC_DIR = Path(os.environ.get('AMSC_DIR') or Path.home() / '.amsc').expanduser() + + +def _wrapper_path(home_dir: str) -> str: + return f"{home_dir.rstrip('/')}/.amsc/ve/bin/{_ENDPOINT_WRAPPER}" + + +# ───────────────────────────────────────────────────────────────────────────── +# Bootstrap: iri / sfapi +# ───────────────────────────────────────────────────────────────────────────── + +def _read_iri_token(endpoint: str) -> str: + path = AMSC_DIR / f'token_{endpoint}' + if not path.exists(): + raise RuntimeError( + f'IRI token file missing: {path} (put your IRI bearer token there)') + token = path.read_text().strip() + if not token: + raise RuntimeError(f'IRI token file is empty: {path}') + return token + + +def _read_sfapi_credentials(endpoint: str) -> tuple: + client_id = os.environ.get('SFAPI_CLIENT_ID', '').strip() + if not client_id: + raise RuntimeError( + "SFAPI_CLIENT_ID env var is required for remote.target.kind='sfapi'") + key_path = AMSC_DIR / f'sfapi_key_{endpoint}.pem' + if not key_path.exists(): + raise RuntimeError(f'SFAPI private key not found: {key_path}') + private_key = key_path.read_text().strip() + if not private_key: + raise RuntimeError(f'SFAPI private key file is empty: {key_path}') + return client_id, private_key + + +def _launch_iri_or_sfapi(rt, target: 'TargetConfig', broker_url: str) -> Dict[str, Any]: + plugin_name = 'sfapi_connect' if target.kind == 'sfapi' else 'iri_connect' + connect_client = rt.get_plugin(_BROKER, plugin_name) + + if target.kind == 'sfapi': + client_id, private_key = _read_sfapi_credentials(target.endpoint) + instance = connect_client.connect(target.endpoint, client_id, private_key) + else: + token = _read_iri_token(target.endpoint) + instance = connect_client.connect(target.endpoint, token) + + endpoint_name = f'rose-{target.endpoint}-{uuid.uuid4().hex[:6]}' + args = ['--name', endpoint_name, '--url', broker_url] + if target.tunnel == 'forward': + if not target.login_host: + raise RuntimeError( + "remote.target.login_host is required when tunnel='forward'") + args += ['--tunnel', 'forward', '--tunnel-via', target.login_host] + elif target.tunnel == 'reverse': + args += ['--tunnel', 'reverse'] + + attrs: Dict[str, Any] = { + 'queue_name': target.queue_name, + 'duration' : target.walltime_min * 60, + 'account' : target.account, + } + if target.constraint: attrs['constraint'] = target.constraint + if target.reservation: attrs['reservation'] = target.reservation + + env = {'RADICAL_ORBIT_BROKER_URL': broker_url} + env.update(target.environment) + if target.setup: + env['RADICAL_ORBIT_ENDPOINT_SETUP'] = '; '.join(target.setup) + + job_spec: Dict[str, Any] = { + 'executable' : _wrapper_path(target.home_dir), + 'arguments' : args, + 'name' : endpoint_name, + 'resources' : {'node_count': target.n_nodes, 'process_count': 1}, + 'attributes' : attrs, + 'environment': env, + } + if target.workdir: + job_spec['directory'] = target.workdir + + print(f'[rose --remote] submitting {target.kind} job ' + f'({target.endpoint} -> {target.resource_id}, ' + f'endpoint: {endpoint_name})...') + job = instance.submit_job(target.resource_id, job_spec) + print(f"[rose --remote] {target.kind} job_id: {job['job_id']}") + + return { + 'kind' : target.kind, + 'client' : instance, + 'connect_client': connect_client, + 'endpoint_key' : target.endpoint, + 'endpoint_name': endpoint_name, + 'resource_id' : target.resource_id, + 'job_id' : job['job_id'], + } + + +# ───────────────────────────────────────────────────────────────────────────── +# Bootstrap: psij (login-node submission) +# ───────────────────────────────────────────────────────────────────────────── + +def _launch_psij(rt, target: 'TargetConfig', broker_url: str, + embedded: bool = False) -> Dict[str, Any]: + if embedded: + # PsiJ runs on the embedded broker itself (this process) — no + # separate login-node endpoint to connect to, and no remote + # sysinfo.homedir() to call: we ARE the host, so use the local home. + edge_name = _BROKER + home = target.home_dir or str(Path.home()) + else: + edge_name = target.edge_name + home = rt.get_plugin(edge_name, 'sysinfo').homedir() + psij = rt.get_plugin(edge_name, 'psij') + executor = target.executor or 'local' + + attrs: Dict[str, Any] = {'duration': target.walltime_min * 60, + 'account' : target.account} + if target.queue_name: + attrs['queue_name'] = target.queue_name + custom_attrs: Dict[str, Any] = {} + if target.constraint: + custom_attrs[f'{executor}.constraint'] = target.constraint + + env = {'RADICAL_ORBIT_BROKER_URL': broker_url} + if target.setup: + env['RADICAL_ORBIT_ENDPOINT_SETUP'] = '; '.join(target.setup) + + child_name = f'rose-{edge_name}-{uuid.uuid4().hex[:6]}' + job_spec = { + 'executable' : _wrapper_path(home), + 'arguments' : ['--name', child_name, '--url', broker_url], + 'attributes' : attrs, + 'custom_attributes' : custom_attrs, + 'resources' : {'node_count': target.n_nodes, 'process_count': 1}, + 'environment' : env, + } + print(f'[rose --remote] submitting psij job via {edge_name} ' + f'(executor: {executor}, endpoint: {child_name})...') + res = psij.submit_tunneled(job_spec, executor=executor, tunnel=target.tunnel) + print(f"[rose --remote] psij job_id: {res['job_id']}") + + return { + 'kind' : 'psij', + 'client' : psij, + 'endpoint_name': res.get('endpoint_name', child_name), + 'job_id' : res['job_id'], + 'parent_edge' : edge_name, + } + + +# ───────────────────────────────────────────────────────────────────────────── +# Wait for the bootstrapped endpoint to register +# ───────────────────────────────────────────────────────────────────────────── + +def _wait_for_endpoint(rt, name: str, timeout: float = 30 * 60, + poll: float = 3.0, heartbeat: float = 30.0) -> None: + print(f'[rose --remote] waiting for endpoint {name!r} to register...') + start = time.time() + last_beat = start + while time.time() - start < timeout: + if name in rt.topology(): + return + time.sleep(poll) + if time.time() - last_beat >= heartbeat: + elapsed = int(time.time() - start) + print(f'[rose --remote] ...{elapsed}s elapsed, ' + f'{int(timeout - elapsed)}s left') + last_beat = time.time() + raise TimeoutError(f'endpoint {name!r} did not register within {timeout}s') + + +# ───────────────────────────────────────────────────────────────────────────── +# Teardown +# ───────────────────────────────────────────────────────────────────────────── + +def _teardown(created: Optional[Dict[str, Any]]) -> None: + if not created: + return + print('[rose --remote] tearing down bootstrap job...') + try: + if created['kind'] == 'psij': + created['client'].cancel_job(created['job_id']) + else: + created['client'].cancel_job(created['resource_id'], created['job_id']) + print(f"[rose --remote] cancelled {created['kind']} job {created['job_id']}") + except Exception as exc: + print(f"[rose --remote] could not cancel {created['kind']} " + f"job {created['job_id']}: {exc}") + + if created['kind'] in ('iri', 'sfapi'): + try: + created['connect_client'].disconnect(created['endpoint_key']) + print(f"[rose --remote] disconnected {created['kind']} " + f"endpoint {created['endpoint_key']}") + except Exception as exc: + print(f"[rose --remote] could not disconnect " + f"{created['endpoint_key']}: {exc}") + + +# ───────────────────────────────────────────────────────────────────────────── +# Entry point +# ───────────────────────────────────────────────────────────────────────────── + +async def run_remote(spec: 'WorkflowSpec') -> None: + """Bootstrap the ``remote.target`` from *spec*, run the workflow against + it, then tear the bootstrap job down. + + Raises ``RuntimeError`` if ``spec.config.remote.target`` is unset — the + CLI is expected to check this and produce a clearer error before calling + in, but this guard keeps the function safe to call directly (e.g. from + tests) too. + + ``spec.config.remote.embedded`` swaps connecting to an already-running + broker for hosting one in this process (``EmbeddedBroker``, loaded with + only the plugin the target needs) — no separate + ``radical-orbit-broker.py`` deployment required. + """ + from radical.orbit import EndpointRuntime + + remote_cfg = spec.config.remote + target = remote_cfg.target + if target is None: + raise RuntimeError( + "run_remote() requires 'remote.target' to be set in the spec") + + eb = None + if remote_cfg.embedded: + from radical.orbit.embedded import EmbeddedBroker + + eb = EmbeddedBroker(plugins=_PLUGIN_FOR_KIND[target.kind]) + eb.start() + print(f'[rose --remote] embedded broker: {eb.url}') + try: + rt = EndpointRuntime(broker_url=eb.url) + rt.start(wait=True) + except Exception: + eb.stop() + raise + else: + rt = EndpointRuntime(broker_url=remote_cfg.broker_url) + rt.start(wait=True) + + broker_url = rt.broker_url + print(f'[rose --remote] broker: {broker_url}') + + created: Optional[Dict[str, Any]] = None + try: + if target.kind == 'psij': + if not remote_cfg.embedded and target.edge_name not in rt.topology(): + raise RuntimeError( + f'remote.target.edge_name {target.edge_name!r} is not ' + f'currently connected to the broker') + created = _launch_psij(rt, target, broker_url, + embedded=remote_cfg.embedded) + else: + created = _launch_iri_or_sfapi(rt, target, broker_url) + + _wait_for_endpoint(rt, created['endpoint_name']) + print(f"[rose --remote] endpoint {created['endpoint_name']!r} is up") + + await spec.workflow(broker_url, created['endpoint_name']) + + finally: + _teardown(created) + rt.stop() + if eb is not None: + eb.stop() + print('[rose --remote] done.') diff --git a/rose/spec/__init__.py b/rose/spec/__init__.py index 39146f91..dbd282fe 100644 --- a/rose/spec/__init__.py +++ b/rose/spec/__init__.py @@ -77,7 +77,8 @@ def _validate_task_imports(cfg: WorkflowConfig) -> None: class WorkflowSpec: - """Validated workflow spec that produces a coroutine compatible with service_utils.run().""" + """Validated workflow spec that produces a coroutine compatible with + ``rose.remote.run_remote()`` (``workflow(broker_url, endpoint_name)``).""" def __init__(self, config: WorkflowConfig) -> None: self.config = config @@ -111,13 +112,13 @@ def workflow_with(self, **overrides: Any) -> WorkflowSpec: def workflow(self) -> Callable[..., Coroutine[Any, Any, None]]: cfg = self.config - async def _workflow(bridge_url: str, endpoint_name: str) -> None: + async def _workflow(broker_url: str, endpoint_name: str) -> None: import rhapsody from radical.asyncflow import WorkflowEngine engine = await rhapsody.get_backend( "orbit", - bridge_url=bridge_url, + broker_url=broker_url, endpoint_name=endpoint_name, backends=cfg.remote.backends, ) diff --git a/rose/spec/schema.py b/rose/spec/schema.py index d8da9cb2..0ea90dcb 100644 --- a/rose/spec/schema.py +++ b/rose/spec/schema.py @@ -76,12 +76,97 @@ def tasks(self) -> dict[str, TaskDef]: # ── Remote / tracking ──────────────────────────────────────────────────────── +class TargetConfig(BaseModel): + """Bootstrap config for `rose run --remote`: how to launch the remote + ORBIT endpoint before the workflow's tasks run on it. + + ``kind`` selects the launch mechanism: + - ``iri``/``sfapi``: submit a job to `resource_id` through the broker's + ``iri_connect``/``sfapi_connect`` plugin (NERSC currently should use + ``sfapi`` — IRI's ``/compute/*`` routes are broken server-side there; + OLCF has no such issue and stays on ``iri``). + - ``psij``: submit a job via PsiJ. Normally on an already-connected + login-node endpoint (``edge_name``); with ``remote.embedded: true``, + ``edge_name`` is not needed — PsiJ runs on the embedded broker itself + (see ``RemoteConfig.embedded``). + """ + + kind: Literal["iri", "sfapi", "psij"] + + # iri / sfapi + endpoint: str | None = None # 'nersc' | 'olcf' + resource_id: str | None = None + home_dir: str | None = None # user $HOME on target; resolves the wrapper path + login_host: str | None = None # for tunnel='forward' + + # psij + edge_name: str | None = None # login-node endpoint already in the topology + executor: str | None = None + + # shared submission attributes + account: str | None = None + queue_name: str | None = None + walltime_min: int = 30 + n_nodes: int = 1 + constraint: str | None = None + reservation: str | None = None + workdir: str | None = None + environment: dict[str, str] = {} + setup: list[str] = [] + tunnel: Literal["none", "forward", "reverse"] = "none" + + model_config = {"extra": "forbid"} + + @model_validator(mode="after") + def _check_fields(self) -> TargetConfig: + if self.kind in ("iri", "sfapi"): + if not self.endpoint: + raise ValueError(f"remote.target.kind={self.kind!r} requires 'endpoint'") + if not self.resource_id: + raise ValueError(f"remote.target.kind={self.kind!r} requires 'resource_id'") + if not self.home_dir: + raise ValueError(f"remote.target.kind={self.kind!r} requires 'home_dir'") + # 'psij' + 'edge_name' is checked at the RemoteConfig level: whether + # edge_name is required depends on RemoteConfig.embedded, a sibling + # field this model can't see. + if not self.account: + raise ValueError(f"remote.target.kind={self.kind!r} requires 'account'") + return self + + class RemoteConfig(BaseModel): pythonpath: list[str] = [] backends: list[str] = ["dragon_v3"] + broker_url: str | None = None + target: TargetConfig | None = None + # Host the ORBIT broker in-process (EmbeddedBroker) instead of connecting + # to an already-running one — no separate radical-orbit-broker.py + # deployment needed. Mutually exclusive with broker_url. Combined with + # target.kind: psij, target.edge_name is not required — PsiJ runs on the + # embedded broker itself. + embedded: bool = False model_config = {"extra": "forbid"} + @model_validator(mode="after") + def _check_embedded(self) -> RemoteConfig: + if self.embedded and self.broker_url: + raise ValueError( + "remote.embedded and remote.broker_url are mutually exclusive " + "— the embedded broker provides its own URL" + ) + if ( + self.target is not None + and self.target.kind == "psij" + and not self.embedded + and not self.target.edge_name + ): + raise ValueError( + "remote.target.kind='psij' requires 'edge_name' " + "(unless remote.embedded is true)" + ) + return self + class TrackingConfig(BaseModel): backend: Literal["mlflow", "clearml", "none"] = "none" diff --git a/tests/unit/spec/test_schema.py b/tests/unit/spec/test_schema.py index beb3e9fb..866cce29 100644 --- a/tests/unit/spec/test_schema.py +++ b/tests/unit/spec/test_schema.py @@ -5,7 +5,7 @@ import pytest from pydantic import ValidationError -from rose.spec.schema import TaskDef, WorkflowConfig +from rose.spec.schema import RemoteConfig, TargetConfig, TaskDef, WorkflowConfig # ── TaskDef ─────────────────────────────────────────────────────────────────── @@ -673,6 +673,210 @@ def test_parameters_defaults_empty(tmp_path): # ── WorkflowConfig — parallel learners (mixed types) ───────────────────────────── +# ── RemoteConfig.target (TargetConfig) ─────────────────────────────────────── + + +def test_remote_target_default_none(): + assert RemoteConfig().target is None + + +def test_remote_embedded_default_false(): + assert RemoteConfig().embedded is False + + +def test_remote_embedded_and_broker_url_mutually_exclusive(): + with pytest.raises(ValueError, match="mutually exclusive"): + RemoteConfig(embedded=True, broker_url="https://broker:8000") + + +def test_remote_embedded_without_broker_url_valid(): + cfg = RemoteConfig(embedded=True) + assert cfg.embedded is True + assert cfg.broker_url is None + + +def test_remote_psij_edge_name_required_unless_embedded(): + with pytest.raises(ValueError, match="requires 'edge_name'"): + RemoteConfig( + embedded=False, + target=TargetConfig(kind="psij", account="amsc007"), + ) + + +def test_remote_psij_edge_name_not_required_when_embedded(): + cfg = RemoteConfig( + embedded=True, + target=TargetConfig(kind="psij", account="amsc007"), + ) + assert cfg.target.edge_name is None + + +def test_remote_psij_edge_name_still_valid_when_embedded(): + """A user may still set edge_name when embedded — harmless, ignored by + the RemoteConfig-level requirement check.""" + cfg = RemoteConfig( + embedded=True, + target=TargetConfig(kind="psij", edge_name="login1", account="amsc007"), + ) + assert cfg.target.edge_name == "login1" + + +def test_remote_iri_target_unaffected_by_embedded_validator(): + cfg = RemoteConfig( + embedded=True, + target=TargetConfig( + kind="iri", endpoint="olcf", resource_id="odo", + account="fus183", home_dir="/home/x", + ), + ) + assert cfg.target.kind == "iri" + + +def test_target_config_sfapi_valid(): + t = TargetConfig( + kind="sfapi", + endpoint="nersc", + resource_id="perlmutter", + account="amsc007", + home_dir="/global/u2/m/merzky", + ) + assert t.kind == "sfapi" + assert t.tunnel == "none" + assert t.walltime_min == 30 + + +def test_target_config_iri_missing_endpoint_raises(): + with pytest.raises(ValueError, match="requires 'endpoint'"): + TargetConfig(kind="iri", resource_id="odo", account="fus183", home_dir="/home/x") + + +def test_target_config_iri_missing_resource_id_raises(): + with pytest.raises(ValueError, match="requires 'resource_id'"): + TargetConfig(kind="iri", endpoint="olcf", account="fus183", home_dir="/home/x") + + +def test_target_config_iri_missing_home_dir_raises(): + with pytest.raises(ValueError, match="requires 'home_dir'"): + TargetConfig(kind="iri", endpoint="olcf", resource_id="odo", account="fus183") + + +def test_target_config_psij_missing_edge_name_allowed_standalone(): + """TargetConfig alone no longer enforces 'edge_name' for kind='psij' — + whether it's required depends on RemoteConfig.embedded, a sibling field + this model can't see. See RemoteConfig-level tests below for the actual + enforcement (edge_name required unless embedded=True).""" + t = TargetConfig(kind="psij", account="amsc007") + assert t.edge_name is None + + +def test_target_config_psij_valid(): + t = TargetConfig(kind="psij", edge_name="perlmutter-login", account="amsc007") + assert t.edge_name == "perlmutter-login" + + +def test_target_config_missing_account_raises(): + with pytest.raises(ValueError, match="requires 'account'"): + TargetConfig(kind="psij", edge_name="login1") + + +def test_target_config_extra_field_rejected(): + with pytest.raises(ValidationError): + TargetConfig( + kind="psij", edge_name="login1", account="a", bogus_field="x" + ) + + +def test_target_config_bad_tunnel_value_rejected(): + with pytest.raises(ValidationError): + TargetConfig(kind="psij", edge_name="login1", account="a", tunnel="yes") + + +def test_remote_target_roundtrip_from_yaml(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: [dragon_v3] + target: + kind: sfapi + endpoint: nersc + resource_id: perlmutter + account: amsc007 + queue_name: debug + walltime_min: 30 + n_nodes: 1 + constraint: cpu + home_dir: /global/u2/m/merzky + """) + p = tmp_path / "remote.yaml" + p.write_text(yaml) + cfg = WorkflowConfig.from_yaml(p) + assert cfg.remote.target is not None + assert cfg.remote.target.kind == "sfapi" + assert cfg.remote.target.endpoint == "nersc" + assert cfg.remote.target.constraint == "cpu" + + +def test_remote_embedded_psij_roundtrip_from_yaml(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: + embedded: true + target: + kind: psij + account: amsc007 + queue_name: debug + """) + p = tmp_path / "embedded.yaml" + p.write_text(yaml) + cfg = WorkflowConfig.from_yaml(p) + assert cfg.remote.embedded is True + assert cfg.remote.target.kind == "psij" + assert cfg.remote.target.edge_name is None + + +def test_remote_without_target_still_valid(tmp_path): + """Specs written before remote.target existed (e.g. the m3dc1 use case) + must keep validating unchanged.""" + p = tmp_path / "spec.yaml" + p.write_text(SEQ_YAML) + cfg = WorkflowConfig.from_yaml(p) + assert cfg.remote.target is None + + 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 a9e7fab8..836769d2 100644 --- a/tests/unit/spec/test_workflow_spec.py +++ b/tests/unit/spec/test_workflow_spec.py @@ -1,6 +1,9 @@ """Unit tests for WorkflowSpec — load_spec and workflow_with().""" +import sys import textwrap +import types +from unittest.mock import AsyncMock, MagicMock import pytest @@ -250,3 +253,45 @@ def test_validate_imports_default_false_skips_check(tmp_path): p.write_text(BAD_MODULE_YAML) spec = load_spec(p) # validate_imports=False by default — should not raise assert spec is not None + + +# ── workflow: calls rhapsody.get_backend with broker_url, not bridge_url ────── +# +# Regression test for the orbit Bridge -> Broker rename: OrbitExecutionBackend +# (rhapsody) takes `broker_url=`, and this call site used to pass the stale +# `bridge_url=` kwarg (a TypeError against current rhapsody). `rhapsody` and +# `radical.asyncflow` are stubbed out and made to fail immediately inside +# get_backend, so this isolates the kwarg names without needing to simulate +# the rest of the execution stack (LearnerBuilder / learner.start / asyncflow). + + +class _Sentinel(Exception): + pass + + +async def test_workflow_calls_get_backend_with_broker_url(tmp_path, monkeypatch): + spec = _make_spec(tmp_path) + captured = {} + + async def fake_get_backend(name, **kwargs): + captured["name"] = name + captured.update(kwargs) + raise _Sentinel + + fake_rhapsody = types.ModuleType("rhapsody") + fake_rhapsody.get_backend = fake_get_backend + + fake_asyncflow_mod = types.ModuleType("radical.asyncflow") + fake_asyncflow_mod.WorkflowEngine = MagicMock() + + monkeypatch.setitem(sys.modules, "rhapsody", fake_rhapsody) + monkeypatch.setitem(sys.modules, "radical.asyncflow", fake_asyncflow_mod) + + with pytest.raises(_Sentinel): + await spec.workflow("https://broker.example:8000", "ep-1") + + assert captured["name"] == "orbit" + assert "bridge_url" not in captured + assert captured["broker_url"] == "https://broker.example:8000" + assert captured["endpoint_name"] == "ep-1" + assert captured["backends"] == spec.config.remote.backends diff --git a/tests/unit/test_remote.py b/tests/unit/test_remote.py new file mode 100644 index 00000000..501d50be --- /dev/null +++ b/tests/unit/test_remote.py @@ -0,0 +1,411 @@ +"""Unit tests for rose.remote — bootstrap/execute/teardown orchestration for +`rose run --remote`. radical.orbit's EndpointRuntime is mocked throughout; +these tests never touch a real broker or HPC system. +""" + +import textwrap +from unittest.mock import AsyncMock, MagicMock + +import pytest + +import rose.remote as remote +from rose.spec import load_spec + +SFAPI_TARGET_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: + target: + kind: sfapi + endpoint: nersc + resource_id: perlmutter + account: amsc007 + home_dir: /global/u2/m/merzky +""") + +EMBEDDED_PSIJ_TARGET_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: + embedded: true + target: + kind: psij + account: amsc007 +""") + +NO_TARGET_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 +""") + + +def _make_spec(tmp_path, yaml): + p = tmp_path / "spec.yaml" + p.write_text(yaml) + return load_spec(p) + + +# ── _wait_for_endpoint ───────────────────────────────────────────────────── + + +def test_wait_for_endpoint_succeeds_when_name_appears(): + rt = MagicMock() + rt.topology.side_effect = [{}, {}, {"ep-1": {}}] + remote._wait_for_endpoint(rt, "ep-1", timeout=10, poll=0) + + +def test_wait_for_endpoint_times_out(): + rt = MagicMock() + rt.topology.return_value = {} + with pytest.raises(TimeoutError, match="ep-1"): + remote._wait_for_endpoint(rt, "ep-1", timeout=0.05, poll=0.01) + + +# ── credential reading ─────────────────────────────────────────────────────── + + +def test_read_iri_token_missing_file_raises(tmp_path, monkeypatch): + monkeypatch.setattr(remote, "AMSC_DIR", tmp_path) + with pytest.raises(RuntimeError, match="token file missing"): + remote._read_iri_token("olcf") + + +def test_read_iri_token_reads_stripped_content(tmp_path, monkeypatch): + monkeypatch.setattr(remote, "AMSC_DIR", tmp_path) + (tmp_path / "token_olcf").write_text(" secret-token \n") + assert remote._read_iri_token("olcf") == "secret-token" + + +def test_read_sfapi_credentials_missing_client_id_raises(tmp_path, monkeypatch): + monkeypatch.setattr(remote, "AMSC_DIR", tmp_path) + monkeypatch.delenv("SFAPI_CLIENT_ID", raising=False) + with pytest.raises(RuntimeError, match="SFAPI_CLIENT_ID"): + remote._read_sfapi_credentials("nersc") + + +def test_read_sfapi_credentials_missing_key_file_raises(tmp_path, monkeypatch): + monkeypatch.setattr(remote, "AMSC_DIR", tmp_path) + monkeypatch.setenv("SFAPI_CLIENT_ID", "abc123") + with pytest.raises(RuntimeError, match="private key not found"): + remote._read_sfapi_credentials("nersc") + + +def test_read_sfapi_credentials_ok(tmp_path, monkeypatch): + monkeypatch.setattr(remote, "AMSC_DIR", tmp_path) + monkeypatch.setenv("SFAPI_CLIENT_ID", "abc123") + (tmp_path / "sfapi_key_nersc.pem").write_text("-----BEGIN KEY-----\n...") + client_id, key = remote._read_sfapi_credentials("nersc") + assert client_id == "abc123" + assert key.startswith("-----BEGIN KEY-----") + + +# ── teardown ────────────────────────────────────────────────────────────── + + +def test_teardown_none_is_noop(): + remote._teardown(None) # must not raise + + +def test_teardown_iri_cancels_and_disconnects(): + instance = MagicMock() + connect_client = MagicMock() + created = { + "kind": "sfapi", "client": instance, "connect_client": connect_client, + "endpoint_key": "nersc", "resource_id": "perlmutter", "job_id": "j1", + } + remote._teardown(created) + instance.cancel_job.assert_called_once_with("perlmutter", "j1") + connect_client.disconnect.assert_called_once_with("nersc") + + +def test_teardown_psij_cancels_only(): + psij = MagicMock() + created = {"kind": "psij", "client": psij, "job_id": "j1", + "parent_edge": "login1", "endpoint_name": "rose-login1-abc"} + remote._teardown(created) + psij.cancel_job.assert_called_once_with("j1") + + +def test_teardown_swallows_cancel_errors(): + instance = MagicMock() + instance.cancel_job.side_effect = RuntimeError("already gone") + connect_client = MagicMock() + created = { + "kind": "iri", "client": instance, "connect_client": connect_client, + "endpoint_key": "olcf", "resource_id": "odo", "job_id": "j1", + } + remote._teardown(created) # must not raise + connect_client.disconnect.assert_called_once_with("olcf") + + +# ── run_remote ──────────────────────────────────────────────────────────── + + +async def test_run_remote_requires_target(tmp_path): + spec = _make_spec(tmp_path, NO_TARGET_YAML) + with pytest.raises(RuntimeError, match="remote.target"): + await remote.run_remote(spec) + + +async def test_run_remote_psij_edge_not_in_topology_raises(tmp_path, monkeypatch): + 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: + target: + kind: psij + edge_name: login1 + account: amsc007 + """) + spec = _make_spec(tmp_path, yaml) + + fake_rt = MagicMock() + fake_rt.start.return_value = fake_rt + fake_rt.broker_url = "https://broker:8000" + fake_rt.topology.return_value = {} + monkeypatch.setattr(remote, "_teardown", MagicMock()) + + import radical.orbit + monkeypatch.setattr(radical.orbit, "EndpointRuntime", MagicMock(return_value=fake_rt)) + + with pytest.raises(RuntimeError, match="login1"): + await remote.run_remote(spec) + fake_rt.stop.assert_called_once() + + +async def test_run_remote_sfapi_end_to_end(tmp_path, monkeypatch): + spec = _make_spec(tmp_path, SFAPI_TARGET_YAML) + monkeypatch.setattr(remote, "AMSC_DIR", tmp_path) + monkeypatch.setenv("SFAPI_CLIENT_ID", "abc123") + (tmp_path / "sfapi_key_nersc.pem").write_text("-----BEGIN KEY-----\n...") + + instance = MagicMock() + instance.submit_job.return_value = {"job_id": "job-42"} + + connect_client = MagicMock() + connect_client.connect.return_value = instance + + fake_rt = MagicMock() + fake_rt.start.return_value = fake_rt + fake_rt.broker_url = "https://broker:8000" + fake_rt.get_plugin.return_value = connect_client + # topology never contains the bootstrapped endpoint (loop exits fast via patched wait) + + monkeypatch.setattr(remote, "_wait_for_endpoint", MagicMock()) + fake_workflow = AsyncMock() + monkeypatch.setattr(type(spec), "workflow", property(lambda self: fake_workflow)) + + import radical.orbit + monkeypatch.setattr(radical.orbit, "EndpointRuntime", MagicMock(return_value=fake_rt)) + + await remote.run_remote(spec) + + fake_rt.get_plugin.assert_called_once_with("broker", "sfapi_connect") + connect_client.connect.assert_called_once_with("nersc", "abc123", + "-----BEGIN KEY-----\n...") + instance.submit_job.assert_called_once() + resource_id, job_spec = instance.submit_job.call_args[0] + assert resource_id == "perlmutter" + assert job_spec["attributes"]["account"] == "amsc007" + + fake_workflow.assert_awaited_once() + call_args = fake_workflow.await_args[0] + assert call_args[0] == "https://broker:8000" + + instance.cancel_job.assert_called_once() + connect_client.disconnect.assert_called_once_with("nersc") + fake_rt.stop.assert_called_once() + + +# ── _launch_psij: embedded vs standalone plugin/home_dir source ─────────── + + +def test_launch_psij_embedded_uses_broker_and_local_home(monkeypatch): + psij = MagicMock() + psij.submit_tunneled.return_value = {"job_id": "j1", "endpoint_name": "rose-broker-abc"} + + rt = MagicMock() + rt.get_plugin.return_value = psij + + from rose.spec.schema import TargetConfig + target = TargetConfig(kind="psij", account="amsc007") + + class _FakeHomePath: + @staticmethod + def home(): + from pathlib import Path as _RealPath + return _RealPath("/home/local") + + # Rebind the name `Path` inside rose.remote only — must not touch the + # real pathlib.Path class, which other code/tests still rely on. + monkeypatch.setattr(remote, "Path", _FakeHomePath) + + created = remote._launch_psij(rt, target, "https://broker:8000", embedded=True) + + rt.get_plugin.assert_called_once_with("broker", "psij") + job_spec = psij.submit_tunneled.call_args[0][0] + assert job_spec["executable"] == "/home/local/.amsc/ve/bin/radical-orbit-endpoint-wrapper.sh" + assert created["endpoint_name"] == "rose-broker-abc" + + +def test_launch_psij_embedded_prefers_explicit_home_dir(): + psij = MagicMock() + psij.submit_tunneled.return_value = {"job_id": "j1", "endpoint_name": "ep"} + rt = MagicMock() + rt.get_plugin.return_value = psij + + from rose.spec.schema import TargetConfig + target = TargetConfig(kind="psij", account="amsc007", home_dir="/custom/home") + + remote._launch_psij(rt, target, "https://broker:8000", embedded=True) + + job_spec = psij.submit_tunneled.call_args[0][0] + assert job_spec["executable"] == "/custom/home/.amsc/ve/bin/radical-orbit-endpoint-wrapper.sh" + + +def test_launch_psij_non_embedded_uses_edge_and_remote_sysinfo(): + psij = MagicMock() + psij.submit_tunneled.return_value = {"job_id": "j1", "endpoint_name": "ep"} + sysinfo = MagicMock() + sysinfo.homedir.return_value = "/remote/home" + + rt = MagicMock() + rt.get_plugin.side_effect = lambda edge, name: {"psij": psij, "sysinfo": sysinfo}[name] + + from rose.spec.schema import TargetConfig + target = TargetConfig(kind="psij", edge_name="login1", account="amsc007") + + remote._launch_psij(rt, target, "https://broker:8000", embedded=False) + + rt.get_plugin.assert_any_call("login1", "psij") + rt.get_plugin.assert_any_call("login1", "sysinfo") + job_spec = psij.submit_tunneled.call_args[0][0] + assert job_spec["executable"] == "/remote/home/.amsc/ve/bin/radical-orbit-endpoint-wrapper.sh" + + +# ── run_remote: embedded broker ───────────────────────────────────────────── + + +async def test_run_remote_embedded_constructs_and_tears_down_broker(tmp_path, monkeypatch): + spec = _make_spec(tmp_path, EMBEDDED_PSIJ_TARGET_YAML) + + psij = MagicMock() + psij.submit_tunneled.return_value = {"job_id": "j1", "endpoint_name": "ep-1"} + + stop_order = [] + + fake_rt = MagicMock() + fake_rt.start.return_value = fake_rt + fake_rt.broker_url = "https://embedded:8000" + fake_rt.get_plugin.return_value = psij + fake_rt.stop.side_effect = lambda: stop_order.append("rt") + + fake_eb = MagicMock() + fake_eb.url = "https://embedded:8000" + fake_eb.stop.side_effect = lambda: stop_order.append("eb") + + monkeypatch.setattr(remote, "_wait_for_endpoint", MagicMock()) + fake_workflow = AsyncMock() + monkeypatch.setattr(type(spec), "workflow", property(lambda self: fake_workflow)) + + import radical.orbit + import radical.orbit.embedded + monkeypatch.setattr(radical.orbit, "EndpointRuntime", MagicMock(return_value=fake_rt)) + fake_embedded_broker_cls = MagicMock(return_value=fake_eb) + monkeypatch.setattr(radical.orbit.embedded, "EmbeddedBroker", fake_embedded_broker_cls) + + await remote.run_remote(spec) + + fake_embedded_broker_cls.assert_called_once_with(plugins="psij") + fake_eb.start.assert_called_once() + radical.orbit.EndpointRuntime.assert_called_once_with(broker_url="https://embedded:8000") + fake_rt.get_plugin.assert_called_once_with("broker", "psij") + + # teardown order: runtime stops before the embedded broker + assert stop_order == ["rt", "eb"] + + +async def test_run_remote_embedded_stops_broker_if_runtime_start_fails(tmp_path, monkeypatch): + spec = _make_spec(tmp_path, EMBEDDED_PSIJ_TARGET_YAML) + + fake_eb = MagicMock() + fake_eb.url = "https://embedded:8000" + + fake_rt = MagicMock() + fake_rt.start.side_effect = RuntimeError("registration failed") + + import radical.orbit + import radical.orbit.embedded + monkeypatch.setattr(radical.orbit, "EndpointRuntime", MagicMock(return_value=fake_rt)) + monkeypatch.setattr(radical.orbit.embedded, "EmbeddedBroker", + MagicMock(return_value=fake_eb)) + + with pytest.raises(RuntimeError, match="registration failed"): + await remote.run_remote(spec) + + fake_eb.start.assert_called_once() + fake_eb.stop.assert_called_once() From 7e07df8fd4df9ce9ab988c06f3dbaa93ad5ebf30 Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Wed, 19 Aug 2026 00:19:56 +0200 Subject: [PATCH 15/19] pre-commit --- rose/cli.py | 2 +- rose/remote.py | 245 +++++++++++++------------- rose/spec/schema.py | 15 +- tests/unit/spec/test_schema.py | 29 +-- tests/unit/spec/test_workflow_spec.py | 8 +- tests/unit/test_remote.py | 46 +++-- 6 files changed, 186 insertions(+), 159 deletions(-) diff --git a/rose/cli.py b/rose/cli.py index 878c1da6..c4181219 100644 --- a/rose/cli.py +++ b/rose/cli.py @@ -69,7 +69,7 @@ def main(): "--remote", action="store_true", help="Run on a remote HPC target via the ORBIT broker " - "(spec-driven; see 'remote.target' in the YAML)", + "(spec-driven; see 'remote.target' in the YAML)", ) run_p.add_argument( "--backend", diff --git a/rose/remote.py b/rose/remote.py index 336bb396..937023d7 100644 --- a/rose/remote.py +++ b/rose/remote.py @@ -1,5 +1,4 @@ -""" -rose.remote — bootstrap and run a ROSE workflow on a remote ORBIT endpoint. +"""rose.remote — bootstrap and run a ROSE workflow on a remote ORBIT endpoint. Three phases, driven by the ``remote.target`` block of the workflow spec (see ``rose.spec.schema.TargetConfig``) instead of interactive prompts: @@ -31,18 +30,18 @@ import time import uuid from pathlib import Path -from typing import TYPE_CHECKING, Any, Dict, Optional +from typing import TYPE_CHECKING, Any if TYPE_CHECKING: from .spec import WorkflowSpec from .spec.schema import TargetConfig -_ENDPOINT_WRAPPER = 'radical-orbit-endpoint-wrapper.sh' -_BROKER = 'broker' # the broker's own participant name (src='broker'), embedded or not +_ENDPOINT_WRAPPER = "radical-orbit-endpoint-wrapper.sh" +_BROKER = "broker" # the broker's own participant name (src='broker'), embedded or not -_PLUGIN_FOR_KIND = {'iri': 'iri_connect', 'sfapi': 'sfapi_connect', 'psij': 'psij'} +_PLUGIN_FOR_KIND = {"iri": "iri_connect", "sfapi": "sfapi_connect", "psij": "psij"} -AMSC_DIR = Path(os.environ.get('AMSC_DIR') or Path.home() / '.amsc').expanduser() +AMSC_DIR = Path(os.environ.get("AMSC_DIR") or Path.home() / ".amsc").expanduser() def _wrapper_path(home_dir: str) -> str: @@ -53,90 +52,92 @@ def _wrapper_path(home_dir: str) -> str: # Bootstrap: iri / sfapi # ───────────────────────────────────────────────────────────────────────────── + def _read_iri_token(endpoint: str) -> str: - path = AMSC_DIR / f'token_{endpoint}' + path = AMSC_DIR / f"token_{endpoint}" if not path.exists(): - raise RuntimeError( - f'IRI token file missing: {path} (put your IRI bearer token there)') + raise RuntimeError(f"IRI token file missing: {path} (put your IRI bearer token there)") token = path.read_text().strip() if not token: - raise RuntimeError(f'IRI token file is empty: {path}') + raise RuntimeError(f"IRI token file is empty: {path}") return token def _read_sfapi_credentials(endpoint: str) -> tuple: - client_id = os.environ.get('SFAPI_CLIENT_ID', '').strip() + client_id = os.environ.get("SFAPI_CLIENT_ID", "").strip() if not client_id: - raise RuntimeError( - "SFAPI_CLIENT_ID env var is required for remote.target.kind='sfapi'") - key_path = AMSC_DIR / f'sfapi_key_{endpoint}.pem' + raise RuntimeError("SFAPI_CLIENT_ID env var is required for remote.target.kind='sfapi'") + key_path = AMSC_DIR / f"sfapi_key_{endpoint}.pem" if not key_path.exists(): - raise RuntimeError(f'SFAPI private key not found: {key_path}') + raise RuntimeError(f"SFAPI private key not found: {key_path}") private_key = key_path.read_text().strip() if not private_key: - raise RuntimeError(f'SFAPI private key file is empty: {key_path}') + raise RuntimeError(f"SFAPI private key file is empty: {key_path}") return client_id, private_key -def _launch_iri_or_sfapi(rt, target: 'TargetConfig', broker_url: str) -> Dict[str, Any]: - plugin_name = 'sfapi_connect' if target.kind == 'sfapi' else 'iri_connect' +def _launch_iri_or_sfapi(rt, target: TargetConfig, broker_url: str) -> dict[str, Any]: + plugin_name = "sfapi_connect" if target.kind == "sfapi" else "iri_connect" connect_client = rt.get_plugin(_BROKER, plugin_name) - if target.kind == 'sfapi': + if target.kind == "sfapi": client_id, private_key = _read_sfapi_credentials(target.endpoint) instance = connect_client.connect(target.endpoint, client_id, private_key) else: token = _read_iri_token(target.endpoint) instance = connect_client.connect(target.endpoint, token) - endpoint_name = f'rose-{target.endpoint}-{uuid.uuid4().hex[:6]}' - args = ['--name', endpoint_name, '--url', broker_url] - if target.tunnel == 'forward': + endpoint_name = f"rose-{target.endpoint}-{uuid.uuid4().hex[:6]}" + args = ["--name", endpoint_name, "--url", broker_url] + if target.tunnel == "forward": if not target.login_host: - raise RuntimeError( - "remote.target.login_host is required when tunnel='forward'") - args += ['--tunnel', 'forward', '--tunnel-via', target.login_host] - elif target.tunnel == 'reverse': - args += ['--tunnel', 'reverse'] - - attrs: Dict[str, Any] = { - 'queue_name': target.queue_name, - 'duration' : target.walltime_min * 60, - 'account' : target.account, + raise RuntimeError("remote.target.login_host is required when tunnel='forward'") + args += ["--tunnel", "forward", "--tunnel-via", target.login_host] + elif target.tunnel == "reverse": + args += ["--tunnel", "reverse"] + + attrs: dict[str, Any] = { + "queue_name": target.queue_name, + "duration": target.walltime_min * 60, + "account": target.account, } - if target.constraint: attrs['constraint'] = target.constraint - if target.reservation: attrs['reservation'] = target.reservation + if target.constraint: + attrs["constraint"] = target.constraint + if target.reservation: + attrs["reservation"] = target.reservation - env = {'RADICAL_ORBIT_BROKER_URL': broker_url} + env = {"RADICAL_ORBIT_BROKER_URL": broker_url} env.update(target.environment) if target.setup: - env['RADICAL_ORBIT_ENDPOINT_SETUP'] = '; '.join(target.setup) - - job_spec: Dict[str, Any] = { - 'executable' : _wrapper_path(target.home_dir), - 'arguments' : args, - 'name' : endpoint_name, - 'resources' : {'node_count': target.n_nodes, 'process_count': 1}, - 'attributes' : attrs, - 'environment': env, + env["RADICAL_ORBIT_ENDPOINT_SETUP"] = "; ".join(target.setup) + + job_spec: dict[str, Any] = { + "executable": _wrapper_path(target.home_dir), + "arguments": args, + "name": endpoint_name, + "resources": {"node_count": target.n_nodes, "process_count": 1}, + "attributes": attrs, + "environment": env, } if target.workdir: - job_spec['directory'] = target.workdir + job_spec["directory"] = target.workdir - print(f'[rose --remote] submitting {target.kind} job ' - f'({target.endpoint} -> {target.resource_id}, ' - f'endpoint: {endpoint_name})...') + print( + f"[rose --remote] submitting {target.kind} job " + f"({target.endpoint} -> {target.resource_id}, " + f"endpoint: {endpoint_name})..." + ) job = instance.submit_job(target.resource_id, job_spec) print(f"[rose --remote] {target.kind} job_id: {job['job_id']}") return { - 'kind' : target.kind, - 'client' : instance, - 'connect_client': connect_client, - 'endpoint_key' : target.endpoint, - 'endpoint_name': endpoint_name, - 'resource_id' : target.resource_id, - 'job_id' : job['job_id'], + "kind": target.kind, + "client": instance, + "connect_client": connect_client, + "endpoint_key": target.endpoint, + "endpoint_name": endpoint_name, + "resource_id": target.resource_id, + "job_id": job["job_id"], } @@ -144,52 +145,55 @@ def _launch_iri_or_sfapi(rt, target: 'TargetConfig', broker_url: str) -> Dict[st # Bootstrap: psij (login-node submission) # ───────────────────────────────────────────────────────────────────────────── -def _launch_psij(rt, target: 'TargetConfig', broker_url: str, - embedded: bool = False) -> Dict[str, Any]: + +def _launch_psij( + rt, target: TargetConfig, broker_url: str, embedded: bool = False +) -> dict[str, Any]: if embedded: # PsiJ runs on the embedded broker itself (this process) — no # separate login-node endpoint to connect to, and no remote # sysinfo.homedir() to call: we ARE the host, so use the local home. edge_name = _BROKER - home = target.home_dir or str(Path.home()) + home = target.home_dir or str(Path.home()) else: edge_name = target.edge_name - home = rt.get_plugin(edge_name, 'sysinfo').homedir() - psij = rt.get_plugin(edge_name, 'psij') - executor = target.executor or 'local' + home = rt.get_plugin(edge_name, "sysinfo").homedir() + psij = rt.get_plugin(edge_name, "psij") + executor = target.executor or "local" - attrs: Dict[str, Any] = {'duration': target.walltime_min * 60, - 'account' : target.account} + attrs: dict[str, Any] = {"duration": target.walltime_min * 60, "account": target.account} if target.queue_name: - attrs['queue_name'] = target.queue_name - custom_attrs: Dict[str, Any] = {} + attrs["queue_name"] = target.queue_name + custom_attrs: dict[str, Any] = {} if target.constraint: - custom_attrs[f'{executor}.constraint'] = target.constraint + custom_attrs[f"{executor}.constraint"] = target.constraint - env = {'RADICAL_ORBIT_BROKER_URL': broker_url} + env = {"RADICAL_ORBIT_BROKER_URL": broker_url} if target.setup: - env['RADICAL_ORBIT_ENDPOINT_SETUP'] = '; '.join(target.setup) + env["RADICAL_ORBIT_ENDPOINT_SETUP"] = "; ".join(target.setup) - child_name = f'rose-{edge_name}-{uuid.uuid4().hex[:6]}' + child_name = f"rose-{edge_name}-{uuid.uuid4().hex[:6]}" job_spec = { - 'executable' : _wrapper_path(home), - 'arguments' : ['--name', child_name, '--url', broker_url], - 'attributes' : attrs, - 'custom_attributes' : custom_attrs, - 'resources' : {'node_count': target.n_nodes, 'process_count': 1}, - 'environment' : env, + "executable": _wrapper_path(home), + "arguments": ["--name", child_name, "--url", broker_url], + "attributes": attrs, + "custom_attributes": custom_attrs, + "resources": {"node_count": target.n_nodes, "process_count": 1}, + "environment": env, } - print(f'[rose --remote] submitting psij job via {edge_name} ' - f'(executor: {executor}, endpoint: {child_name})...') + print( + f"[rose --remote] submitting psij job via {edge_name} " + f"(executor: {executor}, endpoint: {child_name})..." + ) res = psij.submit_tunneled(job_spec, executor=executor, tunnel=target.tunnel) print(f"[rose --remote] psij job_id: {res['job_id']}") return { - 'kind' : 'psij', - 'client' : psij, - 'endpoint_name': res.get('endpoint_name', child_name), - 'job_id' : res['job_id'], - 'parent_edge' : edge_name, + "kind": "psij", + "client": psij, + "endpoint_name": res.get("endpoint_name", child_name), + "job_id": res["job_id"], + "parent_edge": edge_name, } @@ -197,10 +201,12 @@ def _launch_psij(rt, target: 'TargetConfig', broker_url: str, # Wait for the bootstrapped endpoint to register # ───────────────────────────────────────────────────────────────────────────── -def _wait_for_endpoint(rt, name: str, timeout: float = 30 * 60, - poll: float = 3.0, heartbeat: float = 30.0) -> None: - print(f'[rose --remote] waiting for endpoint {name!r} to register...') - start = time.time() + +def _wait_for_endpoint( + rt, name: str, timeout: float = 30 * 60, poll: float = 3.0, heartbeat: float = 30.0 +) -> None: + print(f"[rose --remote] waiting for endpoint {name!r} to register...") + start = time.time() last_beat = start while time.time() - start < timeout: if name in rt.topology(): @@ -208,47 +214,47 @@ def _wait_for_endpoint(rt, name: str, timeout: float = 30 * 60, time.sleep(poll) if time.time() - last_beat >= heartbeat: elapsed = int(time.time() - start) - print(f'[rose --remote] ...{elapsed}s elapsed, ' - f'{int(timeout - elapsed)}s left') + print(f"[rose --remote] ...{elapsed}s elapsed, {int(timeout - elapsed)}s left") last_beat = time.time() - raise TimeoutError(f'endpoint {name!r} did not register within {timeout}s') + raise TimeoutError(f"endpoint {name!r} did not register within {timeout}s") # ───────────────────────────────────────────────────────────────────────────── # Teardown # ───────────────────────────────────────────────────────────────────────────── -def _teardown(created: Optional[Dict[str, Any]]) -> None: + +def _teardown(created: dict[str, Any] | None) -> None: if not created: return - print('[rose --remote] tearing down bootstrap job...') + print("[rose --remote] tearing down bootstrap job...") try: - if created['kind'] == 'psij': - created['client'].cancel_job(created['job_id']) + if created["kind"] == "psij": + created["client"].cancel_job(created["job_id"]) else: - created['client'].cancel_job(created['resource_id'], created['job_id']) + created["client"].cancel_job(created["resource_id"], created["job_id"]) print(f"[rose --remote] cancelled {created['kind']} job {created['job_id']}") except Exception as exc: - print(f"[rose --remote] could not cancel {created['kind']} " - f"job {created['job_id']}: {exc}") + print(f"[rose --remote] could not cancel {created['kind']} job {created['job_id']}: {exc}") - if created['kind'] in ('iri', 'sfapi'): + if created["kind"] in ("iri", "sfapi"): try: - created['connect_client'].disconnect(created['endpoint_key']) - print(f"[rose --remote] disconnected {created['kind']} " - f"endpoint {created['endpoint_key']}") + created["connect_client"].disconnect(created["endpoint_key"]) + print( + f"[rose --remote] disconnected {created['kind']} endpoint {created['endpoint_key']}" + ) except Exception as exc: - print(f"[rose --remote] could not disconnect " - f"{created['endpoint_key']}: {exc}") + print(f"[rose --remote] could not disconnect {created['endpoint_key']}: {exc}") # ───────────────────────────────────────────────────────────────────────────── # Entry point # ───────────────────────────────────────────────────────────────────────────── -async def run_remote(spec: 'WorkflowSpec') -> None: - """Bootstrap the ``remote.target`` from *spec*, run the workflow against - it, then tear the bootstrap job down. + +async def run_remote(spec: WorkflowSpec) -> None: + """Bootstrap the ``remote.target`` from *spec*, run the workflow against it, then tear the + bootstrap job down. Raises ``RuntimeError`` if ``spec.config.remote.target`` is unset — the CLI is expected to check this and produce a clearer error before calling @@ -263,10 +269,9 @@ async def run_remote(spec: 'WorkflowSpec') -> None: from radical.orbit import EndpointRuntime remote_cfg = spec.config.remote - target = remote_cfg.target + target = remote_cfg.target if target is None: - raise RuntimeError( - "run_remote() requires 'remote.target' to be set in the spec") + raise RuntimeError("run_remote() requires 'remote.target' to be set in the spec") eb = None if remote_cfg.embedded: @@ -274,7 +279,7 @@ async def run_remote(spec: 'WorkflowSpec') -> None: eb = EmbeddedBroker(plugins=_PLUGIN_FOR_KIND[target.kind]) eb.start() - print(f'[rose --remote] embedded broker: {eb.url}') + print(f"[rose --remote] embedded broker: {eb.url}") try: rt = EndpointRuntime(broker_url=eb.url) rt.start(wait=True) @@ -286,28 +291,28 @@ async def run_remote(spec: 'WorkflowSpec') -> None: rt.start(wait=True) broker_url = rt.broker_url - print(f'[rose --remote] broker: {broker_url}') + print(f"[rose --remote] broker: {broker_url}") - created: Optional[Dict[str, Any]] = None + created: dict[str, Any] | None = None try: - if target.kind == 'psij': + if target.kind == "psij": if not remote_cfg.embedded and target.edge_name not in rt.topology(): raise RuntimeError( - f'remote.target.edge_name {target.edge_name!r} is not ' - f'currently connected to the broker') - created = _launch_psij(rt, target, broker_url, - embedded=remote_cfg.embedded) + f"remote.target.edge_name {target.edge_name!r} is not " + f"currently connected to the broker" + ) + created = _launch_psij(rt, target, broker_url, embedded=remote_cfg.embedded) else: created = _launch_iri_or_sfapi(rt, target, broker_url) - _wait_for_endpoint(rt, created['endpoint_name']) + _wait_for_endpoint(rt, created["endpoint_name"]) print(f"[rose --remote] endpoint {created['endpoint_name']!r} is up") - await spec.workflow(broker_url, created['endpoint_name']) + await spec.workflow(broker_url, created["endpoint_name"]) finally: _teardown(created) rt.stop() if eb is not None: eb.stop() - print('[rose --remote] done.') + print("[rose --remote] done.") diff --git a/rose/spec/schema.py b/rose/spec/schema.py index 0ea90dcb..ef1dd09d 100644 --- a/rose/spec/schema.py +++ b/rose/spec/schema.py @@ -77,8 +77,8 @@ def tasks(self) -> dict[str, TaskDef]: # ── Remote / tracking ──────────────────────────────────────────────────────── class TargetConfig(BaseModel): - """Bootstrap config for `rose run --remote`: how to launch the remote - ORBIT endpoint before the workflow's tasks run on it. + """Bootstrap config for `rose run --remote`: how to launch the remote ORBIT endpoint before the + workflow's tasks run on it. ``kind`` selects the launch mechanism: - ``iri``/``sfapi``: submit a job to `resource_id` through the broker's @@ -94,13 +94,13 @@ class TargetConfig(BaseModel): kind: Literal["iri", "sfapi", "psij"] # iri / sfapi - endpoint: str | None = None # 'nersc' | 'olcf' + endpoint: str | None = None # 'nersc' | 'olcf' resource_id: str | None = None - home_dir: str | None = None # user $HOME on target; resolves the wrapper path - login_host: str | None = None # for tunnel='forward' + home_dir: str | None = None # user $HOME on target; resolves the wrapper path + login_host: str | None = None # for tunnel='forward' # psij - edge_name: str | None = None # login-node endpoint already in the topology + edge_name: str | None = None # login-node endpoint already in the topology executor: str | None = None # shared submission attributes @@ -162,8 +162,7 @@ def _check_embedded(self) -> RemoteConfig: and not self.target.edge_name ): raise ValueError( - "remote.target.kind='psij' requires 'edge_name' " - "(unless remote.embedded is true)" + "remote.target.kind='psij' requires 'edge_name' (unless remote.embedded is true)" ) return self diff --git a/tests/unit/spec/test_schema.py b/tests/unit/spec/test_schema.py index 866cce29..9452cfb9 100644 --- a/tests/unit/spec/test_schema.py +++ b/tests/unit/spec/test_schema.py @@ -712,8 +712,8 @@ def test_remote_psij_edge_name_not_required_when_embedded(): def test_remote_psij_edge_name_still_valid_when_embedded(): - """A user may still set edge_name when embedded — harmless, ignored by - the RemoteConfig-level requirement check.""" + """A user may still set edge_name when embedded — harmless, ignored by the RemoteConfig-level + requirement check.""" cfg = RemoteConfig( embedded=True, target=TargetConfig(kind="psij", edge_name="login1", account="amsc007"), @@ -725,8 +725,11 @@ def test_remote_iri_target_unaffected_by_embedded_validator(): cfg = RemoteConfig( embedded=True, target=TargetConfig( - kind="iri", endpoint="olcf", resource_id="odo", - account="fus183", home_dir="/home/x", + kind="iri", + endpoint="olcf", + resource_id="odo", + account="fus183", + home_dir="/home/x", ), ) assert cfg.target.kind == "iri" @@ -761,10 +764,12 @@ def test_target_config_iri_missing_home_dir_raises(): def test_target_config_psij_missing_edge_name_allowed_standalone(): - """TargetConfig alone no longer enforces 'edge_name' for kind='psij' — - whether it's required depends on RemoteConfig.embedded, a sibling field - this model can't see. See RemoteConfig-level tests below for the actual - enforcement (edge_name required unless embedded=True).""" + """TargetConfig alone no longer enforces 'edge_name' for kind='psij' — whether it's required + depends on RemoteConfig.embedded, a sibling field this model can't see. + + See RemoteConfig-level tests below for the actual enforcement (edge_name required unless + embedded=True). + """ t = TargetConfig(kind="psij", account="amsc007") assert t.edge_name is None @@ -781,9 +786,7 @@ def test_target_config_missing_account_raises(): def test_target_config_extra_field_rejected(): with pytest.raises(ValidationError): - TargetConfig( - kind="psij", edge_name="login1", account="a", bogus_field="x" - ) + TargetConfig(kind="psij", edge_name="login1", account="a", bogus_field="x") def test_target_config_bad_tunnel_value_rejected(): @@ -869,8 +872,8 @@ def test_remote_embedded_psij_roundtrip_from_yaml(tmp_path): def test_remote_without_target_still_valid(tmp_path): - """Specs written before remote.target existed (e.g. the m3dc1 use case) - must keep validating unchanged.""" + """Specs written before remote.target existed (e.g. the m3dc1 use case) must keep validating + unchanged.""" p = tmp_path / "spec.yaml" p.write_text(SEQ_YAML) cfg = WorkflowConfig.from_yaml(p) diff --git a/tests/unit/spec/test_workflow_spec.py b/tests/unit/spec/test_workflow_spec.py index 836769d2..65c1b65b 100644 --- a/tests/unit/spec/test_workflow_spec.py +++ b/tests/unit/spec/test_workflow_spec.py @@ -3,7 +3,7 @@ import sys import textwrap import types -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import MagicMock import pytest @@ -265,7 +265,7 @@ def test_validate_imports_default_false_skips_check(tmp_path): # the rest of the execution stack (LearnerBuilder / learner.start / asyncflow). -class _Sentinel(Exception): +class _SentinelError(Exception): pass @@ -276,7 +276,7 @@ async def test_workflow_calls_get_backend_with_broker_url(tmp_path, monkeypatch) async def fake_get_backend(name, **kwargs): captured["name"] = name captured.update(kwargs) - raise _Sentinel + raise _SentinelError fake_rhapsody = types.ModuleType("rhapsody") fake_rhapsody.get_backend = fake_get_backend @@ -287,7 +287,7 @@ async def fake_get_backend(name, **kwargs): monkeypatch.setitem(sys.modules, "rhapsody", fake_rhapsody) monkeypatch.setitem(sys.modules, "radical.asyncflow", fake_asyncflow_mod) - with pytest.raises(_Sentinel): + with pytest.raises(_SentinelError): await spec.workflow("https://broker.example:8000", "ep-1") assert captured["name"] == "orbit" diff --git a/tests/unit/test_remote.py b/tests/unit/test_remote.py index 501d50be..04520346 100644 --- a/tests/unit/test_remote.py +++ b/tests/unit/test_remote.py @@ -1,6 +1,7 @@ -"""Unit tests for rose.remote — bootstrap/execute/teardown orchestration for -`rose run --remote`. radical.orbit's EndpointRuntime is mocked throughout; -these tests never touch a real broker or HPC system. +"""Unit tests for rose.remote — bootstrap/execute/teardown orchestration for `rose run --remote`. + +radical.orbit's EndpointRuntime is mocked throughout; these tests never touch a real broker or HPC +system. """ import textwrap @@ -158,8 +159,12 @@ def test_teardown_iri_cancels_and_disconnects(): instance = MagicMock() connect_client = MagicMock() created = { - "kind": "sfapi", "client": instance, "connect_client": connect_client, - "endpoint_key": "nersc", "resource_id": "perlmutter", "job_id": "j1", + "kind": "sfapi", + "client": instance, + "connect_client": connect_client, + "endpoint_key": "nersc", + "resource_id": "perlmutter", + "job_id": "j1", } remote._teardown(created) instance.cancel_job.assert_called_once_with("perlmutter", "j1") @@ -168,8 +173,13 @@ def test_teardown_iri_cancels_and_disconnects(): def test_teardown_psij_cancels_only(): psij = MagicMock() - created = {"kind": "psij", "client": psij, "job_id": "j1", - "parent_edge": "login1", "endpoint_name": "rose-login1-abc"} + created = { + "kind": "psij", + "client": psij, + "job_id": "j1", + "parent_edge": "login1", + "endpoint_name": "rose-login1-abc", + } remote._teardown(created) psij.cancel_job.assert_called_once_with("j1") @@ -179,8 +189,12 @@ def test_teardown_swallows_cancel_errors(): instance.cancel_job.side_effect = RuntimeError("already gone") connect_client = MagicMock() created = { - "kind": "iri", "client": instance, "connect_client": connect_client, - "endpoint_key": "olcf", "resource_id": "odo", "job_id": "j1", + "kind": "iri", + "client": instance, + "connect_client": connect_client, + "endpoint_key": "olcf", + "resource_id": "odo", + "job_id": "j1", } remote._teardown(created) # must not raise connect_client.disconnect.assert_called_once_with("olcf") @@ -230,6 +244,7 @@ async def test_run_remote_psij_edge_not_in_topology_raises(tmp_path, monkeypatch monkeypatch.setattr(remote, "_teardown", MagicMock()) import radical.orbit + monkeypatch.setattr(radical.orbit, "EndpointRuntime", MagicMock(return_value=fake_rt)) with pytest.raises(RuntimeError, match="login1"): @@ -260,13 +275,13 @@ async def test_run_remote_sfapi_end_to_end(tmp_path, monkeypatch): monkeypatch.setattr(type(spec), "workflow", property(lambda self: fake_workflow)) import radical.orbit + monkeypatch.setattr(radical.orbit, "EndpointRuntime", MagicMock(return_value=fake_rt)) await remote.run_remote(spec) fake_rt.get_plugin.assert_called_once_with("broker", "sfapi_connect") - connect_client.connect.assert_called_once_with("nersc", "abc123", - "-----BEGIN KEY-----\n...") + connect_client.connect.assert_called_once_with("nersc", "abc123", "-----BEGIN KEY-----\n...") instance.submit_job.assert_called_once() resource_id, job_spec = instance.submit_job.call_args[0] assert resource_id == "perlmutter" @@ -292,12 +307,14 @@ def test_launch_psij_embedded_uses_broker_and_local_home(monkeypatch): rt.get_plugin.return_value = psij from rose.spec.schema import TargetConfig + target = TargetConfig(kind="psij", account="amsc007") class _FakeHomePath: @staticmethod def home(): from pathlib import Path as _RealPath + return _RealPath("/home/local") # Rebind the name `Path` inside rose.remote only — must not touch the @@ -319,6 +336,7 @@ def test_launch_psij_embedded_prefers_explicit_home_dir(): rt.get_plugin.return_value = psij from rose.spec.schema import TargetConfig + target = TargetConfig(kind="psij", account="amsc007", home_dir="/custom/home") remote._launch_psij(rt, target, "https://broker:8000", embedded=True) @@ -337,6 +355,7 @@ def test_launch_psij_non_embedded_uses_edge_and_remote_sysinfo(): rt.get_plugin.side_effect = lambda edge, name: {"psij": psij, "sysinfo": sysinfo}[name] from rose.spec.schema import TargetConfig + target = TargetConfig(kind="psij", edge_name="login1", account="amsc007") remote._launch_psij(rt, target, "https://broker:8000", embedded=False) @@ -374,6 +393,7 @@ async def test_run_remote_embedded_constructs_and_tears_down_broker(tmp_path, mo import radical.orbit import radical.orbit.embedded + monkeypatch.setattr(radical.orbit, "EndpointRuntime", MagicMock(return_value=fake_rt)) fake_embedded_broker_cls = MagicMock(return_value=fake_eb) monkeypatch.setattr(radical.orbit.embedded, "EmbeddedBroker", fake_embedded_broker_cls) @@ -400,9 +420,9 @@ async def test_run_remote_embedded_stops_broker_if_runtime_start_fails(tmp_path, import radical.orbit import radical.orbit.embedded + monkeypatch.setattr(radical.orbit, "EndpointRuntime", MagicMock(return_value=fake_rt)) - monkeypatch.setattr(radical.orbit.embedded, "EmbeddedBroker", - MagicMock(return_value=fake_eb)) + monkeypatch.setattr(radical.orbit.embedded, "EmbeddedBroker", MagicMock(return_value=fake_eb)) with pytest.raises(RuntimeError, match="registration failed"): await remote.run_remote(spec) From 3a47c551b7a7f0a4ead0bc73f57fd19cf5f755ba Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Wed, 19 Aug 2026 01:37:01 +0200 Subject: [PATCH 16/19] ruff --- tests/unit/test_learner_stop.py | 2 +- tests/unit/test_parallel_learner.py | 2 +- tests/unit/test_rl_par_learner.py | 4 ++-- tests/unit/test_rl_seq_learner.py | 2 +- tests/unit/test_sequential_learner.py | 2 +- tests/unit/test_uq_learner.py | 4 ++-- 6 files changed, 8 insertions(+), 8 deletions(-) diff --git a/tests/unit/test_learner_stop.py b/tests/unit/test_learner_stop.py index 9c18ef70..c2935fe2 100644 --- a/tests/unit/test_learner_stop.py +++ b/tests/unit/test_learner_stop.py @@ -3,8 +3,8 @@ import pytest from radical.asyncflow import WorkflowEngine - from rose.al.active_learner import ParallelActiveLearner, SequentialActiveLearner + from rose.learner import IterationState diff --git a/tests/unit/test_parallel_learner.py b/tests/unit/test_parallel_learner.py index 3bc7c524..96db285a 100644 --- a/tests/unit/test_parallel_learner.py +++ b/tests/unit/test_parallel_learner.py @@ -4,12 +4,12 @@ # Assuming these imports based on the code structure from radical.asyncflow import WorkflowEngine - from rose.al.active_learner import ( LearnerConfig, ParallelActiveLearner, SequentialActiveLearner, ) + from rose.learner import IterationState diff --git a/tests/unit/test_rl_par_learner.py b/tests/unit/test_rl_par_learner.py index d1978748..6010fe99 100644 --- a/tests/unit/test_rl_par_learner.py +++ b/tests/unit/test_rl_par_learner.py @@ -2,13 +2,13 @@ import pytest from radical.asyncflow import WorkflowEngine - -from rose.learner import IterationState, LearnerConfig from rose.rl.reinforcement_learner import ( ParallelReinforcementLearner, SequentialReinforcementLearner, ) +from rose.learner import IterationState, LearnerConfig + class TestParallelReinforcementLearner: """Test cases for ParallelReinforcementLearner class.""" diff --git a/tests/unit/test_rl_seq_learner.py b/tests/unit/test_rl_seq_learner.py index 6eb99339..cd3c1b12 100644 --- a/tests/unit/test_rl_seq_learner.py +++ b/tests/unit/test_rl_seq_learner.py @@ -2,9 +2,9 @@ import pytest from radical.asyncflow import WorkflowEngine +from rose.rl.reinforcement_learner import SequentialReinforcementLearner from rose.learner import TaskConfig -from rose.rl.reinforcement_learner import SequentialReinforcementLearner @pytest.mark.filterwarnings("ignore::RuntimeWarning") diff --git a/tests/unit/test_sequential_learner.py b/tests/unit/test_sequential_learner.py index 4956cc55..938d5206 100644 --- a/tests/unit/test_sequential_learner.py +++ b/tests/unit/test_sequential_learner.py @@ -4,8 +4,8 @@ # Assuming these imports based on the code structure from radical.asyncflow import WorkflowEngine - from rose.al.active_learner import SequentialActiveLearner + from rose.learner import TaskConfig diff --git a/tests/unit/test_uq_learner.py b/tests/unit/test_uq_learner.py index d02398aa..4acb6cbb 100644 --- a/tests/unit/test_uq_learner.py +++ b/tests/unit/test_uq_learner.py @@ -5,12 +5,12 @@ # Assuming these imports based on the code structure from radical.asyncflow import WorkflowEngine - -from rose.learner import IterationState from rose.uq import UQScorer from rose.uq.uq_active_learner import ParallelUQLearner, SeqUQLearner from rose.uq.uq_learner import UQLearnerConfig +from rose.learner import IterationState + async def mock_start_iterator(*args, **kwargs): """Helper to mock an async iterator.""" From 0a374f80b0518c3f5bfb2df7510875d85f653aea Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Wed, 19 Aug 2026 02:55:51 +0200 Subject: [PATCH 17/19] adding pre-commit to tox --- tox.ini | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/tox.ini b/tox.ini index a9b97835..4ccbf787 100644 --- a/tox.ini +++ b/tox.ini @@ -1,14 +1,25 @@ [tox] -envlist = py310,py311,py312,py313 +envlist = pre-commit,py310,py311,py312,py313 isolated_build = true +# Pre-commit must pass before any other env runs — same recipe as +# radical.asyncflow's tox.ini: this runs the identical `pre-commit run +# --all-files` CI's pre-commit job runs, so failures show up locally +# (whole repo, not just files you touched) instead of only in CI. +[testenv:pre-commit] +extras = dev +allowlist_externals = pre-commit +commands = pre-commit run --all-files + # Unit test env [testenv] +depends = pre-commit extras = dev commands = pytest tests/unit {posargs} # Integration test [testenv:{py310,py311,py312,py313}-all] +depends = pre-commit extras = dev setenv = RADICAL_VERBOSE=DEBUG @@ -30,5 +41,6 @@ commands = # Coverage [testenv:coverage] +depends = pre-commit extras = dev commands = pytest --cov=rose --cov-report=html --cov-report=term {posargs} From 1ae6633a344f6cf43af1c9fd49cce69386a2a814 Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Wed, 19 Aug 2026 03:54:18 +0200 Subject: [PATCH 18/19] another attempt to fix pre-commit --- pyproject.toml | 9 +++++++++ tests/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 +- tests/integration/test_run_sequential_learner.py | 2 +- tests/integration/test_run_uq_learner.py | 4 ++-- tests/unit/test_learner_stop.py | 2 +- tests/unit/test_parallel_learner.py | 2 +- tests/unit/test_rl_par_learner.py | 4 ++-- tests/unit/test_rl_seq_learner.py | 2 +- tests/unit/test_sequential_learner.py | 2 +- tests/unit/test_uq_learner.py | 4 ++-- 12 files changed, 23 insertions(+), 14 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index f404a4b7..3147aff3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -77,6 +77,15 @@ select = ["E", "F", "W", "B", "I", "N", "UP"] # N812/N813/N816/N817: allow ML import aliases (functional as F, etc.) ignore = ["N803", "N806", "N801", "N812", "N813", "N816", "N817"] +[tool.ruff.lint.isort] +# Without this, ruff's I001 (import sort) falls back to auto-detecting +# first-party packages, which is environment-dependent (e.g. whether `rose` +# happens to be resolvable/installed in the linting environment) — the exact +# cause of local pre-commit runs and CI's isolated hook venv disagreeing on +# whether `rose.*` imports belong in their own group. Pinning it explicitly +# makes the classification deterministic everywhere. +known-first-party = ["rose"] + [tool.ruff.lint.per-file-ignores] # task_description={"shell": True} is a ROSE API pattern introspected by decorators "examples/**/run_me.py" = ["B006"] diff --git a/tests/integration/test_run_parallel_learner.py b/tests/integration/test_run_parallel_learner.py index 8e7c76d1..c180ccc9 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.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 0abc4bc6..ad244013 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 16372a63..f61f4495 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 62c2029e..d7a75efd 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.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 8bcf63d4..86533e3f 100644 --- a/tests/integration/test_run_uq_learner.py +++ b/tests/integration/test_run_uq_learner.py @@ -5,10 +5,10 @@ import pytest from radical.asyncflow import WorkflowEngine from rhapsody.backends import ConcurrentExecutionBackend -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 +from rose.uq import UQ_REGISTRY, register_uq +from rose.uq.uq_active_learner import ParallelUQLearner, SeqUQLearner @pytest.mark.asyncio diff --git a/tests/unit/test_learner_stop.py b/tests/unit/test_learner_stop.py index c2935fe2..9c18ef70 100644 --- a/tests/unit/test_learner_stop.py +++ b/tests/unit/test_learner_stop.py @@ -3,8 +3,8 @@ import pytest from radical.asyncflow import WorkflowEngine -from rose.al.active_learner import ParallelActiveLearner, SequentialActiveLearner +from rose.al.active_learner import ParallelActiveLearner, SequentialActiveLearner from rose.learner import IterationState diff --git a/tests/unit/test_parallel_learner.py b/tests/unit/test_parallel_learner.py index 96db285a..3bc7c524 100644 --- a/tests/unit/test_parallel_learner.py +++ b/tests/unit/test_parallel_learner.py @@ -4,12 +4,12 @@ # Assuming these imports based on the code structure from radical.asyncflow import WorkflowEngine + from rose.al.active_learner import ( LearnerConfig, ParallelActiveLearner, SequentialActiveLearner, ) - from rose.learner import IterationState diff --git a/tests/unit/test_rl_par_learner.py b/tests/unit/test_rl_par_learner.py index 6010fe99..d1978748 100644 --- a/tests/unit/test_rl_par_learner.py +++ b/tests/unit/test_rl_par_learner.py @@ -2,13 +2,13 @@ import pytest from radical.asyncflow import WorkflowEngine + +from rose.learner import IterationState, LearnerConfig from rose.rl.reinforcement_learner import ( ParallelReinforcementLearner, SequentialReinforcementLearner, ) -from rose.learner import IterationState, LearnerConfig - class TestParallelReinforcementLearner: """Test cases for ParallelReinforcementLearner class.""" diff --git a/tests/unit/test_rl_seq_learner.py b/tests/unit/test_rl_seq_learner.py index cd3c1b12..6eb99339 100644 --- a/tests/unit/test_rl_seq_learner.py +++ b/tests/unit/test_rl_seq_learner.py @@ -2,9 +2,9 @@ import pytest from radical.asyncflow import WorkflowEngine -from rose.rl.reinforcement_learner import SequentialReinforcementLearner from rose.learner import TaskConfig +from rose.rl.reinforcement_learner import SequentialReinforcementLearner @pytest.mark.filterwarnings("ignore::RuntimeWarning") diff --git a/tests/unit/test_sequential_learner.py b/tests/unit/test_sequential_learner.py index 938d5206..4956cc55 100644 --- a/tests/unit/test_sequential_learner.py +++ b/tests/unit/test_sequential_learner.py @@ -4,8 +4,8 @@ # Assuming these imports based on the code structure from radical.asyncflow import WorkflowEngine -from rose.al.active_learner import SequentialActiveLearner +from rose.al.active_learner import SequentialActiveLearner from rose.learner import TaskConfig diff --git a/tests/unit/test_uq_learner.py b/tests/unit/test_uq_learner.py index 4acb6cbb..d02398aa 100644 --- a/tests/unit/test_uq_learner.py +++ b/tests/unit/test_uq_learner.py @@ -5,12 +5,12 @@ # Assuming these imports based on the code structure from radical.asyncflow import WorkflowEngine + +from rose.learner import IterationState from rose.uq import UQScorer from rose.uq.uq_active_learner import ParallelUQLearner, SeqUQLearner from rose.uq.uq_learner import UQLearnerConfig -from rose.learner import IterationState - async def mock_start_iterator(*args, **kwargs): """Helper to mock an async iterator.""" From d6ff2feb0c153eda371949f914fe0f1138b384c4 Mon Sep 17 00:00:00 2001 From: AymenFJA Date: Wed, 19 Aug 2026 15:02:12 +0200 Subject: [PATCH 19/19] adding orbit as dep --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 3147aff3..45abf411 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,6 +20,7 @@ requires-python = ">=3.10" dependencies = [ "numpy", "radical.asyncflow", + "radical.orbit", "rhapsody-py[radical_pilot]", 'rhapsody-py[dragon]; python_version >= "3.10" and python_version <= "3.12"', ]