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/_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/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/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..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 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 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. @@ -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/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/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/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..66c46fb5 100644 --- a/docs/user-guide/advanced-rl-workflow.md +++ b/docs/user-guide/advanced-rl-workflow.md @@ -1,40 +1,28 @@ +# 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: ```python from radical.asyncflow import WorkflowEngine -from rhapsody.backends import RadicalExecutionBackend -from rose.rl.reinforcement_learner import SequentialReinforcementLearner +from rhapsody.backends import DragonExecutionBackendV3 +from rose.rl import SequentialReinforcementLearner ``` @@ -44,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) @@ -80,8 +65,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..18db108e 100644 --- a/docs/user-guide/basic-acl-workflow.md +++ b/docs/user-guide/basic-acl-workflow.md @@ -3,17 +3,19 @@ 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 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) @@ -38,8 +40,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 +94,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..7e1329ca 100644 --- a/docs/user-guide/basic-rl-workflow.md +++ b/docs/user-guide/basic-rl-workflow.md @@ -1,17 +1,22 @@ +# Basic Reinforcement Learning Workflow + ## Define your target machine to run on 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 +from rose.rl 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) @@ -33,7 +38,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..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) @@ -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..e70e1060 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 @@ -79,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 new file mode 100644 index 00000000..c12c6981 --- /dev/null +++ b/docs/user-guide/rose-aas.md @@ -0,0 +1,84 @@ +# 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 **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: + +- 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 | +|---|---|---|---| +| `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, `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`). + +--- + +## How ROSE operates within a job + +The orbit backend connects two things over the ORBIT broker: + +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 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. + +--- + +## 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..ebf9b664 --- /dev/null +++ b/docs/user-guide/spec-api.md @@ -0,0 +1,754 @@ +# 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. + +--- + + +!!! note + 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. + + +## 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 WorkflowConfig object +``` + +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. + +--- + +## 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 + +
+
+ +**Python API** + +```python +from rose.al 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.builder import LearnerBuilder + +builder = LearnerBuilder("workflow.yaml", asyncflow) +cfg = builder.config +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 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 import ParallelActiveLearner + +acl = ParallelActiveLearner(asyncflow) + +@acl.simulation_task(as_executable=False) +async def simulate(*args, **kwargs): + ... + +@acl.training_task(as_executable=False) +async def train(sim_result, **kwargs): + ... + +@acl.active_learn_task(as_executable=False) +async def active_learn(sim_result, model, **kwargs): + ... + +@acl.as_stop_criterion( + metric_name="mse", + threshold=0.01, +) +async def check_mse(*args, **kwargs): + ... + +async for state in acl.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 +builder = LearnerBuilder("workflow.yaml", asyncflow) +cfg = builder.config +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 +acl = ParallelActiveLearner(asyncflow) + +# Manual routing by learner_id +TRAIN_CMD = { + 0: "python train.py --model linear", + 1: "python train.py --model ridge", +} + +@acl.simulation_task +async def simulate(*args, **kwargs): + return "python sim.py" + +@acl.training_task +async def train(*args, **kwargs): + return TRAIN_CMD[kwargs["learner_id"]] + +@acl.active_learn_task +async def active_learn(*args, **kwargs): + return "python active_learn.py" + +@acl.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 acl.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 +builder = LearnerBuilder("workflow.yaml", asyncflow) +cfg = builder.config +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 must be async. The return value is passed as the first positional argument to the next task in the chain. + +
+
+ +**`type: shell`** + +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: + 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. + +
+
+ +### `task_description` example + +
+
+ +**Python API** + +```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: + 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. + +--- + +## 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 + 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 +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. + +### `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 | +|---|---|---| +| `dask` | `DaskExecutionBackend` | Dask-based execution | +| `concurrent` | `ConcurrentExecutionBackend` | CPU-only runs, no Dragon needed | +| `dragon` | `DragonExecutionBackendV3` | For highly compute intensive surrogates | + + + +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`. + +### `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 + +```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 | +| `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 | + + +### 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/docs/user-guide/target-resources.md b/docs/user-guide/target-resources.md index eb6add8c..84e3decd 100644 --- a/docs/user-guide/target-resources.md +++ b/docs/user-guide/target-resources.md @@ -1,19 +1,19 @@ # 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: ```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 +from rose.al import SequentialActiveLearner -engine = await RadicalExecutionBackend( - {'runtime': 30, - 'resource': 'local.localhost'}) +engine = await ConcurrentExecutionBackend(ProcessPoolExecutor()) asyncflow = await WorkflowEngine.create(engine) @@ -21,20 +21,24 @@ acl = SequentialActiveLearner(asyncflow) ``` ## HPC Resources -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): + +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): + + +!!! 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 +from rose.al 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/docs/user-guide/uq_based-acl-workflow.md b/docs/user-guide/uq_based-acl-workflow.md index 24cfc90f..e9a8e5e7 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' ``` @@ -94,8 +94,8 @@ 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 -learner = ParallelUQLearner(asyncflow) +from rose.uq import ParallelUQLearner +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 c35fd6a4..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 @@ -66,7 +63,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 @@ -75,6 +76,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 @@ -86,15 +89,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 + - 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 diff --git a/pyproject.toml b/pyproject.toml index e8fcca05..45abf411 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,10 +20,14 @@ 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"', ] +[project.scripts] +rose = "rose.cli:main" + [project.urls] Homepage = "https://github.com/radical-cybertools/ROSE" Issues = "https://github.com/radical-cybertools/ROSE/issues" @@ -36,6 +40,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", @@ -72,6 +78,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/rose/__init__.py b/rose/__init__.py index 0038267b..8a48e1d6 100644 --- a/rose/__init__.py +++ b/rose/__init__.py @@ -1,11 +1,38 @@ -from rose.al import active_learner, selector +import sys as _sys + 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 + +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 + +_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", @@ -21,4 +48,7 @@ # Tracking "TrackerBase", "PipelineManifest", + # YAML spec layer + "load_spec", + "WorkflowSpec", ] 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..c4181219 --- /dev/null +++ b/rose/cli.py @@ -0,0 +1,93 @@ +#!/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 _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") + + 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( + "--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="dragon_v3", + help="Rhapsody backend name for --local (default: concurrent)", + ) + + args = parser.parse_args() + + if args.command == "run": + 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/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/remote.py b/rose/remote.py new file mode 100644 index 00000000..937023d7 --- /dev/null +++ b/rose/remote.py @@ -0,0 +1,318 @@ +"""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 + +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, {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: dict[str, Any] | None) -> 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']} 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']} endpoint {created['endpoint_key']}" + ) + except Exception as 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. + + 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: dict[str, Any] | None = 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/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/__init__.py b/rose/spec/__init__.py new file mode 100644 index 00000000..dbd282fe --- /dev/null +++ b/rose/spec/__init__.py @@ -0,0 +1,155 @@ +from __future__ import annotations + +from collections.abc import Callable, Coroutine +from pathlib import Path +from typing import Any + +from .builder import LearnerBuilder +from .schema import WorkflowConfig + +__all__ = ["load_spec", "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: + 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(WorkflowConfig.from_yaml(path)) + if validate_imports: + _validate_task_imports(spec.config) + return spec + + +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 ld in cfg.learners: + for td in ld.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: WorkflowConfig) -> 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: + """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 + + 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 ``WorkflowConfig`` 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(WorkflowConfig.model_validate(data)) + + @property + def workflow(self) -> Callable[..., Coroutine[Any, Any, None]]: + cfg = self.config + + async def _workflow(broker_url: str, endpoint_name: str) -> None: + import rhapsody + from radical.asyncflow import WorkflowEngine + + engine = await rhapsody.get_backend( + "orbit", + broker_url=broker_url, + endpoint_name=endpoint_name, + backends=cfg.remote.backends, + ) + 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": + 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): + 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..e8910108 --- /dev/null +++ b/rose/spec/adapters.py @@ -0,0 +1,111 @@ +from __future__ import annotations + +from collections.abc import Callable + +from .schema import RemoteConfig, TaskDef + + +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, td) + return _make_python_closure(task_def.function, remote.pythonpath, td) + + @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-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( + {i: tdi.command for i, tdi in enumerate(task_defs)}, slot_name, td + ) + return _make_python_dispatch( + {i: tdi.function for i, tdi in enumerate(task_defs)}, + list(remote.pythonpath), + slot_name, + td, + ) + + +def _make_shell_closure(command: str, task_description: dict) -> Callable: + _cmd = command + _td = task_description + + async def _task(*args, task_description=_td, **kwargs) -> str: + return _cmd.format_map(kwargs) + + _task.__name__ = "shell_task" + return _task + + +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, task_description=_td, **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, task_description: dict) -> Callable: + _cmds = dict(cmds) + _td = task_description + + async def _dispatch(*args, task_description=_td, **kwargs) -> str: + cmd = _cmds[kwargs.get("learner_id", 0)] + return cmd.format_map(kwargs) + + _dispatch.__name__ = f"{slot_name}_dispatch" + return _dispatch + + +def _make_python_dispatch( + 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, task_description=_td, **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..471b5fae --- /dev/null +++ b/rose/spec/builder.py @@ -0,0 +1,162 @@ +from __future__ import annotations + +from pathlib import Path + +from rose.learner import Learner + +from .adapters import TaskAdapterFactory +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. +_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": + from rose.active_learning.active_learner import SequentialActiveLearner + + return SequentialActiveLearner + if learner_type == "parallel_active_learner": + from rose.active_learning.active_learner import ParallelActiveLearner + + return ParallelActiveLearner + if learner_type == "sequential_reinforcement_learner": + from rose.reinforcement_learning.reinforcement_learner import SequentialReinforcementLearner + + return SequentialReinforcementLearner + if learner_type == "uq_active_learner": + from rose.uncertainty_quantification.uq_active_learner import SeqUQLearner + + return SeqUQLearner + raise ValueError(f"No learner class registered for type '{learner_type}'") + + +class LearnerBuilder: + 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: + cfg = self.config + learner = _get_learner_class(cfg.learner.type)(self.asyncflow) + + if cfg.learners 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: 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: WorkflowConfig) -> None: + required = _REQUIRED_SLOTS[cfg.learner.type] + for slot_name in required: + 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) + + 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): + """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.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) + params["pythonpath"] = list(cfg.remote.pythonpath) + 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 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}) + configs.append( + LearnerConfig(**{slot: schedule for slot in lc_slots}, 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..ef1dd09d --- /dev/null +++ b/rose/spec/schema.py @@ -0,0 +1,266 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Any, 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" + 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: + 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 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 + 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 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" + experiment: str = "ROSE-Spec" + run_name: str | None = None + + 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 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. + 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) + + 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) -> WorkflowConfig: + ltype = self.learner.type + required = _REQUIRED_SLOTS.get(ltype) + if required is None: + raise ValueError( + 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 ld in self.learners: + missing = required - set(ld.tasks.keys()) + if missing: + raise ValueError(f"Learner '{ld.label}' missing: {sorted(missing)}") + extra = set(ld.tasks.keys()) - required + if extra: + raise ValueError(f"Learner '{ld.label}' unexpected fields: {sorted(extra)}") + for slot in required: + 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(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. " + "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()) + 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 + + @model_validator(mode="after") + def _validate_parameters(self) -> WorkflowConfig: + 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) -> WorkflowConfig: + import yaml + + raw = yaml.safe_load(Path(path).read_text()) + return cls.model_validate(raw) 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", -] diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py new file mode 100644 index 00000000..e69de29b 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..45c60d70 --- /dev/null +++ b/tests/integration/spec/helpers.py @@ -0,0 +1,40 @@ +"""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] + + +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 + + +# ── 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 new file mode 100644 index 00000000..7ce139b3 --- /dev/null +++ b/tests/integration/spec/test_yaml_workflow.py @@ -0,0 +1,217 @@ +"""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 WorkflowConfig + +# ── 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 = WorkflowConfig.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) + + +# ── 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 = WorkflowConfig.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 = WorkflowConfig.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/__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..c6608c27 --- /dev/null +++ b/tests/unit/spec/test_adapters.py @@ -0,0 +1,168 @@ +"""Unit tests for rose.spec.adapters β€” closure factories, no HPC needed.""" + +import asyncio +import inspect + +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" + + +# ── 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} + + +# ── 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}), + 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_builder.py b/tests/unit/spec/test_builder.py new file mode 100644 index 00000000..9a893293 --- /dev/null +++ b/tests/unit/spec/test_builder.py @@ -0,0 +1,310 @@ +"""Unit tests for rose.spec.builder β€” LearnerConfig construction, no HPC needed.""" + +import textwrap +from unittest.mock import MagicMock + +from rose.spec.schema import WorkflowConfig + +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 = 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 + + +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 + + +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("""\ + 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 new file mode 100644 index 00000000..9452cfb9 --- /dev/null +++ b/tests/unit/spec/test_schema.py @@ -0,0 +1,918 @@ +"""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 RemoteConfig, TargetConfig, TaskDef, WorkflowConfig + +# ── 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(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} + ) + 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 = WorkflowConfig.from_yaml(p) + assert cfg.simulation.task_description == {"cpu_count": 8} + assert cfg.training.task_description is None + + +# ── WorkflowConfig β€” 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 = WorkflowConfig.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"): + WorkflowConfig.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"): + WorkflowConfig.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"): + WorkflowConfig.from_yaml(p) + + +# ── WorkflowConfig β€” parallel learner with learners ────────────────────────────── + +PAR_YAML = textwrap.dedent("""\ + learner: + type: parallel_active_learner + max_iter: 3 + + learners: + - 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_learners_roundtrip(tmp_path): + p = tmp_path / "par.yaml" + p.write_text(PAR_YAML) + 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" + + +def test_parallel_learner_missing_slot(tmp_path): + yaml = textwrap.dedent("""\ + learner: + type: parallel_active_learner + learners: + - 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"): + WorkflowConfig.from_yaml(p) + + +# ── WorkflowConfig β€” parameters block ──────────────────────────────────────────── + +# ── WorkflowConfig β€” reserved parameter keys ────────────────────────────────────── + +# ── WorkflowConfig β€” 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 = WorkflowConfig.from_yaml(p) + 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: + 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"): + WorkflowConfig.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"): + 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: + 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"): + WorkflowConfig.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"): + WorkflowConfig.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"): + WorkflowConfig.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 = WorkflowConfig.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 = 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 = WorkflowConfig.from_yaml(p) + assert cfg.parameters == {} + + +# ── 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: + type: parallel_active_learner + learners: + - 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"): + WorkflowConfig.from_yaml(p) diff --git a/tests/unit/spec/test_workflow_spec.py b/tests/unit/spec/test_workflow_spec.py new file mode 100644 index 00000000..65c1b65b --- /dev/null +++ b/tests/unit/spec/test_workflow_spec.py @@ -0,0 +1,297 @@ +"""Unit tests for WorkflowSpec β€” load_spec and workflow_with().""" + +import sys +import textwrap +import types +from unittest.mock import MagicMock + +import pytest + +from rose.spec import WorkflowSpec, load_spec + +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 + + +# ── 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 _SentinelError(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 _SentinelError + + 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(_SentinelError): + 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..04520346 --- /dev/null +++ b/tests/unit/test_remote.py @@ -0,0 +1,431 @@ +"""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() 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}