diff --git a/dev_tools/agent/skills/checks/frontmatter.py b/dev_tools/agent/skills/checks/frontmatter.py index 93356c7066..0c5c2601f5 100644 --- a/dev_tools/agent/skills/checks/frontmatter.py +++ b/dev_tools/agent/skills/checks/frontmatter.py @@ -34,6 +34,7 @@ "submits_production", } ) +NVFLARE_METADATA_FIELDS = ("min_flare_version", "max_flare_version", "blast_radius", "skill_version", "status") @dataclass(frozen=True) @@ -83,6 +84,25 @@ def parse_skill_frontmatter(skill_file: Path | str) -> dict[str, Any]: return metadata +def normalize_skill_metadata(frontmatter: Mapping[str, Any]) -> dict[str, Any]: + """Return NVFLARE metadata from spec-compliant or legacy frontmatter.""" + normalized = dict(frontmatter) + extension = frontmatter.get("metadata") + if extension is None: + return normalized + if not isinstance(extension, dict): + raise SkillFrontmatterError("SKILL.md frontmatter field 'metadata' must be a mapping") + for field in NVFLARE_METADATA_FIELDS: + nested_value = extension.get(field) + if nested_value is None: + continue + legacy_value = frontmatter.get(field) + if legacy_value is not None and legacy_value != nested_value: + raise SkillFrontmatterError(f"conflicting values for frontmatter field '{field}'") + normalized[field] = nested_value + return normalized + + def validate_skill_dir(skill_dir: Path | str) -> SkillValidationResult: """Validate one guide-compatible skill directory.""" path = Path(skill_dir) @@ -104,7 +124,7 @@ def validate_skill_dir(skill_dir: Path | str) -> SkillValidationResult: return SkillValidationResult(str(path), metadata, tuple(issues)) try: - metadata = parse_skill_frontmatter(skill_file) + metadata = normalize_skill_metadata(parse_skill_frontmatter(skill_file)) except SkillFrontmatterError as e: issues.append(_issue("skill-frontmatter-invalid", str(e), skill_file)) return SkillValidationResult(str(path), metadata, tuple(issues)) diff --git a/docs/design/autofl_skill.md b/docs/design/autofl_skill.md new file mode 100644 index 0000000000..ab8e5b9dd0 --- /dev/null +++ b/docs/design/autofl_skill.md @@ -0,0 +1,199 @@ +# NVFlare Auto-FL Skill Design + +## Summary + +Auto-FL should enter NVFlare as a skill-first product experience. Users select +an official NVFlare Auto-FL skill in a coding agent, point it at an existing +`job.py`, and state the optimization objective, environment, and budget. NVFlare +owns deterministic import of campaign-relevant settings, execution truth, policy +boundaries, artifacts, and reproducibility. The agent owns candidate planning, +code edits within allowed paths, experiment execution through the existing +`job.py`, comparison, and narrative reporting. + +This avoids introducing a new public Auto-FL command tree while still making +Auto-FL an NVFlare-owned feature. + +## Product Boundary + +The first production-oriented slice includes: + +- A root `skills/nvflare-autofl` agent skill that follows the NVFLARE skills + layout used by the general agent-skills work. +- A deterministic, skill-private `job.py` importer that emits reviewable + `autofl.yaml` for the Auto-FL campaign without executing user code. +- A trust contract in `autofl.yaml` showing editable campaign settings, + unresolved fields, fixed-budget constraints, and allowed edit paths. +- A skill-local candidate lifecycle that snapshots the current best source, + gives the agent an isolated draft, validates the resulting patch, and keeps or + restores source according to the campaign metric. +- Documentation for using the skill with simulation, POC, and production + environments through existing NVFlare surfaces. + +The first version does not embed or vendor a coding agent, and it does not add a +public Auto-FL command family. Users install and select the skill, then express +their intent in a prompt; they do not invoke its helper scripts directly. + +## Role of autofl.yaml + +`autofl.yaml` is not a replacement for `job.py` and is not a second exported +job format. The original `job.py` remains the experiment entry point the agent +uses to run candidates, and exported job folders remain the NVFlare execution +and submission artifacts. + +The purpose of `autofl.yaml` is to expose the human-reviewable Auto-FL campaign +layer: + +- Objective metric, requested environment, and candidate budget. +- Editable search-space settings discovered from `job.py` and related train + scripts. +- Fixed-budget constraints that must remain comparable across candidates. +- Allowed edit paths and files that are out of scope for the agent. +- Existing preferred source targets declared by task-local + `mutation_schema.yaml`, once they resolve inside the job workspace. +- Allowed creation patterns for new Python modules under the job root. +- Artifact, ledger, and report locations for the campaign. +- Provenance and unresolved fields that need user review before safe execution. + +By default, users should not need to edit `autofl.yaml`. They review or modify +it only when the importer surfaces unresolved settings or when they want to +override campaign knobs explicitly. + +## Deterministic Import + +The importer parses Python source with `ast`; it does not import or execute +user code. It resolves direct imports, import aliases, and module aliases for +known Recipe surfaces and NVFlare-distributed classes whose names end in +`Job`. The generic API `Job`, local subclasses, and non-NVFlare subclasses stay +explicitly unresolved. The importer focuses +on campaign-relevant settings rather than duplicating the full exported job: + +- Recipe/FedJob constructor and class import. +- `SimEnv`, `PocEnv`, and `ProdEnv` references. +- `train_script` resolution for literal and argparse-derived values, or for one + unambiguous NVFlare `ScriptRunner(script=...)` call. +- Objective metric from user request, `key_metric`, or explicit unresolved + default. +- Fixed-budget fields such as rounds, clients, and candidate budget. +- Common argparse tunables from `job.py` and the resolved train script. + +The exported job folder remains useful as execution truth once the job is +materialized, because it contains resolved NVFlare app and component configs. +However, it does not reliably preserve all authoring intent needed for an +Auto-FL campaign, such as editable source files, train-argument construction, +tunable-versus-fixed intent, and source provenance. Therefore the importer uses +deterministic Python/static parsing for the campaign layer and may use exported +config inspection as a validation aid when available. + +Unsupported or dynamic fields are carried forward as unresolved review items +instead of being guessed by the importer or the agent. +The runner writes this reviewable `autofl.yaml` before admission and refuses to +start a baseline when the job surface or fixed comparison budget remains +safety-critical and unresolved. + +## Trust Contract + +Every import result includes: + +- `import`: importer version, source path, source hash, support status, and + confidence. +- `job`: surface, entrypoint, train script, and call arguments with provenance. +- `objective`, `budget`, `environment`, and `search_space`. +- `trust_contract`: extracted facts, unresolved fields, allowed edit paths, and + allowed creation patterns and agent controls. This is the sole source of + candidate source permissions. + +The skill must present editable, unresolved, and allowed sections before it runs +candidates. This is the core product guardrail: NVFlare makes the campaign +reviewable and reproducible; the agent makes it interactive and exploratory. +During campaign initialization, the runner merges existing, workspace-local +`mutation_schema.yaml` `preferred_targets` into +`trust_contract.allowed_edit_paths`. Missing, symlinked, reserved, or +out-of-workspace targets remain unresolved rather than being silently +authorized. + +## Candidate Contract + +The agent, rather than the deterministic runner, owns search policy. It may +change tunables, edit the imported job's allowed source files, or implement new +client and server algorithms as Python modules. This includes creating or +editing server aggregator modules and registering them through `job.py`; the +agent is not limited to pre-enumerated FedAvg, FedAvgM, FedAdam, FedOpt, or +SCAFFOLD choices. Each attempt starts from the retained best source +in `.nvflare/autofl/candidates//source` and has a generated +`candidate_manifest.json` containing its hypothesis, base candidate, run +arguments, changed files, source and budget hashes, patch hash, artifacts, and +result. + +NVFlare computes the manifest's evidence fields; the agent does not assert them. +Before execution, the helper rejects stale candidates, path traversal, symlink +escapes, unauthorized existing-file edits, and detectable fixed-budget drift. +It applies the candidate transactionally to the real job workspace, retains a +new best, and restores the previous best after a discard or crash. This works +without requiring a Git repository and leaves the best source ready for the +standard NVFlare job lifecycle. + +The built-in parameter candidates are suggestion seeds only. They are returned +as machine-readable hypotheses and arguments when requested, but are not the +default search loop and are never executed without agent selection. +After each literature-triggered plateau, campaign state requests at least one +source-backed server aggregation candidate under the same comparison budget. +When the job contract makes that impossible, the agent records the reason in +the literature event instead of silently omitting aggregation exploration. + +## Execution Model + +The skill uses existing NVFlare execution surfaces: + +- Simulation: initialize a baseline, prepare an agent-authored candidate draft, + and evaluate it through the existing `job.py` and configured `SimEnv`. +- POC: use the existing job authoring/export flow, startup kits, and standard + `nvflare job` commands, then record the job ID, artifacts, and metric against + the candidate manifest. +- Production: use standard startup-kit authentication, site policy, job submit, + wait, download, and inspection commands with the same manifest and result + recording contract. + +Production is a valid optimization environment. The best candidate may later be +submitted or reused through the standard NVFlare job lifecycle; no separate +promotion command is needed. + +The runner is the sole writer of `.nvflare/autofl/campaign_state.json`. Its +`status` action rescans the ledger, pending manifests, stop files, and cap and +only rewrites state when its semantic contents change. It does not regenerate +the ledger, plot, or report during an unchanged status check. The standalone +campaign guard is a read-only diagnostic and cannot overwrite runner metadata. +A stop file takes immediate precedence: pending candidates must be safely +abandoned before final reporting. Without a stop request, pending prepared or +externally ready candidates take precedence over cap exhaustion and reporting. + +Every score records its metric name, extraction source, and artifact. Structured +metric artifacts take precedence over exact, provenance-labelled text fallback. +The ledger is replaced atomically so a failed write cannot leave a partial +campaign record. + +## Skill Implementation Boundary + +The deterministic importer and campaign runner live under +`skills/nvflare-autofl/scripts/`. They are private executable resources of the +Agent Skill, resolved relative to `SKILL.md` by the activated coding agent. The +human-facing workflow exposes only skill installation, selection, and an intent +prompt; it does not document direct Python helper invocation. + +This placement deliberately keeps the unreleased `autofl.yaml` contract out of +`nvflare.app_common` and the public `nvflare agent` CLI. A general NVFlare job +contract command should be considered only after another concrete workflow +needs the same interface and the schema has proved stable. The general, +read-only `nvflare agent inspect` surface does not acquire an Auto-FL-specific +profile in this proposal. + +## Review Questions + +- Are the supported `job.py` patterns sufficient for an initial prototype? +- Are the edit and creation permissions in `autofl.yaml` appropriate for + algorithm-level candidates while preserving candidate comparability? +- Which exported-job fields should be used as validation evidence versus static + `job.py` parsing for authoring intent? +- Does the Auto-FL skill pass the general NVFLARE skill frontmatter, trigger, + and eval checks after it lands under `skills/nvflare-autofl`? +- Which candidate-manifest and metric/artifact fields should become stable + NVFlare APIs after the skill-local contract proves itself? diff --git a/docs/index.rst b/docs/index.rst index 40dc298d7c..9f5616d78e 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -38,6 +38,7 @@ NVIDIA FLARE programming_guide/experiment_tracking Federated XGBoost user_guide/data_scientist_guide/data_preparation + Agent Skills CLI Tools .. toctree:: diff --git a/docs/user_guide/agent_skills/autofl_skill.rst b/docs/user_guide/agent_skills/autofl_skill.rst new file mode 100644 index 0000000000..81827f7c8c --- /dev/null +++ b/docs/user_guide/agent_skills/autofl_skill.rst @@ -0,0 +1,105 @@ +.. _autofl_skill: + +############################ +NVFlare Auto-FL Agent Skill +############################ + +The NVFlare Auto-FL agent skill optimizes an existing NVFlare ``job.py`` +through a coding agent. It is not an ``nvflare autofl`` command and does not +add an Auto-FL command family to the NVFlare CLI. + +Install the NVFlare-owned skill for the coding agent you use. For Codex: + +.. code-block:: shell + + nvflare agent skills install --agent codex --skill nvflare-autofl + +For Claude Code, use ``--agent claude``. The installer places the complete +skill, including its deterministic helper scripts, in the selected agent's +managed skill directory. + +User Experience +=============== + +Select the installed skill in the coding agent and provide the optimization +intent: + +.. code-block:: text + + Select: NVFlare Auto-FL skill + Prompt: Optimize ./job.py for accuracy in sim. + +An explicit candidate budget is optional. Without one, the campaign continues +until the user interrupts it or a hard safety or runtime blocker prevents +further comparable execution. + +Users do not invoke scripts from the installed skill directory. The activated +coding agent resolves and runs those bundled resources internally. Bundled +scripts are private implementation details, not public NVFlare commands or +Python APIs. + +Deterministic Import +==================== + +The skill first parses ``job.py`` without importing or executing it. Its private +importer recognizes supported NVFlare Recipe and FedJob construction patterns, +aliases, ``SimEnv`` configuration, fixed training budgets, metrics, train +scripts, and common argparse tunables. + +The result is a reviewable ``autofl.yaml`` containing: + +- the optimization metric, direction, environment, and candidate budget; +- the fixed comparison budget that candidates must preserve; +- ``trust_contract.allowed_edit_paths`` and allowed Python creation patterns; +- source and importer provenance; +- unresolved dynamic or unsupported fields requiring review. + +When the user does not name a metric, a deterministic ``key_metric`` extracted +from ``job.py`` takes precedence. The default user experience does not require +editing ``autofl.yaml``. + +Candidate Lifecycle +=================== + +The coding agent forms a hypothesis and asks the private skill runner to create +an isolated candidate source tree. The agent may edit allowed existing files or +add Python modules, including new client algorithms and server aggregators. + +For every candidate, NVFlare-owned helper code: + +- recomputes changed files rather than trusting agent-written manifest paths; +- rejects edits outside the trust contract and fixed-budget drift; +- runs simulation through the configured ``SimEnv`` or prepares standard POC + or production submission; +- records the score, metric provenance, source patch, command, artifacts, and + failure evidence; +- retains an improved candidate or restores the previous best source. + +POC and production candidates use the normal ``nvflare job submit``, ``job +wait``, and ``job download`` lifecycle with configured startup-kit policy. The +skill does not bypass authentication or site policy. + +Campaign Artifacts +================== + +The job directory contains the human-reviewable and reproducibility artifacts: + +- ``autofl.yaml``: imported campaign and trust contract; +- ``results.tsv``: atomic candidate ledger with metric provenance; +- ``progress.png``: campaign trajectory; +- ``autofl_report.md``: current campaign summary; +- ``.nvflare/autofl/campaign_state.json``: next action and stop status; +- ``.nvflare/autofl/candidates//candidate_manifest.json``: candidate + hypothesis, source hashes, changed files, result, and artifacts. + +A manual stop takes precedence over pending execution. If a candidate is +pending, the agent abandons it safely before generating the final report. + +Supported First Version +======================= + +The first version supports statically recognizable NVFlare Recipe constructors +and NVFlare-distributed classes ending in ``Job``. Generic, local, and +non-NVFlare job or recipe classes remain unresolved. Ambiguous scripts and +dynamic safety-critical comparison fields block baseline execution rather than +being guessed. diff --git a/docs/user_guide/agent_skills/index.rst b/docs/user_guide/agent_skills/index.rst new file mode 100644 index 0000000000..8b9327ab9f --- /dev/null +++ b/docs/user_guide/agent_skills/index.rst @@ -0,0 +1,13 @@ +############ +Agent Skills +############ + +NVFlare-owned agent skills package repeatable coding-agent workflows with +deterministic helpers and reviewable contracts. Install and list them through +``nvflare agent skills``; interact with the installed skill through the coding +agent rather than by invoking its private scripts directly. + +.. toctree:: + :maxdepth: 1 + + autofl_skill diff --git a/skills/nvflare-autofl/BENCHMARK.md b/skills/nvflare-autofl/BENCHMARK.md new file mode 100644 index 0000000000..e8e9174f69 --- /dev/null +++ b/skills/nvflare-autofl/BENCHMARK.md @@ -0,0 +1,24 @@ +# Benchmark Summary + +Status: draft/internal pending runtime evaluation. + +Skill version: 0.1.0 +FLARE version: 2.8.0 minimum + +## Initial Checks + +| Check | Status | Notes | +| --- | --- | --- | +| Positive trigger | Draft | `autofl-optimize-existing-job` defines the primary Auto-FL prompt. | +| Adjacent negative trigger | Draft | PyTorch conversion routes to `nvflare-convert-pytorch`. | +| Diagnosis negative trigger | Draft | Failed-job diagnosis routes to `nvflare-diagnose-job`. | +| Global negative trigger | Draft | Non-FLARE prompts route to no skill. | +| Mandatory behavior | Draft | Behavior IDs cover deterministic import, agent-authored code candidates, candidate manifests, bounded edits, and existing FLARE execution. | +| Prohibited behavior | Draft | Behavior IDs prohibit bypassing policy, editing outside allowed paths, and treating `autofl.yaml` as an exported job. | +| Process evaluation | Draft | Metrics cover import-first behavior, contract preservation, score extraction, and unwanted production actions. | + +## Known Gaps + +- Runtime agent-performance scoring has not been run yet. +- Adjacent negatives should be expanded after more NVFLARE workflow skills land. +- Production execution behavior needs site-policy fixture coverage before the skill graduates from draft status. diff --git a/skills/nvflare-autofl/SKILL.md b/skills/nvflare-autofl/SKILL.md new file mode 100644 index 0000000000..dfac62916a --- /dev/null +++ b/skills/nvflare-autofl/SKILL.md @@ -0,0 +1,199 @@ +--- +name: nvflare-autofl +description: "Optimize an existing NVFLARE job.py through an agent-assisted Auto-FL campaign that preserves FLARE execution, policy, artifacts, and reproducibility." +license: Apache-2.0 +compatibility: "Requires NVFLARE 2.8.0+, Python, and permission to run NVFLARE jobs in the selected environment." +metadata: + author: "nvflare" + min_flare_version: "2.8.0" + blast_radius: submits_production + category: Optimization + skill_version: "0.1.0" +--- + +# NVFLARE Auto-FL + +## Use When +Use this skill when the user asks to optimize an existing NVFLARE `job.py` for +accuracy, AUC, loss, runtime, robustness, or another metric in simulation, POC, +or production. + +## Do Not Use When +Do not use for converting non-FL training code into NVFLARE, diagnosing failed +jobs without an optimization goal, production deployment setup, or generic +hyperparameter tuning outside an NVFLARE job. + +## Workflow +Use this skill to optimize an existing NVFLARE `job.py` without asking the user +to learn a new Auto-FL command tree. The user selects this skill, points to a +job, and states the objective, environment, and optional budget. NVFLARE +provides the deterministic campaign import, execution substrate, policy +boundaries, artifacts, and machine-readable contracts. The coding agent owns +hypotheses, source edits, new algorithm implementations, and candidate choice. + +Resolve [run_job_campaign.py](scripts/run_job_campaign.py) relative to this +`SKILL.md`, store its absolute path as `RUNNER`, and initialize the campaign: + +```bash +python "$RUNNER" initialize ./job.py [--metric ] --mode --env [--max-candidates ] +``` + +Read `autofl.yaml` and the JSON response, then prepare an agent-authored +candidate with a short hypothesis and optional candidate-only arguments: + +```bash +python "$RUNNER" prepare ./job.py --name --hypothesis "" [--run-args ""] +``` + +Edit only the returned candidate source directory. Modify existing allowed +files or add Python modules under the job root; do not edit the live best source +directly. Then evaluate the manifest: + +```bash +python "$RUNNER" evaluate ./job.py --manifest +``` + +Simulation evaluation runs the candidate immediately. POC and production +evaluation validates and materializes the candidate; submit it with standard +`nvflare job` commands, then call `record` with the manifest, job ID, artifacts, +and score. Use `abandon` to restore a pending candidate. Use `suggest` only when +you want deterministic tunable seeds; suggestions are never executed +automatically and do not limit agent-authored code candidates. + +If the job directory contains a task-local `mutation_schema.yaml`, treat its +`comparison_budget_args.default_candidate_budget` and mutation bounds as +authoritative. Invalid generated proposals are product friction, not campaign +blockers; preserve the same campaign and continue with another same-budget +candidate. + +Treat `job.py`, generated `autofl.yaml`, and optional job-local +`mutation_schema.yaml` as campaign inputs. Do not require example-specific +runbooks, branches, or initialization scripts. + +Request escalated execution for the runner command because NVFLARE simulator +runs create local sockets that fail inside the restricted Codex sandbox. If a +runner command reports a sandbox/socket permission failure, treat it as an +infrastructure retry, not as a candidate result, and rerun the same command with +escalated execution. + +The helper owns deterministic import, source snapshots, candidate validation, +execution, restoration, counting, ledger updates, campaign state, plotting, +and reports. After each lifecycle action, read +`.nvflare/autofl/campaign_state.json` and only finalize when +`final_response_allowed=true`. For long-running and simulator-stall handling, read +[continuous-campaigns.md](references/continuous-campaigns.md). +For comparison budgets, data distributions, and rerun evidence, read +[experiment comparability](references/experiment-comparability.md). + +Read `autofl.yaml` and show the user a concise campaign summary: + +- **Editable**: metric, environment, candidate budget, tunables, artifact + locations, `objective.optimization_metric`, metric source, source hash, and importer version. +- **Unresolved**: dynamic defaults, unsupported Python semantics, missing + metric sources, unknown data paths, or any low-confidence fields. +- **Allowed**: files the agent may edit, Python source it may create, + fixed-budget fields it must preserve, and policy boundaries for the requested + environment. + +Treat `autofl.yaml` as the human-reviewable Auto-FL campaign config, not as a +replacement for `job.py` or an exported NVFLARE job folder. Use the original +`job.py` as the runnable experiment entry point throughout the candidate loop. + +If `autofl.yaml` contains unresolved fields that affect execution safety, +candidate comparability, or production submission, ask the user to resolve those +specific fields before running candidates. + +## Requirements +- Edit existing files only through candidate drafts and within + `trust_contract.allowed_edit_paths`. New Python modules may match + `trust_contract.allowed_create_patterns` under the job root. +- Use existing `mutation_schema.yaml` `preferred_targets` only after the runner + reflects them in the trust contract; surface unresolved targets. +- You may create and register new Python server aggregators through `job.py`; + do not limit exploration to existing FedAvg/FedAvgM/FedAdam/FedOpt/SCAFFOLD choices. +- Preserve `budget.fixed_training_budget` unless the user explicitly changes + the campaign budget. +- If the environment provides `PYTHON`, `VIRTUAL_ENV`, or a venv on `PATH`, + treat that prepared runtime as authoritative. Verify it, then use it for + import, validation, execution, metric extraction, plotting, and reporting. + Do not search for alternate interpreters or install dependencies unless the + user explicitly asks you to prepare the environment. +- Treat generated `autofl.yaml`, task-local `mutation_schema.yaml`, and + existing NVFLARE job/runtime configuration as authoritative. In the default + simulation product flow, do not require task-local prose profiles, special + branch setup, or harness initialization before invoking the runner. +- Use NVFLARE's existing execution surfaces: + - For simulation, run the imported job with its configured `SimEnv`. + - For POC and production, use standard `nvflare job submit`, `job wait`, + `job download`, and related job/status commands. +- Record every candidate in a ledger such as `results.tsv` with a short name, + changed files, diff summary, run command, metric result, artifacts, and + failure reason when applicable. +- Prefer small, reviewable edits over broad rewrites. +- Treat production as an available execution environment, but never bypass + startup-kit authentication, site policy, or normal NVFLARE job submission. + +## Candidate Loop +1. Inspect `autofl.yaml`, current best source, prior manifests, and results. +2. Form a concrete hypothesis. Use literature, framework knowledge, source + edits, new client or server algorithms, or a fallback tunable suggestion as + appropriate. +3. Prepare a candidate, edit its draft, and evaluate its manifest. +4. Let the helper validate paths and fixed-budget comparability, compute the + patch hash, execute or materialize it, extract metrics, and keep or restore. +5. Read campaign state and execute `next_action`. Run a source-backed literature + pass when requested, then implement its strongest compatible idea. + +## Continuous Campaign Rule +For uncapped campaigns, continue proposing and evaluating same-budget candidates +after setup and baseline until manually interrupted. Do not ask whether to keep +going or finalize while campaign state says `final_response_allowed=false`. + +A kept improvement, refreshed plot, updated report, local commit, first plateau +check, or encoded `job.py` default is a checkpoint, not completion. Treat the +campaign state as authoritative: if `final_response_allowed=false`, execute +`next_action` and keep the same `job.py`, `autofl.yaml`, metric, environment, +ledger, and comparison budget. Use +[continuous-campaigns.md](references/continuous-campaigns.md) for simulator +watchdogs, campaign-state handling, and recovery rules. + +## Candidate Caps + +Campaigns are uncapped by default. If the user says "optimize this job" without +an explicit candidate budget, continue the campaign until manually interrupted +or blocked. Do not stop after the first baseline, first batch, first successful +candidate, first kept improvement, first local commit, or first plateau +checkpoint. Do not stop after a first sweep of tunables; broaden into +agent-authored code or literature-derived algorithm candidates. Uncapped progress +reports must not ask whether to continue; continue unless the +user explicitly interrupts or the code-owned state permits finalization. +Do not invent a replacement campaign or new objective after a recoverable +failure. Keep the current campaign identity and artifacts coherent unless the +human explicitly requests a new campaign. + +If the user provides an `N`-candidate budget, pass it only through the runner's +explicit `--max-candidates` argument and count up to `N` comparable attempts +after baseline. Never infer a cap from an inherited environment variable. Do +not count import, validation, smoke runs, plotting, reporting, baseline, or +infrastructure-only retries. Count a real candidate crash after execution +starts. State must report `candidate_cap_source=explicit` or `uncapped`. + +Treat plateau as a decision checkpoint, not an automatic stop. Summarize the +plateau in the running report, refresh `progress.png`, run the runner's `status` +action to refresh `.nvflare/autofl/campaign_state.json`, choose the returned +next mode, and continue unless the state reports `final_response_allowed=true`. +Use `campaign_guard.py` only for read-only ledger diagnostics; it never updates +authoritative campaign state. +After a source-backed review, record it with the helper's `record --literature +--hypothesis ""` action before preparing its candidate. +After every literature-triggered plateau, evaluate at least one source-backed +server aggregation candidate, or record why the job contract is incompatible. + +## Stop Handling + +Only produce a final answer for a campaign when the code-owned campaign state +reports `final_response_allowed=true`, for example because the user manually +stopped it, an explicit cap is exhausted, production policy blocks execution, or +a hard safety/runtime blocker prevents further comparable runs. Then finalize +`results.tsv`, `progress.png`, and a concise report with baseline, best score, +metric source, failures, friction, commands, and absolute artifact paths. diff --git a/skills/nvflare-autofl/evals/evals.json b/skills/nvflare-autofl/evals/evals.json new file mode 100644 index 0000000000..e7dfb55e1a --- /dev/null +++ b/skills/nvflare-autofl/evals/evals.json @@ -0,0 +1,131 @@ +{ + "skill_name": "nvflare-autofl", + "evals": [ + { + "id": "autofl-optimize-existing-job", + "prompt": "Use NVFLARE Auto-FL to optimize ./job.py for validation accuracy in simulation with an 8-candidate budget.", + "expected_output": "The agent invokes the NVFLARE Auto-FL skill, imports job.py into autofl.yaml, summarizes editable/unresolved/allowed campaign settings, then prepares and evaluates agent-authored source or tunable candidates through candidate manifests and existing FLARE execution surfaces.", + "files": [], + "assertions": [ + "The agent generates autofl.yaml before editing files.", + "The agent summarizes editable settings, unresolved fields, fixed-budget constraints, and allowed edit paths.", + "The agent forms hypotheses and edits isolated candidate source rather than limiting the campaign to built-in tunable sweeps.", + "The agent records changed files, patch hash, base candidate, and artifacts in candidate_manifest.json.", + "The agent runs candidates using the existing job.py rather than replacing it with autofl.yaml.", + "The agent records candidate results and reports the best reproducible candidate." + ], + "nvflare": { + "expected_skill": "nvflare-autofl", + "mandatory_behavior": [ + { + "id": "deterministic-import-first", + "description": "runs deterministic job import before candidate edits" + }, + { + "id": "campaign-summary-before-runs", + "description": "shows editable, unresolved, allowed, and fixed-budget campaign fields before running candidates" + }, + { + "id": "bounded-edit-surface", + "description": "edits only files allowed by autofl.yaml" + }, + { + "id": "agent-authored-code-candidates", + "description": "lets the coding agent implement source and algorithm candidates in an isolated draft" + }, + { + "id": "candidate-manifest-provenance", + "description": "computes a manifest with base source, changed files, patch hash, budget hash, artifacts, and result" + }, + { + "id": "existing-flare-execution", + "description": "uses job.py and standard FLARE execution surfaces for candidate runs" + } + ], + "prohibited_behavior": [ + { + "id": "no-autofl-yaml-as-exported-job", + "description": "does not treat autofl.yaml as a replacement for job.py or an exported FLARE job" + }, + { + "id": "no-policy-bypass", + "description": "does not bypass startup-kit authentication, site policy, or normal job submission for production runs" + }, + { + "id": "no-out-of-scope-edits", + "description": "does not edit outside the allowed edit paths" + }, + { + "id": "no-default-generated-search-policy", + "description": "does not let built-in tunable suggestions replace agent candidate planning" + } + ], + "process_metrics": [ + { + "id": "import_before_edit", + "description": "whether deterministic import occurs before candidate code edits" + }, + { + "id": "fixed_budget_violation_count", + "description": "number of candidate runs that changed fixed-budget fields without user approval" + }, + { + "id": "metric_extraction_success", + "description": "whether the requested metric is extracted from NVFLARE artifacts or logs" + }, + { + "id": "candidate_manifest_completeness", + "description": "whether code candidates include deterministic source, budget, patch, and artifact provenance" + }, + { + "id": "unwanted_production_action_count", + "description": "number of production submissions or policy-sensitive actions performed without explicit user context" + }, + { + "id": "report_reproducibility", + "description": "whether the final report includes commands, artifacts, changed files, and the best candidate" + } + ] + } + }, + { + "id": "autofl-negative-pytorch-conversion", + "prompt": "Convert this standalone PyTorch training script into an NVFLARE federated job.", + "expected_output": "The Auto-FL skill should not be the lead; a conversion skill should handle the request before Auto-FL is applicable.", + "files": [], + "assertions": [ + "The selected skill is nvflare-convert-pytorch, not nvflare-autofl." + ], + "nvflare": { + "expected_skill": "nvflare-convert-pytorch", + "negative_for": "nvflare-autofl" + } + }, + { + "id": "autofl-negative-diagnose-job", + "prompt": "My NVFLARE job failed with EXECUTION_EXCEPTION. Diagnose the client logs and tell me what went wrong.", + "expected_output": "The Auto-FL skill should not be the lead; failed-job diagnosis should route to nvflare-diagnose-job.", + "files": [], + "assertions": [ + "The selected skill is nvflare-diagnose-job, not nvflare-autofl." + ], + "nvflare": { + "expected_skill": "nvflare-diagnose-job", + "negative_for": "nvflare-autofl" + } + }, + { + "id": "autofl-global-negative-web-app", + "prompt": "Build a React landing page for a bakery.", + "expected_output": "No FLARE skill should trigger.", + "files": [], + "assertions": [ + "The selected skill is none." + ], + "nvflare": { + "expected_skill": null, + "negative_for": "*" + } + } + ] +} diff --git a/skills/nvflare-autofl/references/continuous-campaigns.md b/skills/nvflare-autofl/references/continuous-campaigns.md new file mode 100644 index 0000000000..5be0dd4351 --- /dev/null +++ b/skills/nvflare-autofl/references/continuous-campaigns.md @@ -0,0 +1,101 @@ +# Continuous Campaigns + +Use this reference when an Auto-FL campaign is uncapped, long-running, or +recovering from a simulator stall. The top-level skill owns the interaction +contract; this file carries the operational detail that keeps the campaign from +prematurely stopping. + +## Lifecycle State + +Each `scripts/run_job_campaign.py` lifecycle action exits with a JSON envelope; +the campaign continues through `.nvflare/autofl/campaign_state.json`. In +uncapped mode, a completed action, current best, plot, report, or exhausted local +tunable sweep is a checkpoint only. Execute `next_action` while +`final_response_allowed=false`. + +A prepared manifest is pending work. Edit its candidate source and evaluate it, +or abandon it explicitly; do not silently start another candidate. Invalid +drafts are product friction to repair and reevaluate, not a reason to terminate +the campaign. + +During a long `evaluate` action, monitor the process plus +`autofl_runs//run.log`. A live process with no final ledger row is a +running candidate. If logs are temporarily quiet but CPU or GPU use and the +child process remain active, keep waiting. + +## Campaign Guards + +The product runner is the sole writer of +`.nvflare/autofl/campaign_state.json`; read that state before any final +response. `scripts/campaign_guard.py` is a read-only ledger diagnostic. To +rescan pending manifests, stop files, caps, and the ledger and refresh +authoritative state, run the runner's `status` action. If state has +`final_response_allowed=false`, execute `next_action` immediately; the skill +text is only the interaction layer. + +Common next actions: + +- `repair_baseline`: fix the deterministic import, execution, or metric issue, + then retry initialization. +- `edit_candidate`: finish the pending candidate draft, then invoke the runner's + `evaluate` lifecycle action. +- `abandon_candidate`: abandon pending candidate work after a manual stop, then + refresh status before final reporting. +- `propose_candidate`: form a hypothesis, prepare its manifest, and edit the + returned candidate source directory. +- `submit_baseline` or `submit_candidate`: use the standard POC/production job + lifecycle, then call `record` with its job ID and artifacts. +- `rerun_with_escalated_execution`: retry the same lifecycle action after + repairing runtime permissions; do not count it as a candidate. +- `run_literature_loop`: run a short source-backed literature pass, record a + non-scored `literature` row when a ledger is available, then launch the next + compatible same-budget candidates. Include at least one source-backed server + aggregation candidate; if that is incompatible with the job contract, + record the reason in the literature event. +- `final_report`: generate final artifacts only after the runner permits a final + response. + +After every finalized batch, run the available plateau or progress watchdog when +the task provides one. If it recommends `continue`, refresh `progress.png` and +keep iterating locally. If it recommends a literature or exploration mode, record +that decision in the ledger/report, refresh `progress.png`, and launch the top +compatible candidate batch next. If no non-duplicate safe local axis remains, +switch mode rather than stopping: broaden the search within `autofl.yaml`, run a +literature-inspired proposal pass, implement a compatible algorithm change, or +request deterministic tunable suggestions as seeds. + +Server aggregation is an open code-search surface. The agent may create a new +Python aggregator module, edit an existing allowed aggregator module, and +register it through `job.py`; it must not limit exploration to the job's +pre-existing FedAvg, FedAvgM, FedAdam, FedOpt, or SCAFFOLD options. + +## Simulator Recovery + +For NVFLARE simulator runs, the server log can be quiet after it dispatches a +round while individual clients are still training. Before declaring a stall, +inspect the active simulator directory under `/tmp/nvflare/simulation/` and +check `site-*/log.txt` or `site-*/log_fl.txt` for epoch, finished-training, +download, or task-completion progress. If any client log or server aggregation +marker advances within the expected candidate runtime, continue the same +candidate; do not stop the runner, final-answer, or start a new campaign. + +If the active NVFLARE simulator logs show a hard child-process connection +failure such as +`Failed to create connection to the child process in SimulatorClientRunner`, or +if a dispatched simulator round has no advancing server/client progress markers +past the configured no-progress watchdog timeout, do not start a new campaign and +do not produce a final report. Treat only that active candidate as crashed, +terminate the stuck `job.py` child if the runner has not already done so, +preserve the same `job.py`, `autofl.yaml`, metric, environment, ledger, and +comparison budget, then continue the same campaign. + +The product runner includes these simulator-stall watchdogs; prefer letting it +record the crash row, refresh artifacts, and launch the next candidate. For +legitimately long quiet tasks, raise `--simulator-no-progress-timeout`, set +`AUTOFL_SIMULATOR_NO_PROGRESS_TIMEOUT_SECONDS`, or set +`simulator_no_progress_timeout_seconds` in the task profile. + +If the current process has already stopped but the user did not ask to stop, do +not leave the campaign in a terminal state. Inspect the campaign state and +ledger, fix the recoverable cause, and continue the same optimization with the +same `job.py`, `autofl.yaml`, metric, environment, and comparison budget. diff --git a/skills/nvflare-autofl/references/experiment-comparability.md b/skills/nvflare-autofl/references/experiment-comparability.md new file mode 100644 index 0000000000..04b2ddb3f1 --- /dev/null +++ b/skills/nvflare-autofl/references/experiment-comparability.md @@ -0,0 +1,113 @@ +# Auto-FL Experiment Comparability + +Use this reference for Auto-FL reruns, recipe comparisons, data replacement, +data distribution, synthetic data, and site heterogeneity requests. + +## Iterative Reruns + +When the user asks to change batch size, train args, number of rounds, +`min_clients`, site count, or recipe, update only the affected training args or +job configuration. For recipe changes, rerun +`nvflare recipe show --format json` and verify the new recipe's +parameters before editing. + +After each change, rerun local validation when possible. If an exported job was +previously produced and export remains in scope, export again so the job folder +matches the updated source. + +## Recipe Search And Accuracy Comparison + +When the user asks for the best recipe or best accuracy, first define the target +metric, validation split, maximum run budget, and compatible recipe set. Do not +promise that one recipe is best without measured evidence from comparable runs. + +Keep dataset split, number of sites, rounds, epochs, seed, and evaluation metric +comparable unless the user asks to tune them. Report a small results table with +recipe, settings, metric value, command, status, and result path. If a candidate +cannot run, report the blocker instead of silently dropping it. + +## Data Distribution Experiments + +When the user asks to compare IID and heterogeneous data splits, define the split +strategy before editing. Examples include equal random IID shards, label-skewed +non-IID shards, quantity-skewed shards, or user-provided per-site partitions. + +Keep recipe, rounds, epochs, batch size, seed, and metric comparable unless the +user asks to tune them. Do not copy private data into generated artifacts; prefer +split indices, deterministic samplers, or site-local path arguments. + +Report split strategy, per-site sample counts or label summary when available, +metric value, command, status, and result path. + +## Dataset Replacement Experiments + +When the user provides a dataset URL and asks to repeat an experiment, record +the URL, dataset name when known, version or timestamp when available, expected +download size if known, license or access constraints when visible, and local +cache/path used for validation. + +Do not hide download, preprocessing, schema, or label-mapping assumptions. Keep +recipe, site count, rounds, epochs, batch size, seed, split policy, and metric +comparable unless the user asks to tune them. If the new dataset requires a data +loader or preprocessing change, keep it scoped and report the changed files. + +Follow the project's existing data-prep structure. If it already has +`download_data.py`, `prepare_data.py`, `prepare_data.sh`, or equivalent helpers, +extend those rather than creating a parallel structure. If no helper exists and +a new one is needed, use the established NVFLARE example convention of separate +download and prepare/split steps, and keep download paths, cache paths, and +per-site output directories explicit. + +Use hello-world examples as the first convention reference for new helpers: +`examples/hello-world/hello-lr/download_data.py`, +`examples/hello-world/hello-lr/prepare_data.py`, +`examples/hello-world/hello-jax/prepare_data.py`, and shell-based examples such +as `examples/hello-world/hello-cyclic/prepare_data.sh`. + +## Synthetic Per-Site Data + +When the user asks for synthetic data per site, add a deterministic data +generation step only after the expected data schema is clear from the model +input, transforms, loss function, and data loader, or from a user-provided data +generation spec. + +If modality, schema, label semantics, target distribution, missingness/noise, +site heterogeneity, expected metric, or generation library is not clear, ask the +user for a data generation spec or an approved generator/library before creating +data. The spec should name modality, shape or columns, label or target +definition, class balance or value distributions, missing-data and noise +assumptions, per-site distribution differences, sample counts, seed, and +expected metric interpretation. + +Do not invent labels, features, class balance, missingness, site skew, or +expected accuracy. If the user supplies a generator or domain-specific synthetic +data tool, wire it into the existing data-prep flow and record the tool, version +or command, seed, and parameters used. + +Prefer extending existing `prepare_data.py`, `prepare_data.sh`, or equivalent +helpers. If a separate generator is needed, keep it under the same data-prep +structure and call it from the prepare step. + +Generated site data should be written to explicit per-site outputs that the job +can pass as site-specific data paths. Treat synthetic validation as a smoke test +of wiring and training execution unless the user provides a synthetic data spec +with meaningful expected metrics. + +## Site-Specific Training Heterogeneity + +When the user asks to simulate different site speeds or training +hyperparameters, prefer per-site arguments or per-site config in `job.py`. +Examples include per-site learning rate, batch size, local epochs, sleep/delay +for speed simulation, dataset shard, or workload size. + +Only create site-specific training scripts when arguments/config cannot express +the requested behavior. If scripts are split, keep shared model and training +helpers common and report why script splitting was necessary. Report each +site's settings, command, status, metric, and result path. + +## Common Gaps To Report + +- The source training script has side effects at import time. +- The model has non-serializable state outside framework-native model state. +- The dataset path is site-specific and cannot be validated locally. +- The job file has no export path yet. diff --git a/skills/nvflare-autofl/scripts/campaign_guard.py b/skills/nvflare-autofl/scripts/campaign_guard.py new file mode 100644 index 0000000000..c69a1b66e1 --- /dev/null +++ b/skills/nvflare-autofl/scripts/campaign_guard.py @@ -0,0 +1,377 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Diagnose the product Auto-FL campaign continuation decision. + +The skill runner executes candidates and writes ``results.tsv``. This guard +turns that ledger into a machine-readable decision. The campaign runner is the +only writer of authoritative campaign state. +""" + +from __future__ import annotations + +import argparse +import csv +import json +import math +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple + +DEFAULT_HARD_CRASH_THRESHOLD = 6 +DEFAULT_MIN_DELTA = 0.0005 +DEFAULT_PLATEAU_THRESHOLD = 8 +DEFAULT_STOP_FILES = ("STOP_AUTOFL", ".nvflare/autofl/STOP") +ATTEMPT_STATUSES = {"candidate", "keep", "discard", "crash"} +SCORED_ATTEMPT_STATUSES = {"keep", "discard"} +LITERATURE_EVENT_STATUSES = {"event", "literature", "checkpoint"} + + +def utc_now() -> str: + return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z") + + +def parse_score(value: Any) -> Optional[float]: + try: + score = float(value) + except (TypeError, ValueError): + return None + if math.isnan(score) or math.isinf(score): + return None + return score + + +def load_rows(path: Path) -> List[Dict[str, str]]: + if not path.exists(): + return [] + with path.open("r", encoding="utf-8", newline="") as f: + return list(csv.DictReader(f, delimiter="\t")) + + +def normalize_status(row: Dict[str, str]) -> str: + return (row.get("status", "") or "").strip().lower() + + +def row_text(row: Dict[str, str]) -> str: + return " ".join(str(value or "") for value in row.values()).lower() + + +def is_baseline(row: Dict[str, str]) -> bool: + return normalize_status(row) == "baseline" + + +def is_literature_event(row: Dict[str, str]) -> bool: + status = normalize_status(row) + if status in LITERATURE_EVENT_STATUSES and "literature" in row_text(row): + return True + return "literature" in (row.get("name", "") or "").lower() + + +def comparable_attempts(rows: List[Dict[str, str]]) -> List[Dict[str, str]]: + return [row for row in rows if normalize_status(row) in ATTEMPT_STATUSES and not is_baseline(row)] + + +def pending_candidates(rows: List[Dict[str, str]]) -> List[Dict[str, str]]: + return [row for row in rows if normalize_status(row) == "candidate"] + + +def scored_attempts_with_index(rows: List[Dict[str, str]]) -> List[Tuple[int, Dict[str, str], float]]: + scored = [] + for idx, row in enumerate(rows): + if normalize_status(row) not in SCORED_ATTEMPT_STATUSES or is_baseline(row): + continue + score = parse_score(row.get("score", "")) + if score is not None: + scored.append((idx, row, score)) + return scored + + +def better(new_score: float, old_score: Optional[float], mode: str, min_delta: float = 0.0) -> bool: + if old_score is None: + return True + if mode == "min": + return new_score < old_score - min_delta + return new_score > old_score + min_delta + + +def best_score(rows: List[Dict[str, str]], mode: str) -> Optional[float]: + best = None + for row in rows: + status = normalize_status(row) + retained = status == "keep" or (is_baseline(row) and status not in ATTEMPT_STATUSES) + if not retained: + continue + score = parse_score(row.get("score", "")) + if score is None: + continue + if better(score, best, mode): + best = score + return best + + +def plateau_status(rows: List[Dict[str, str]], threshold: int, min_delta: float, mode: str) -> Dict[str, Any]: + retained = [] + for idx, row in enumerate(rows): + status = normalize_status(row) + is_retained = status == "keep" or (is_baseline(row) and status not in ATTEMPT_STATUSES) + if not is_retained: + continue + score = parse_score(row.get("score", "")) + if score is not None: + retained.append((idx, row, score)) + scored_attempts = scored_attempts_with_index(rows) + if threshold <= 0 or not retained: + return { + "available": True, + "recommendation": "continue", + "scored_since_reset": 0, + "threshold": threshold, + "min_delta": min_delta, + } + + best = None + best_row_idx = -1 + best_scored_idx = -1 + best_name = "" + for scored_idx, (row_idx, row, score) in enumerate(retained): + if better(score, best, mode, min_delta): + best = score + best_row_idx = row_idx + best_scored_idx = scored_idx + best_name = row.get("name", "") + + last_literature_idx = max((idx for idx, row in enumerate(rows) if is_literature_event(row)), default=-1) + reset_row_idx = max(best_row_idx, last_literature_idx) + scored_since_reset = sum(1 for row_idx, _, _ in scored_attempts if row_idx > reset_row_idx) + recommendation = "literature" if scored_since_reset >= threshold else "continue" + return { + "available": True, + "recommendation": recommendation, + "best_name": best_name, + "best_score": best, + "best_scored_index": best_scored_idx, + "last_literature_event_index": last_literature_idx, + "min_delta": min_delta, + "reset_row_index": reset_row_idx, + "scored_since_reset": scored_since_reset, + "threshold": threshold, + } + + +def parse_max_candidates(value: Optional[str]) -> Optional[int]: + if value is None or str(value).strip() == "": + return None + try: + parsed = int(value) + except (TypeError, ValueError): + return None + if parsed <= 0: + return None + return parsed + + +def parse_max_candidates_arg(value: str) -> int: + parsed = parse_max_candidates(value) + if parsed is None: + raise argparse.ArgumentTypeError("must be a positive integer") + return parsed + + +def existing_stop_files(paths: List[str]) -> List[str]: + return [path for path in paths if Path(path).exists()] + + +def repeated_crash_blocker(attempts: List[Dict[str, str]], threshold: int) -> bool: + if threshold <= 0 or len(attempts) < threshold: + return False + return all(normalize_status(row) == "crash" for row in attempts[-threshold:]) + + +def guard_state_for_rows( + rows: List[Dict[str, str]], + *, + results_path: str = "results.tsv", + max_candidates: Optional[int] = None, + stop_files: Optional[List[str]] = None, + plateau_threshold: int = DEFAULT_PLATEAU_THRESHOLD, + min_delta: float = DEFAULT_MIN_DELTA, + hard_crash_threshold: int = DEFAULT_HARD_CRASH_THRESHOLD, + mode: str = "max", + pending_manifest_count: int = 0, +) -> Dict[str, Any]: + attempts = comparable_attempts(rows) + pending = pending_candidates(rows) + cap = max_candidates + cap_source = "explicit" if cap is not None else "uncapped" + stop_file_hits = existing_stop_files(stop_files or list(DEFAULT_STOP_FILES)) + plateau = plateau_status(rows, plateau_threshold, min_delta, mode) + + decision = "continue" + reason = "continue" + next_action = "propose_candidate" + final_response_allowed = False + + pending_count = len(pending) + pending_manifest_count + if stop_file_hits: + decision = "stop" + if pending_count: + reason = "manual_stop_pending_candidate" + next_action = "abandon_candidate" + else: + reason = "manual_stop_file" + next_action = "final_report" + final_response_allowed = True + elif pending_count: + reason = "pending_candidates" + next_action = "edit_candidate" + elif cap is not None and len(attempts) >= cap: + decision = "stop" + reason = "candidate_cap_exhausted" + next_action = "final_report" + final_response_allowed = True + elif repeated_crash_blocker(attempts, hard_crash_threshold): + decision = "stop" + reason = "hard_repeated_crash_blocker" + next_action = "final_report" + final_response_allowed = True + elif plateau.get("recommendation") == "literature": + reason = "plateau_literature" + next_action = "run_literature_loop" + + if final_response_allowed: + instruction = "Final report is allowed because authoritative campaign state reached a stop condition." + elif next_action == "abandon_candidate": + instruction = ( + "Manual stop requested. Do not execute the pending candidate. Run the runner abandon action, " + "then refresh status and generate the final report." + ) + elif next_action == "edit_candidate": + instruction = "Do not produce a final answer. Finish the pending candidate, then run the runner status action." + elif next_action == "run_literature_loop": + instruction = ( + "Do not produce a final answer. Run the literature loop, record a literature event, " + "then launch source-backed candidates under the same comparison budget. Include at least one " + "server aggregation candidate when compatible with the job contract; otherwise record why it is incompatible." + ) + else: + instruction = "Do not produce a final answer. Propose and prepare the next same-budget candidate now." + + return { + "schema_version": "nvflare.autofl.campaign_state.v1", + "updated_at": utc_now(), + "results": results_path, + "decision": decision, + "reason": reason, + "next_action": next_action, + "final_response_allowed": final_response_allowed, + "candidate_cap": cap, + "candidate_cap_source": cap_source, + "candidate_attempts": len(attempts), + "pending_candidates": pending_count, + "scored_attempts": len(scored_attempts_with_index(rows)), + "best_score": best_score(rows, mode), + "stop_files": stop_file_hits, + "plateau": plateau, + "required_exploration": "source_backed_server_aggregation" if next_action == "run_literature_loop" else None, + "agent_instruction": instruction, + } + + +def guard_state( + results_path: Path, + *, + max_candidates: Optional[int] = None, + stop_files: Optional[List[str]] = None, + plateau_threshold: int = DEFAULT_PLATEAU_THRESHOLD, + min_delta: float = DEFAULT_MIN_DELTA, + hard_crash_threshold: int = DEFAULT_HARD_CRASH_THRESHOLD, + mode: str = "max", + pending_manifest_count: int = 0, +) -> Dict[str, Any]: + return guard_state_for_rows( + load_rows(results_path), + results_path=str(results_path), + max_candidates=max_candidates, + stop_files=stop_files, + plateau_threshold=plateau_threshold, + min_delta=min_delta, + hard_crash_threshold=hard_crash_threshold, + mode=mode, + pending_manifest_count=pending_manifest_count, + ) + + +def print_text(state: Dict[str, Any]) -> None: + for key in [ + "decision", + "reason", + "next_action", + "final_response_allowed", + "candidate_cap", + "candidate_cap_source", + "candidate_attempts", + "pending_candidates", + "scored_attempts", + "best_score", + "agent_instruction", + ]: + value = state.get(key) + if isinstance(value, bool): + value = str(value).lower() + elif value is None: + value = "" + print(f"{key}={value}") + + +def main(argv: Optional[List[str]] = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("results", nargs="?", default="results.tsv") + parser.add_argument("--max-candidates", type=parse_max_candidates_arg) + parser.add_argument("--stop-file", action="append") + parser.add_argument("--plateau-threshold", type=int, default=DEFAULT_PLATEAU_THRESHOLD) + parser.add_argument("--min-delta", type=float, default=DEFAULT_MIN_DELTA) + parser.add_argument("--hard-crash-threshold", type=int, default=DEFAULT_HARD_CRASH_THRESHOLD) + parser.add_argument("--mode", choices=["max", "min"], default="max") + parser.add_argument("--format", choices=["text", "json"], default="text") + args = parser.parse_args(argv) + + if args.plateau_threshold <= 0: + raise ValueError("--plateau-threshold must be positive") + if args.min_delta < 0: + raise ValueError("--min-delta must be non-negative") + + results_path = Path(args.results).resolve() + stop_files = args.stop_file if args.stop_file is not None else list(DEFAULT_STOP_FILES) + resolved_stop_files = [ + str(path if path.is_absolute() else results_path.parent / path) for path in map(Path, stop_files) + ] + state = guard_state( + results_path, + max_candidates=args.max_candidates, + stop_files=resolved_stop_files, + plateau_threshold=args.plateau_threshold, + min_delta=args.min_delta, + hard_crash_threshold=args.hard_crash_threshold, + mode=args.mode, + ) + if args.format == "json": + print(json.dumps(state, sort_keys=True)) + else: + print_text(state) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/skills/nvflare-autofl/scripts/job_importer.py b/skills/nvflare-autofl/scripts/job_importer.py new file mode 100644 index 0000000000..d2e2ca09d6 --- /dev/null +++ b/skills/nvflare-autofl/scripts/job_importer.py @@ -0,0 +1,898 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Deterministically import supported NVFlare job scripts into ``autofl.yaml``. + +The Auto-FL skill uses this module as its trust layer. The importer parses +Python source with ``ast``; it never imports or executes the user's ``job.py``. +Supported Recipe/FedJob patterns are converted into a reviewable config, while +dynamic or unsupported fields are surfaced under ``unresolved``. +""" + +from __future__ import annotations + +import ast +import hashlib +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Dict, Iterable, List, Optional, Tuple + +import yaml + +AUTOFL_CONFIG_SCHEMA_VERSION = "nvflare.autofl.config.v1" +IMPORTER_VERSION = "nvflare-autofl-job-importer/v1" +ALLOWED_CREATE_PATTERNS = ["**/*.py"] + +SUPPORTED_ENV_NAMES = {"PocEnv", "ProdEnv", "SimEnv"} +TUNABLE_ARG_NAMES = { + "aggregation_epochs", + "alpha", + "batch_size", + "cosine_lr_eta_min_factor", + "eval_batch_size", + "epochs", + "fedopt_beta1", + "fedopt_beta2", + "fedopt_tau", + "fedproxloss_mu", + "local_epochs", + "local_train_steps", + "lr", + "max_model_params", + "model_arch", + "momentum", + "num_workers", + "server_lr", + "server_momentum", + "weight_decay", +} + + +@dataclass(frozen=True) +class ArgSpec: + """Static argparse field extracted from source.""" + + name: str + flags: Tuple[str, ...] + default: Any = None + default_source: str = "argparse_default" + default_unresolved: bool = False + value_type: Optional[str] = None + choices: Optional[List[Any]] = None + action: Optional[str] = None + + +@dataclass(frozen=True) +class ResolvedValue: + """Resolved expression value plus provenance and confidence.""" + + value: Any + source: str + confidence: str = "high" + unresolved: bool = False + + +@dataclass(frozen=True) +class CallInfo: + """Supported call found in ``job.py`` with source-local resolution context.""" + + name: str + full_name: str + keywords: Dict[str, ast.AST] + assignments: Dict[str, ast.AST] + source: str + function_name: Optional[str] = None + + +class JobImportError(ValueError): + """Raised when the importer cannot read or parse a job file.""" + + +class DeterministicJobImporter: + """Rule-based importer for supported NVFlare Recipe and FedJob scripts.""" + + def __init__(self, workspace_root: Optional[str] = None): + self.workspace_root = Path(workspace_root or ".").resolve() + + def import_job( + self, + job_path: str, + *, + metric: Optional[str] = None, + mode: str = "max", + target_env: Optional[str] = None, + max_candidates: Optional[int] = None, + ) -> Dict[str, Any]: + """Return an ``autofl.yaml``-shaped config for ``job_path``. + + Args: + job_path: Path to ``job.py`` or a directory containing ``job.py``. + metric: Optional optimization metric requested by the user. + mode: ``max`` or ``min`` objective direction. + target_env: Optional target environment, such as ``sim`` or ``prod``. + max_candidates: Optional fixed candidate budget. + + Returns: + A deterministic, YAML-serializable mapping. + """ + + if mode not in {"max", "min"}: + raise JobImportError("mode must be 'max' or 'min'") + + source_path = self._resolve_job_path(job_path) + source_text = source_path.read_text(encoding="utf-8") + try: + tree = ast.parse(source_text, filename=str(source_path)) + except SyntaxError as e: + raise JobImportError(f"failed to parse {source_path}: {e}") from e + + index = _ImportIndex.from_tree(tree, source_text) + job_call = index.first_job_call() + env_call = index.first_env_call() + train_script = self._resolve_train_script(source_path, job_call, index, source_text) + train_args = _collect_argparse_args_from_file(train_script) if train_script else {} + unresolved: List[Dict[str, str]] = [] + + if not job_call: + unsupported = index.first_unsupported_job_call() + reason = "no supported Recipe or NVFlare FedJob constructor was found" + if unsupported: + reason = ( + f"constructor {unsupported.full_name} is a local or non-NVFlare Job subclass; " + "its contract cannot be imported deterministically" + ) + unresolved.append(_unresolved("job.surface", reason)) + if not train_script: + unresolved.append(_unresolved("job.train_script", "no train_script was found or resolved")) + + metric_name, metric_source, metric_issue = self._resolve_metric(metric, job_call, index.parser_args) + objective = _objective_contract(metric_name, mode, metric_source) + if metric_issue: + unresolved.append(metric_issue) + + budget, budget_issues = self._resolve_budget(max_candidates, job_call, env_call, index.parser_args, source_text) + unresolved.extend(budget_issues) + + allowed_edit_paths = self._allowed_edit_paths(source_path, train_script) + search_space, search_issues = self._suggest_search_space(index.parser_args, train_args) + unresolved.extend(search_issues) + + job_payload = { + "source": self._display_path(source_path), + "surface": _surface_name(job_call), + "entrypoint": "main" if _has_main_entrypoint(tree) else "unresolved", + } + if job_call: + call_args, call_issues = self._resolved_call_keywords(job_call, index.parser_args, source_text) + unresolved.extend(call_issues) + if _is_recipe_call(job_call): + job_payload.update( + { + "recipe": job_call.name, + "recipe_class": job_call.full_name, + "recipe_args": call_args, + } + ) + else: + job_payload.update( + { + "fed_job": job_call.name, + "fed_job_class": job_call.full_name, + "fed_job_args": call_args, + } + ) + if train_script: + job_payload["train_script"] = self._display_path(train_script) + + environment = self._environment_profile(target_env, env_call, index.parser_args, source_text) + if env_call: + env_args, env_issues = self._resolved_call_keywords(env_call, index.parser_args, source_text) + unresolved.extend(env_issues) + environment["discovered"] = {"name": env_call.name, "args": env_args} + + config = { + "schema_version": AUTOFL_CONFIG_SCHEMA_VERSION, + "import": { + "importer_version": IMPORTER_VERSION, + "source": self._display_path(source_path), + "source_sha256": _sha256_text(source_text), + "confidence": _overall_confidence(unresolved, job_call), + "support": { + "status": "supported" if job_call else "partial", + "patterns": _support_patterns(job_call, env_call), + }, + }, + "job": job_payload, + "objective": objective, + "budget": budget, + "environment": environment, + "search_space": {"suggested": search_space}, + "artifacts": { + "collect": ["logs", "metrics", "job_config", "candidate_diff", "candidate_manifest"], + "result_root": "autofl_runs", + }, + "trust_contract": { + "extracted": _trust_extracted( + job_call, + env_call, + train_script, + budget, + objective, + search_space, + ), + "unresolved": list(unresolved), + "allowed_edit_paths": allowed_edit_paths, + "allowed_create_patterns": list(ALLOWED_CREATE_PATTERNS), + "agent_controls": { + "must_not_edit_outside_allowed_paths": True, + "must_preserve_fixed_training_budget": bool(budget.get("fixed_training_budget")), + "must_report_candidate_diffs": True, + }, + }, + "unresolved": list(unresolved), + } + return config + + def dump_yaml(self, config: Dict[str, Any]) -> str: + """Return deterministic YAML for an imported Auto-FL config.""" + + return dump_autofl_yaml(config) + + def _resolve_job_path(self, job_path: str) -> Path: + path = Path(job_path) + if not path.is_absolute(): + path = self.workspace_root / path + if path.is_dir(): + path = path / "job.py" + if not path.exists(): + raise JobImportError(f"job.py not found: {job_path}") + if not path.is_file(): + raise JobImportError(f"job path must be a file or directory containing job.py: {job_path}") + return path.resolve() + + def _display_path(self, path: Path) -> str: + try: + return path.resolve().relative_to(self.workspace_root).as_posix() + except ValueError: + return path.resolve().as_posix() + + def _resolve_train_script( + self, + source_path: Path, + job_call: Optional[CallInfo], + index: "_ImportIndex", + source_text: str, + ) -> Optional[Path]: + if not job_call: + return None + + train_script_node = job_call.keywords.get("train_script") + if train_script_node: + resolved = _resolve_value(train_script_node, job_call.assignments, index.parser_args, source_text) + value = resolved.value + if isinstance(value, str) and _is_resolved_path_string(resolved): + return _existing_path((source_path.parent / value).resolve()) + return None + + script_paths = set() + for runner_call in index.script_runner_calls: + script_node = runner_call.keywords.get("script") + if not script_node: + continue + resolved = _resolve_value(script_node, runner_call.assignments, index.parser_args, source_text) + if isinstance(resolved.value, str) and _is_resolved_path_string(resolved): + path = _existing_path((source_path.parent / resolved.value).resolve()) + if path: + script_paths.add(path) + return next(iter(script_paths)) if len(script_paths) == 1 else None + + def _resolve_metric( + self, + requested_metric: Optional[str], + job_call: Optional[CallInfo], + parser_args: Dict[str, ArgSpec], + ) -> Tuple[str, str, Optional[Dict[str, str]]]: + if requested_metric: + return requested_metric, "user_request", None + if job_call and "key_metric" in job_call.keywords: + resolved = _resolve_value(job_call.keywords["key_metric"], job_call.assignments, parser_args, "") + if isinstance(resolved.value, str) and not resolved.unresolved: + return resolved.value, resolved.source, None + return "accuracy", "default", _unresolved("objective.metric", resolved.source) + if "key_metric" in parser_args and isinstance(parser_args["key_metric"].default, str): + return parser_args["key_metric"].default, "arg:key_metric", None + return ( + "accuracy", + "default", + _unresolved("objective.metric", "metric defaulted to accuracy; validation metric source is unknown"), + ) + + def _resolve_budget( + self, + max_candidates: Optional[int], + job_call: Optional[CallInfo], + env_call: Optional[CallInfo], + parser_args: Dict[str, ArgSpec], + source_text: str, + ) -> Tuple[Dict[str, Any], List[Dict[str, str]]]: + budget: Dict[str, Any] = {} + fixed_training_budget: Dict[str, Any] = {} + unresolved = [] + if max_candidates is not None: + budget["max_candidates"] = max_candidates + + if job_call: + for output_key, job_key in (("num_rounds", "num_rounds"), ("min_clients", "min_clients")): + if job_key in job_call.keywords: + resolved = _resolve_value( + job_call.keywords[job_key], job_call.assignments, parser_args, source_text + ) + if resolved.unresolved: + unresolved.append(_unresolved(f"budget.fixed_training_budget.{output_key}", resolved.source)) + else: + fixed_training_budget[output_key] = resolved.value + if "n_clients" in job_call.keywords: + resolved = _resolve_value( + job_call.keywords["n_clients"], job_call.assignments, parser_args, source_text + ) + if resolved.unresolved: + unresolved.append(_unresolved("budget.fixed_training_budget.num_clients", resolved.source)) + else: + fixed_training_budget["num_clients"] = resolved.value + + if env_call and env_call.name == "SimEnv" and "num_clients" in env_call.keywords: + resolved = _resolve_value(env_call.keywords["num_clients"], env_call.assignments, parser_args, source_text) + if resolved.unresolved: + unresolved.append(_unresolved("budget.fixed_training_budget.num_clients", resolved.source)) + else: + fixed_training_budget["num_clients"] = resolved.value + + if fixed_training_budget: + budget["fixed_training_budget"] = fixed_training_budget + else: + unresolved.append(_unresolved("budget.fixed_training_budget", "no fixed training budget was resolved")) + return budget, unresolved + + def _environment_profile( + self, + target_env: Optional[str], + env_call: Optional[CallInfo], + parser_args: Dict[str, ArgSpec], + source_text: str, + ) -> Dict[str, Any]: + requested = target_env or (_env_name_to_profile(env_call.name) if env_call else "sim") + environment: Dict[str, Any] = {"requested": requested, "profiles": {}} + if env_call and env_call.name == "SimEnv": + sim_profile: Dict[str, Any] = {} + if "num_clients" in env_call.keywords: + resolved = _resolve_value( + env_call.keywords["num_clients"], env_call.assignments, parser_args, source_text + ) + if not resolved.unresolved: + sim_profile["num_clients"] = resolved.value + environment["profiles"]["sim"] = sim_profile + return environment + + def _allowed_edit_paths(self, source_path: Path, train_script: Optional[Path]) -> List[str]: + candidates = [source_path] + if train_script: + candidates.append(train_script) + for filename in ("model.py", "mutation_schema.yaml", "requirements.txt"): + path = source_path.parent / filename + if path.exists(): + candidates.append(path.resolve()) + return list(dict.fromkeys(self._display_path(path) for path in candidates)) + + def _suggest_search_space( + self, + job_args: Dict[str, ArgSpec], + train_args: Dict[str, ArgSpec], + ) -> Tuple[Dict[str, Any], List[Dict[str, str]]]: + suggested = {} + unresolved = [] + for arg_name, arg_spec in sorted({**job_args, **train_args}.items()): + if arg_name not in TUNABLE_ARG_NAMES: + continue + item = { + "type": arg_spec.value_type or _type_name(arg_spec.default), + "default": arg_spec.default, + "source": f"argparse:{arg_name}", + "confidence": "low" if arg_spec.default_unresolved else "high", + } + if arg_spec.choices: + item["values"] = arg_spec.choices + if arg_spec.default_unresolved: + item["default_source"] = arg_spec.default_source + item["unresolved"] = True + unresolved.append( + _unresolved( + f"search_space.suggested.{arg_name}.default", + f"default is dynamic expression: {arg_spec.default}", + ) + ) + suggested[arg_name] = item + + if not suggested: + unresolved.append(_unresolved("search_space.suggested", "no supported tunable argparse fields were found")) + return suggested, unresolved + + def _resolved_call_keywords( + self, + call_info: CallInfo, + parser_args: Dict[str, ArgSpec], + source_text: str, + ) -> Tuple[Dict[str, Any], List[Dict[str, str]]]: + resolved = {} + unresolved = [] + for key, value_node in sorted(call_info.keywords.items()): + value = _resolve_value(value_node, call_info.assignments, parser_args, source_text) + resolved[key] = { + "value": value.value, + "source": value.source, + "confidence": value.confidence, + } + if value.unresolved and key != "model": + unresolved.append(_unresolved(f"job.{call_info.name}.{key}", value.source)) + return resolved, unresolved + + +def import_job_to_autofl_config( + job_path: str, + *, + workspace_root: Optional[str] = None, + metric: Optional[str] = None, + mode: str = "max", + target_env: Optional[str] = None, + max_candidates: Optional[int] = None, +) -> Dict[str, Any]: + """Convenience wrapper for deterministic job import.""" + + importer = DeterministicJobImporter(workspace_root=workspace_root) + return importer.import_job( + job_path, + metric=metric, + mode=mode, + target_env=target_env, + max_candidates=max_candidates, + ) + + +def dump_autofl_yaml(config: Dict[str, Any]) -> str: + """Return deterministic YAML for an imported Auto-FL config.""" + + return yaml.dump(config, Dumper=_NoAliasSafeDumper, sort_keys=False) + + +class _NoAliasSafeDumper(yaml.SafeDumper): + def ignore_aliases(self, data): + return True + + +class _ImportIndex(ast.NodeVisitor): + def __init__(self, source_text: str): + self.source_text = source_text + self.imports: Dict[str, str] = {} + self.parser_args: Dict[str, ArgSpec] = {} + self.module_assignments: Dict[str, ast.AST] = {} + self._local_assignments_stack: List[Dict[str, ast.AST]] = [] + self._function_stack: List[str] = [] + self._argparse_parser_names: set[str] = set() + self._argparse_subparser_names: set[str] = set() + self.job_calls: List[CallInfo] = [] + self.unsupported_job_calls: List[CallInfo] = [] + self.script_runner_calls: List[CallInfo] = [] + self.env_calls: List[CallInfo] = [] + + @classmethod + def from_tree(cls, tree: ast.AST, source_text: str) -> "_ImportIndex": + index = cls(source_text) + index.visit(tree) + return index + + def first_job_call(self) -> Optional[CallInfo]: + return self.job_calls[0] if self.job_calls else None + + def first_env_call(self) -> Optional[CallInfo]: + return self.env_calls[0] if self.env_calls else None + + def first_unsupported_job_call(self) -> Optional[CallInfo]: + return self.unsupported_job_calls[0] if self.unsupported_job_calls else None + + def visit_Import(self, node: ast.Import) -> None: + for alias in node.names: + self.imports[alias.asname or alias.name] = alias.name + + def visit_ImportFrom(self, node: ast.ImportFrom) -> None: + if node.module: + for alias in node.names: + self.imports[alias.asname or alias.name] = f"{node.module}.{alias.name}" + + def visit_FunctionDef(self, node: ast.FunctionDef) -> None: + self._visit_function(node) + + def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> None: + self._visit_function(node) + + def _visit_function(self, node: ast.FunctionDef | ast.AsyncFunctionDef) -> None: + self._function_stack.append(node.name) + self._local_assignments_stack.append({}) + self.generic_visit(node) + self._local_assignments_stack.pop() + self._function_stack.pop() + + def visit_Assign(self, node: ast.Assign) -> None: + for target in node.targets: + if isinstance(target, ast.Name): + self._current_assignments()[target.id] = node.value + self._track_argparse_assignment(target.id, node.value) + self.generic_visit(node) + + def visit_AnnAssign(self, node: ast.AnnAssign) -> None: + if isinstance(node.target, ast.Name) and node.value: + self._current_assignments()[node.target.id] = node.value + self._track_argparse_assignment(node.target.id, node.value) + self.generic_visit(node) + + def visit_AugAssign(self, node: ast.AugAssign) -> None: + if isinstance(node.target, ast.Name): + # The result depends on runtime state. Preserve the expression so + # downstream resolution marks the value as unresolved instead of + # reusing an earlier static assignment. + self._current_assignments()[node.target.id] = node + self.generic_visit(node) + + def visit_Call(self, node: ast.Call) -> None: + call_name = _call_name(node.func) + if self._is_argparse_add_argument_call(call_name): + arg_spec = _arg_spec_from_call(node) + if arg_spec: + self.parser_args[arg_spec.name] = arg_spec + + resolved_name = self._resolve_import_path(call_name) + short_name = resolved_name.split(".")[-1] + is_supported_job = _is_supported_job_call(short_name, resolved_name) + is_environment = short_name in SUPPORTED_ENV_NAMES + is_script_runner = short_name == "ScriptRunner" and resolved_name.startswith("nvflare.") + is_unsupported_job = short_name.endswith("Job") and not is_supported_job and short_name != "Job" + if is_supported_job or is_environment or is_script_runner or is_unsupported_job: + call_info = CallInfo( + name=short_name, + full_name=resolved_name, + keywords={keyword.arg: keyword.value for keyword in node.keywords if keyword.arg}, + assignments=self._resolution_assignments(), + source=_source_segment(self.source_text, node), + function_name=self._function_stack[-1] if self._function_stack else None, + ) + if is_environment: + self.env_calls.append(call_info) + elif is_script_runner: + self.script_runner_calls.append(call_info) + elif is_supported_job: + self.job_calls.append(call_info) + else: + self.unsupported_job_calls.append(call_info) + self.generic_visit(node) + + def _track_argparse_assignment(self, target: str, value: ast.AST) -> None: + if not isinstance(value, ast.Call): + return + call_name = self._resolve_import_path(_call_name(value.func)) + if call_name in {"argparse.ArgumentParser", "argparse._ArgumentGroup"}: + self._argparse_parser_names.add(target) + return + owner, _, method = call_name.rpartition(".") + if method == "add_subparsers" and owner in self._argparse_parser_names: + self._argparse_subparser_names.add(target) + elif method == "add_parser" and owner in self._argparse_subparser_names: + self._argparse_parser_names.add(target) + + def _is_argparse_add_argument_call(self, call_name: str) -> bool: + owner, _, method = call_name.rpartition(".") + return method == "add_argument" and owner in self._argparse_parser_names + + def _resolve_import_path(self, call_name: str) -> str: + if not call_name: + return call_name + root, separator, remainder = call_name.partition(".") + imported = self.imports.get(root) + if not imported: + return call_name + return f"{imported}.{remainder}" if separator else imported + + def _current_assignments(self) -> Dict[str, ast.AST]: + if self._local_assignments_stack: + return self._local_assignments_stack[-1] + return self.module_assignments + + def _resolution_assignments(self) -> Dict[str, ast.AST]: + assignments = dict(self.module_assignments) + if self._local_assignments_stack: + assignments.update(self._local_assignments_stack[-1]) + return assignments + + +def _collect_argparse_args_from_file(path: Optional[Path]) -> Dict[str, ArgSpec]: + if not path or not path.exists(): + return {} + try: + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + except SyntaxError: + return {} + return _ImportIndex.from_tree(tree, "").parser_args + + +def _arg_spec_from_call(node: ast.Call) -> Optional[ArgSpec]: + flags = [] + for arg in node.args: + is_literal, value = _literal_value(arg) + if is_literal and isinstance(value, str) and value.startswith("-"): + flags.append(value) + if not flags: + return None + + keywords = {keyword.arg: keyword.value for keyword in node.keywords if keyword.arg} + name = _literal_keyword_value(keywords.get("dest")) or _name_from_flags(flags) + if not name: + return None + + action = _literal_keyword_value(keywords.get("action")) + default, default_source, default_unresolved = _arg_default_from_keywords(keywords) + if not default_unresolved and default is None and action == "store_true": + default = False + default_source = "argparse_action" + elif not default_unresolved and default is None and action == "store_false": + default = True + default_source = "argparse_action" + + return ArgSpec( + name=name, + flags=tuple(flags), + default=default, + default_source=default_source, + default_unresolved=default_unresolved, + value_type=_call_name(keywords["type"]) if "type" in keywords else None, + choices=_literal_sequence(keywords.get("choices")), + action=action, + ) + + +def _arg_default_from_keywords(keywords: Dict[str, ast.AST]) -> Tuple[Any, str, bool]: + if "default" not in keywords: + return None, "argparse_default", False + + node = keywords["default"] + is_literal, literal = _literal_value(node) + if is_literal: + return literal, "literal", False + return _unparse(node), "expression", True + + +def _name_from_flags(flags: Iterable[str]) -> Optional[str]: + long_flags = [flag for flag in flags if flag.startswith("--")] + selected = long_flags[0] if long_flags else next(iter(flags), None) + if not selected: + return None + return selected.lstrip("-").replace("-", "_") + + +def _resolve_value( + node: ast.AST, + assignments: Dict[str, ast.AST], + parser_args: Dict[str, ArgSpec], + source_text: str, + seen: Optional[set[str]] = None, +) -> ResolvedValue: + seen = seen or set() + is_literal, literal = _literal_value(node) + if is_literal: + return ResolvedValue(literal, "literal") + + if isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name) and node.value.id == "args": + arg_spec = parser_args.get(node.attr) + if arg_spec: + return _resolve_arg_default(node.attr, arg_spec) + return ResolvedValue(None, f"unresolved arg:{node.attr}", "low", True) + + if isinstance(node, ast.Name): + if node.id in seen: + return ResolvedValue(None, f"recursive reference:{node.id}", "low", True) + if node.id in parser_args: + return _resolve_arg_default(node.id, parser_args[node.id]) + if node.id in assignments: + return _resolve_value(assignments[node.id], assignments, parser_args, source_text, seen | {node.id}) + return ResolvedValue(node.id, f"name:{node.id}", "low", True) + + if isinstance(node, ast.Call): + call_name = _call_name(node.func) + if call_name in {"Path", "pathlib.Path", "os.path.join"}: + arg_value = _first_resolved_argparse_string(node, assignments, parser_args, source_text) + if arg_value is not None: + return arg_value + return ResolvedValue(_source_segment(source_text, node) or call_name, f"call:{call_name}", "low", True) + + return ResolvedValue(_source_segment(source_text, node) or type(node).__name__, "expression", "low", True) + + +def _resolve_arg_default(name: str, arg_spec: ArgSpec) -> ResolvedValue: + if arg_spec.default_unresolved: + return ResolvedValue(arg_spec.default, f"arg:{name}:{arg_spec.default_source}", "low", True) + return ResolvedValue(arg_spec.default, f"arg:{name}") + + +def _first_resolved_argparse_string( + node: ast.Call, + assignments: Dict[str, ast.AST], + parser_args: Dict[str, ArgSpec], + source_text: str, +) -> Optional[ResolvedValue]: + for arg in node.args: + resolved = _resolve_value(arg, assignments, parser_args, source_text) + if resolved.source.startswith("arg:") and isinstance(resolved.value, str): + return resolved + return None + + +def _call_name(node: Optional[ast.AST]) -> str: + if isinstance(node, ast.Name): + return node.id + if isinstance(node, ast.Attribute): + parent = _call_name(node.value) + return f"{parent}.{node.attr}" if parent else node.attr + return "" + + +def _literal_value(node: Optional[ast.AST]) -> Tuple[bool, Any]: + if isinstance(node, ast.Constant): + return True, node.value + if isinstance(node, (ast.List, ast.Tuple)): + values = [] + for item in node.elts: + is_literal, value = _literal_value(item) + if not is_literal: + return False, None + values.append(value) + return True, values + return False, None + + +def _literal_keyword_value(node: Optional[ast.AST]) -> Any: + is_literal, value = _literal_value(node) + return value if is_literal else None + + +def _literal_sequence(node: Optional[ast.AST]) -> Optional[List[Any]]: + is_literal, value = _literal_value(node) + return value if is_literal and isinstance(value, list) else None + + +def _source_segment(source_text: str, node: ast.AST) -> str: + if not source_text: + return "" + return ast.get_source_segment(source_text, node) or "" + + +def _unparse(node: ast.AST) -> str: + try: + return ast.unparse(node) + except Exception: + return type(node).__name__ + + +def _is_supported_job_call(name: str, resolved_name: str) -> bool: + if not resolved_name.startswith("nvflare."): + return False + if name.endswith("Recipe"): + return True + return name.endswith("Job") and name != "Job" + + +def _is_recipe_call(call_info: CallInfo) -> bool: + return call_info.name.endswith("Recipe") + + +def _surface_name(call_info: Optional[CallInfo]) -> str: + if not call_info: + return "unknown" + return "recipe" if _is_recipe_call(call_info) else "fed_job" + + +def _env_name_to_profile(env_name: str) -> str: + return env_name.removesuffix("Env").lower() + + +def _support_patterns(job_call: Optional[CallInfo], env_call: Optional[CallInfo]) -> List[str]: + patterns = [] + if job_call: + patterns.append(f"{_surface_name(job_call)}:{job_call.name}") + if env_call: + patterns.append(f"env:{env_call.name}") + return patterns + + +def _trust_extracted( + job_call: Optional[CallInfo], + env_call: Optional[CallInfo], + train_script: Optional[Path], + budget: Dict[str, Any], + objective: Dict[str, Any], + search_space: Dict[str, Any], +) -> List[Dict[str, Any]]: + extracted = [] + if job_call: + extracted.append({"field": "job.surface", "value": _surface_name(job_call)}) + extracted.append({"field": f"job.{_surface_name(job_call)}", "value": job_call.name}) + if env_call: + extracted.append({"field": "environment.discovered", "value": env_call.name}) + if train_script: + extracted.append({"field": "job.train_script", "value": train_script.name}) + extracted.append({"field": "objective.metric", "value": objective["metric"]}) + extracted.append({"field": "objective.optimization_metric", "value": objective["optimization_metric"]}) + if "fixed_training_budget" in budget: + extracted.append({"field": "budget.fixed_training_budget", "value": budget["fixed_training_budget"]}) + if search_space: + extracted.append({"field": "search_space.suggested", "value": sorted(search_space)}) + return extracted + + +def _objective_contract(metric_name: str, mode: str, source: str) -> Dict[str, Any]: + return { + "metric": metric_name, + "requested_metric": metric_name, + "optimization_metric": metric_name, + "metric_extraction_order": [metric_name], + "mode": mode, + "metric_contract_source": source, + } + + +def _has_main_entrypoint(tree: ast.AST) -> bool: + return any( + isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == "main" for node in ast.walk(tree) + ) + + +def _existing_path(path: Path) -> Optional[Path]: + return path.resolve() if path.exists() else None + + +def _is_resolved_path_string(value: ResolvedValue) -> bool: + return not value.unresolved and (value.source == "literal" or value.source.startswith("arg:")) + + +def _sha256_text(text: str) -> str: + return hashlib.sha256(text.encode("utf-8")).hexdigest() + + +def _type_name(value: Any) -> str: + if isinstance(value, bool): + return "bool" + if isinstance(value, int): + return "int" + if isinstance(value, float): + return "float" + if isinstance(value, str): + return "str" + if value is None: + return "unknown" + return type(value).__name__ + + +def _overall_confidence(unresolved: List[Dict[str, str]], job_call: Optional[CallInfo]) -> str: + if not job_call: + return "low" + return "medium" if unresolved else "high" + + +def _unresolved(field: str, reason: str) -> Dict[str, str]: + return {"field": field, "reason": reason} diff --git a/skills/nvflare-autofl/scripts/plot_progress.py b/skills/nvflare-autofl/scripts/plot_progress.py new file mode 100755 index 0000000000..9a03933a16 --- /dev/null +++ b/skills/nvflare-autofl/scripts/plot_progress.py @@ -0,0 +1,538 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Generate a readable Auto-FL progress plot from the campaign ledger.""" + +from __future__ import annotations + +import argparse +import csv +import math +import os +import tempfile +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Iterable, List, Optional, Sequence, Tuple + + +class PlotDependencyError(RuntimeError): + """Raised when the rich plotting dependency is unavailable.""" + + +class NoScoredResultsError(ValueError): + """Raised when a campaign has no scored rows to plot yet.""" + + +@dataclass +class ProgressRecord: + index: int + status: str + name: str + score: Optional[float] + runtime_seconds: float + description: str + + +def parse_score(value: Any) -> Optional[float]: + try: + score = float(value) + except (TypeError, ValueError): + return None + if math.isnan(score) or math.isinf(score): + return None + return score + + +def parse_runtime(value: Any) -> float: + try: + seconds = float(value) + except (TypeError, ValueError): + return 0.0 + if math.isnan(seconds) or math.isinf(seconds): + return 0.0 + return max(0.0, seconds) + + +def format_runtime(seconds: float) -> str: + if seconds <= 0: + return "" + if seconds < 3600: + return f"{seconds / 60:.1f}m" + return f"{seconds / 3600:.2f}h" + + +def truncate(text: Any, limit: int) -> str: + compact = " ".join(str(text or "").strip().split()) + if len(compact) <= limit: + return compact + return compact[: max(0, limit - 3)] + "..." + + +def compact_label(text: Any, limit: int) -> str: + compact = " ".join(str(text or "").strip().split()) + replacements = { + "server_momentum": "smom", + "server_lr": "slr", + "clip_grad_norm": "clip", + "eta_min_factor": "eta_min", + "label_smoothing": "ls", + } + for old, new in replacements.items(): + compact = compact.replace(old, new) + for marker in [" (", ";"]: + if marker in compact: + compact = compact.split(marker, 1)[0] + return truncate(compact, limit) + + +def record_label(record: ProgressRecord, limit: int) -> str: + return compact_label(record.name or record.description, limit) + + +def load_results(path: Path) -> List[ProgressRecord]: + records = [] + with path.open("r", encoding="utf-8", newline="") as f: + for index, row in enumerate(csv.DictReader(f, delimiter="\t")): + records.append( + ProgressRecord( + index=index, + status=(row.get("status", "") or "").strip().lower(), + name=row.get("name", "") or row.get("candidate", "") or row.get("commit", ""), + score=parse_score(row.get("score", "")), + runtime_seconds=parse_runtime(row.get("runtime_seconds", "")), + description=row.get("diff_summary", "") or row.get("description", ""), + ) + ) + return records + + +def normalize_records(records: Iterable[Any]) -> List[ProgressRecord]: + normalized = [] + for index, record in enumerate(records): + normalized.append( + ProgressRecord( + index=index, + status=str(getattr(record, "status", "") or "").strip().lower(), + name=str(getattr(record, "name", "") or ""), + score=parse_score(getattr(record, "score", None)), + runtime_seconds=parse_runtime(getattr(record, "runtime_seconds", 0.0)), + description=str(getattr(record, "diff_summary", "") or getattr(record, "description", "") or ""), + ) + ) + return normalized + + +def better(value: float, incumbent: Optional[float], mode: str) -> bool: + if incumbent is None: + return True + return value > incumbent if mode == "max" else value < incumbent + + +def cumulative_best(values: Sequence[float], mode: str) -> List[float]: + incumbent = None + result = [] + for value in values: + if better(value, incumbent, mode): + incumbent = value + result.append(incumbent) + return result + + +def percentile(values: Sequence[float], fraction: float) -> float: + if not values: + raise ValueError("percentile requires at least one value") + if len(values) == 1: + return values[0] + ordered = sorted(values) + position = min(max(fraction, 0.0), 1.0) * (len(ordered) - 1) + lower_index = math.floor(position) + upper_index = math.ceil(position) + if lower_index == upper_index: + return ordered[lower_index] + weight = position - lower_index + return ordered[lower_index] + (ordered[upper_index] - ordered[lower_index]) * weight + + +def default_y_limits( + scores: Sequence[float], baseline: float, mode: str, full_y_range: bool = False +) -> Tuple[float, float]: + score_min = min(scores) + score_max = max(scores) + if full_y_range: + span = max(score_max - score_min, 0.01) + return score_min - max(0.01, span * 0.08), score_max + max(0.01, span * 0.16) + + if mode == "max": + useful_min = min(baseline, percentile(scores, 0.20)) + useful_span = max(score_max - useful_min, 0.01) + lower = useful_min - max(0.01, useful_span * 0.20) + upper = score_max + max(0.015, useful_span * 0.35) + if score_min >= lower: + full_span = max(score_max - score_min, 0.01) + lower = score_min - max(0.01, full_span * 0.08) + return lower, upper + + useful_max = max(baseline, percentile(scores, 0.80)) + useful_span = max(useful_max - score_min, 0.01) + lower = score_min - max(0.015, useful_span * 0.35) + upper = useful_max + max(0.01, useful_span * 0.20) + if score_max <= upper: + full_span = max(score_max - score_min, 0.01) + upper = score_max + max(0.01, full_span * 0.08) + return lower, upper + + +def select_observed_milestones( + valid: Sequence[ProgressRecord], mode: str, max_labels: int +) -> List[Tuple[float, ProgressRecord]]: + if max_labels <= 0: + return [] + incumbent = None + milestones = [] + final_running_best = None + for record in valid: + if record.score is None or not better(record.score, incumbent, mode): + continue + delta = 0.0 if incumbent is None else abs(record.score - incumbent) + final_running_best = (delta, record) + if incumbent is not None: + milestones.append((delta, record)) + incumbent = record.score + + if final_running_best is None: + return [] + if not milestones: + return [final_running_best] + + final_index = final_running_best[1].index + non_final = [item for item in milestones if item[1].index != final_index] + selected = sorted(non_final, key=lambda item: item[0], reverse=True)[: max_labels - 1] + selected_indices = {record.index for _, record in selected} + selected_indices.add(final_index) + return [(delta, record) for delta, record in milestones if record.index in selected_indices] + + +def select_literature_labels(records: Sequence[ProgressRecord], max_labels: int) -> List[ProgressRecord]: + if max_labels <= 0 or not records: + return [] + if len(records) <= max_labels: + return list(records) + latest = records[-1] + selected = {latest.index} + for record in sorted(records, key=lambda item: item.runtime_seconds, reverse=True): + if len(selected) >= max_labels: + break + selected.add(record.index) + return [record for record in records if record.index in selected] + + +def label_placement( + label_number: int, + record: ProgressRecord, + x_limits: Tuple[float, float], + y_limits: Tuple[float, float], +) -> Tuple[Tuple[int, int], str, str]: + x_span = max(x_limits[1] - x_limits[0], 1.0) + y_span = max(y_limits[1] - y_limits[0], 1e-9) + x_fraction = (record.index - x_limits[0]) / x_span + score = record.score if record.score is not None else y_limits[0] + y_fraction = (score - y_limits[0]) / y_span + near_right = x_fraction > 0.72 + near_top = y_fraction > 0.78 + x_offset = -10 if near_right else 10 + y_base = -12 if near_top else 12 + y_step = (label_number % 3) * 8 + return ( + (x_offset, y_base - y_step if near_top else y_base + y_step), + "right" if near_right else "left", + "top" if near_top else "bottom", + ) + + +def plot_progress( + records: Iterable[Any], + output: Path, + mode: str, + metric_label: str, + max_labels: int = 6, + max_literature_labels: int = 4, + full_y_range: bool = False, +) -> Tuple[float, float]: + try: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + except ImportError as exc: + raise PlotDependencyError("matplotlib is required for the rich Auto-FL progress plot") from exc + + rows = normalize_records(records) + valid = [record for record in rows if record.score is not None and record.status != "crash"] + literature_rows = [record for record in rows if record.status == "literature"] + if not valid: + raise NoScoredResultsError("No non-crash rows with numeric scores found") + + baseline_row = next((record for record in valid if record.status == "baseline"), valid[0]) + baseline = baseline_row.score + best_row = valid[0] + for record in valid[1:]: + if better(record.score, best_row.score, mode): + best_row = record + best_score = best_row.score + + fig, ax = plt.subplots(figsize=(16, 8)) + fig.patch.set_facecolor("white") + ax.set_facecolor("white") + styles = { + "discard": {"color": "#cccccc", "size": 18, "alpha": 0.55, "label": "Discarded"}, + "candidate": {"color": "#3498db", "size": 32, "alpha": 0.75, "label": "Candidate"}, + "keep": {"color": "#2ecc71", "size": 72, "alpha": 0.95, "label": "Kept"}, + } + groups = { + "discard": [record for record in valid if record.status == "discard"], + "candidate": [record for record in valid if record.status == "candidate"], + "keep": [record for record in valid if record.status in {"baseline", "keep"}], + } + for status, group in groups.items(): + if not group: + continue + style = styles[status] + ax.scatter( + [record.index for record in group], + [record.score for record in group], + c=style["color"], + s=style["size"], + alpha=style["alpha"], + zorder=4 if status == "keep" else 2, + label=style["label"], + edgecolors="black" if status == "keep" else "none", + linewidths=0.5 if status == "keep" else 0, + ) + + known_statuses = {"baseline", "keep", "discard", "candidate", "crash", "literature"} + other = [record for record in valid if record.status not in known_statuses] + if other: + ax.scatter( + [record.index for record in other], + [record.score for record in other], + c="#9b59b6", + s=28, + alpha=0.65, + zorder=2, + label="Other", + ) + + observed_scores = [record.score for record in valid] + ax.step( + [record.index for record in valid], + cumulative_best(observed_scores, mode), + where="post", + color="#27ae60", + linewidth=2.2, + alpha=0.75, + zorder=3, + label="Running best observed", + ) + + runtime_rows = [record for record in rows if record.runtime_seconds > 0 and record.status != "literature"] + total_runtime = sum(record.runtime_seconds for record in runtime_rows) + average_runtime = total_runtime / len(runtime_rows) if runtime_rows else 0.0 + literature_runtime = sum(record.runtime_seconds for record in literature_rows) + runtime_title = f", {format_runtime(total_runtime)} total" if total_runtime else "" + if average_runtime: + runtime_title += f", {format_runtime(average_runtime)} avg/candidate" + if literature_rows: + runtime_title += f", {len(literature_rows)} lit" + if literature_runtime: + runtime_title += f" ({format_runtime(literature_runtime)})" + + direction = "higher" if mode == "max" else "lower" + ax.set_xlabel("Experiment #", fontsize=12) + ax.set_ylabel(f"{metric_label} ({direction} is better)", fontsize=12) + ax.set_title( + f"Auto-FL Progress ({metric_label}): {len(rows)} rows, {len(valid)} scored, " + f"{sum(record.status == 'keep' for record in rows)} kept, " + f"{sum(record.status == 'candidate' for record in rows)} candidate, " + f"{sum(record.status == 'discard' for record in rows)} discarded, " + f"{sum(record.status == 'crash' for record in rows)} crash{runtime_title}", + fontsize=14, + ) + ax.axhline( + baseline, + color="#7f8c8d", + linewidth=1.2, + alpha=0.5, + linestyle="--", + label="Baseline", + ) + ax.grid(True, alpha=0.2) + + y_limits = default_y_limits(observed_scores, baseline, mode, full_y_range=full_y_range) + ax.set_ylim(*y_limits) + ax.set_xlim(-0.5, max(len(rows) - 0.5, max(record.index for record in valid) + 0.5)) + + if literature_rows: + event_color = "#8e44ad" + event_y = y_limits[1] - max(y_limits[1] - y_limits[0], 1e-9) * 0.035 + for record in literature_rows: + ax.axvline(record.index, color=event_color, linestyle=":", linewidth=1.1, alpha=0.45, zorder=1) + ax.scatter( + [record.index for record in literature_rows], + [event_y for _ in literature_rows], + marker="v", + c=event_color, + s=50, + alpha=0.85, + zorder=5, + label="Literature review", + edgecolors="white", + linewidths=0.4, + ) + event_x_limits = ax.get_xlim() + event_x_span = max(event_x_limits[1] - event_x_limits[0], 1.0) + for record in select_literature_labels(literature_rows, max_literature_labels): + runtime = format_runtime(record.runtime_seconds) + near_right = (record.index - event_x_limits[0]) / event_x_span > 0.88 + annotation = ax.annotate( + f"lit #{record.index}{f' {runtime}' if runtime else ''}: {compact_label(record.description, 30)}", + (record.index, event_y), + textcoords="offset points", + xytext=(-4, -4) if near_right else (4, -4), + fontsize=7.0, + color=event_color, + alpha=0.9, + rotation=90, + ha="right" if near_right else "left", + va="top", + annotation_clip=True, + ) + annotation.set_clip_on(True) + + for label_number, (_, record) in enumerate(select_observed_milestones(valid, mode, max_labels)): + offset, horizontal_align, vertical_align = label_placement(label_number, record, ax.get_xlim(), ax.get_ylim()) + annotation = ax.annotate( + f"#{record.index} {record.score:.4f}: {record_label(record, 28)}", + (record.index, record.score), + textcoords="offset points", + xytext=offset, + fontsize=8.0, + color="#1a7a3a", + alpha=0.9, + ha=horizontal_align, + va=vertical_align, + annotation_clip=True, + arrowprops={ + "arrowstyle": "-", + "color": "#1a7a3a", + "alpha": 0.35, + "linewidth": 0.8, + "shrinkA": 0, + "shrinkB": 4, + }, + ) + annotation.set_clip_on(True) + + summary_lines = [ + f"Baseline: {baseline:.6f}", + f"Best: {best_score:.6f}", + f"Delta: {best_score - baseline:+.6f}", + ] + if total_runtime: + summary_lines.append(f"Runtime: {format_runtime(total_runtime)}") + if average_runtime: + summary_lines.append(f"Avg/candidate: {format_runtime(average_runtime)}") + if literature_rows: + literature_summary = f"Lit reviews: {len(literature_rows)}" + if literature_runtime: + literature_summary += f" ({format_runtime(literature_runtime)})" + summary_lines.append(literature_summary) + summary_lines.append(f"Best run: #{best_row.index} {record_label(best_row, 36)}") + ax.text( + 0.015, + 0.985, + "\n".join(summary_lines), + transform=ax.transAxes, + ha="left", + va="top", + fontsize=9, + bbox={ + "boxstyle": "round,pad=0.35", + "facecolor": "white", + "edgecolor": "#dddddd", + "alpha": 0.9, + }, + ) + + clipped = sum(score < y_limits[0] or score > y_limits[1] for score in observed_scores) + if clipped: + ax.text( + 0.99, + 0.015, + f"{clipped} outlier{'s' if clipped != 1 else ''} outside displayed range", + transform=ax.transAxes, + ha="right", + va="bottom", + fontsize=8, + color="#666666", + ) + + ax.legend(loc="best", fontsize=9) + output.parent.mkdir(parents=True, exist_ok=True) + temp_path = None + try: + with tempfile.NamedTemporaryFile( + dir=output.parent, prefix=f".{output.stem}.", suffix=".png", delete=False + ) as f: + temp_path = Path(f.name) + plt.tight_layout() + fig.savefig(temp_path, dpi=150, facecolor="white", transparent=False) + os.replace(temp_path, output) + finally: + plt.close(fig) + if temp_path and temp_path.exists(): + temp_path.unlink() + return baseline, best_score + + +def main(argv: Optional[Sequence[str]] = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("path", nargs="?", default="results.tsv", help="path to the Auto-FL TSV ledger") + parser.add_argument("--output", default="progress.png", help="output PNG path") + parser.add_argument("--mode", choices=["max", "min"], default="max") + parser.add_argument("--metric", default="score", help="metric label shown in the plot") + parser.add_argument("--max-labels", type=int, default=6) + parser.add_argument("--max-literature-labels", type=int, default=4) + parser.add_argument("--full-y-range", action="store_true") + args = parser.parse_args(argv) + + records = load_results(Path(args.path)) + baseline, best = plot_progress( + records, + Path(args.output), + args.mode, + args.metric, + max_labels=args.max_labels, + max_literature_labels=args.max_literature_labels, + full_y_range=args.full_y_range, + ) + print(f"Saved {args.output}") + print(f"baseline={baseline:.6f}") + print(f"best={best:.6f}") + print(f"delta={best - baseline:+.6f}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/skills/nvflare-autofl/scripts/run_job_campaign.py b/skills/nvflare-autofl/scripts/run_job_campaign.py new file mode 100644 index 0000000000..1f6a5282f7 --- /dev/null +++ b/skills/nvflare-autofl/scripts/run_job_campaign.py @@ -0,0 +1,3070 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Manage agent-authored Auto-FL candidates for an existing NVFlare job.py. + +The coding agent owns hypotheses and source edits. This helper snapshots the +current best source, validates and evaluates candidate manifests, restores +discarded candidates, and records reproducible campaign state and artifacts. +""" + +from __future__ import annotations + +import argparse +import codecs +import csv +import difflib +import hashlib +import importlib.util +import io +import json +import os +import queue +import re +import shlex +import shutil +import signal +import subprocess +import sys +import threading +import time +import uuid +from dataclasses import dataclass, field +from datetime import datetime, timezone +from pathlib import Path, PurePosixPath +from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple + +try: + import yaml +except ImportError: # pragma: no cover - NVFlare installs PyYAML + yaml = None + + +RESULT_FIELDS = [ + "status", + "name", + "score", + "runtime_seconds", + "changed_files", + "diff_summary", + "run_command", + "artifacts", + "failure_reason", + "candidate_manifest", + "base_candidate", + "patch_sha256", + "metric_name", + "metric_source", + "metric_artifact", +] + +CANDIDATE_MANIFEST_SCHEMA_VERSION = "nvflare.autofl.candidate.v1" +CANDIDATE_MANIFEST_STATUSES = {"prepared", "ready_for_external_execution", "keep", "discard", "crash", "abandoned"} +CAMPAIGN_METADATA_SCHEMA_VERSION = "nvflare.autofl.campaign.v1" +CAMPAIGN_METADATA_PATH = ".nvflare/autofl/campaign.json" +CANDIDATE_ROOT = ".nvflare/autofl/candidates" +BEST_SNAPSHOT_ROOT = ".nvflare/autofl/snapshots/best" +ALLOWED_CREATE_PATTERNS = ["**/*.py"] +RESERVED_CANDIDATE_PATH_PARTS = {".git", ".nvflare", "__pycache__", "autofl_runs"} + +INFRASTRUCTURE_RETRY = "infrastructure_retry" +SIMULATOR_STALL_EXIT_CODE = 125 +SIMULATOR_STALL_PATTERNS = ( + "Failed to create connection to the child process in SimulatorClientRunner", + "SimulatorClientRunner - ERROR - run_client_thread error", +) +SIMULATOR_STALL_LOG_LIMIT = 65536 +SIMULATOR_AGGREGATION_RE = re.compile(r"Aggregated\s+(\d+)/(\d+)\s+results") +SIMULATOR_PROGRESS_PATTERNS = ( + "Round ", + "Aggregated ", + "Beginning model validation", + "Saved validation result", + "Finished FedAvg", + "Finished ScatterAndGather", + "cross_site_val", + "round=", + "epoch", + "returning result", + "sent result", + "task completed", +) +SIMULATOR_PROGRESS_LOG_LIMIT = 131072 +DEFAULT_SIMULATOR_NO_PROGRESS_TIMEOUT = 240 +DEFAULT_JOB_HELP_TIMEOUT = 30 +DEFAULT_HARD_CRASH_THRESHOLD = 6 +DEFAULT_PLATEAU_MIN_DELTA = 0.0005 +DEFAULT_PLATEAU_THRESHOLD = 8 + +FIXED_BUDGET_TO_CLI = { + "num_clients": "n_clients", + "min_clients": "min_clients", + "num_rounds": "num_rounds", +} + +COMPARISON_BUDGET_TO_CLI = { + "n_clients": "n_clients", + "num_rounds": "num_rounds", + "aggregation_epochs": "aggregation_epochs", + "local_train_steps": "local_train_steps", + "batch_size": "batch_size", + "eval_batch_size": "eval_batch_size", + "alpha": "alpha", + "seed": "seed", + "model_arch": "model_arch", + "max_model_params": "max_model_params", + "aggregator": "aggregator", + "final_eval_clients": "final_eval_clients", +} + + +@dataclass +class RunRecord: + status: str + name: str + score: Optional[float] + runtime_seconds: float + changed_files: str + diff_summary: str + run_command: str + artifacts: str + failure_reason: str = "" + candidate_manifest: str = "" + base_candidate: str = "" + patch_sha256: str = "" + metric_name: str = "" + metric_source: str = "" + metric_artifact: str = "" + + +@dataclass(frozen=True) +class MetricEvidence: + score: float + metric_name: str + source: str + artifact: str + + +@dataclass +class JobRun: + name: str + args: List[str] + description: str + status: str = "candidate" + score: Optional[float] = None + runtime_seconds: float = 0.0 + artifacts: str = "" + failure_reason: str = "" + command: List[str] = field(default_factory=list) + + +def env_int(name: str, default: int) -> int: + try: + value = int(os.environ.get(name, default)) + except (TypeError, ValueError): + return default + return value + + +def env_float(name: str, default: float) -> float: + try: + value = float(os.environ.get(name, default)) + except (TypeError, ValueError): + return default + return value + + +def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "action", + choices=["initialize", "prepare", "evaluate", "abandon", "suggest", "record", "status"], + help="skill-internal campaign lifecycle action", + ) + parser.add_argument("job", help="NVFlare job.py to optimize") + parser.add_argument("--metric", help="optimization metric; omit to use job.py key_metric") + parser.add_argument("--mode", choices=["max", "min"], default="max") + parser.add_argument("--env", dest="target_env", choices=["sim", "poc", "prod"], default="sim") + cap_group = parser.add_mutually_exclusive_group() + cap_group.add_argument( + "--max-candidates", + type=int, + help="optional candidate cap; omit to continue until interrupted or blocked", + ) + cap_group.add_argument("--uncapped", action="store_true", help="remove a previously configured candidate cap") + parser.add_argument("--autofl-yaml", default="autofl.yaml") + parser.add_argument("--results", default="results.tsv") + parser.add_argument("--state", default=".nvflare/autofl/campaign_state.json") + parser.add_argument("--progress", default="progress.png") + parser.add_argument("--report", default="autofl_report.md") + parser.add_argument("--output-root", default="autofl_runs") + parser.add_argument( + "--plateau-threshold", + type=int, + default=env_int("AUTOFL_PLATEAU_THRESHOLD", DEFAULT_PLATEAU_THRESHOLD), + help=( + "scored comparable candidate attempts after the last material improvement or literature event " + "before campaign state requests run_literature_loop" + ), + ) + parser.add_argument( + "--plateau-min-delta", + type=float, + default=env_float("AUTOFL_PLATEAU_MIN_DELTA", DEFAULT_PLATEAU_MIN_DELTA), + help="minimum metric delta required to reset the plateau clock", + ) + parser.add_argument( + "--hard-crash-threshold", + type=int, + default=env_int("AUTOFL_HARD_CRASH_THRESHOLD", DEFAULT_HARD_CRASH_THRESHOLD), + help="stop after this many consecutive real candidate crashes; set 0 to disable", + ) + parser.add_argument( + "--stop-file", + action="append", + help="manual stop-file path; defaults to STOP_AUTOFL and .nvflare/autofl/STOP under the job directory", + ) + parser.add_argument("--base-args", default="", help="extra job args applied to baseline and all candidates") + parser.add_argument("--timeout", type=int, default=900) + parser.add_argument( + "--simulator-no-progress-timeout", + type=int, + default=env_int("AUTOFL_SIMULATOR_NO_PROGRESS_TIMEOUT_SECONDS", DEFAULT_SIMULATOR_NO_PROGRESS_TIMEOUT), + help=( + "candidate-level simulator no-progress timeout in seconds; set 0 to disable. " + "This is separate from the full run timeout and only applies after progress markers appear." + ), + ) + parser.add_argument("--python", default=os.environ.get("PYTHON") or sys.executable) + parser.add_argument("--prefer-synthetic", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--synthetic-train-size", type=int, default=2048) + parser.add_argument("--synthetic-test-size", type=int, default=256) + parser.add_argument("--name", help="candidate name for prepare") + parser.add_argument("--hypothesis", help="candidate hypothesis for prepare") + parser.add_argument("--manifest", help="candidate_manifest.json path") + parser.add_argument("--run-args", default="", help="candidate-only job.py arguments") + parser.add_argument("--score", type=float, help="externally measured metric for record") + parser.add_argument("--artifacts", dest="external_artifacts", help="external POC/production artifacts") + parser.add_argument("--job-id", help="standard NVFlare job ID for an external result") + parser.add_argument("--failure-reason", default="", help="external execution failure") + parser.add_argument("--baseline", action="store_true", help="record an externally executed baseline") + parser.add_argument("--literature", action="store_true", help="record a literature-review checkpoint") + parser.add_argument("--limit", type=int, default=10, help="maximum fallback suggestions") + args = parser.parse_args(argv) + tokens = list(argv) if argv is not None else sys.argv[1:] + explicit_settings = set() + for action in parser._actions: + for option in action.option_strings: + if any(token == option or token.startswith(f"{option}=") for token in tokens): + explicit_settings.add(action.dest) + break + args._explicit_settings = explicit_settings + if args.uncapped: + args.max_candidates = None + args._explicit_settings.add("max_candidates") + return args + + +def terminate_process(process: subprocess.Popen) -> None: + if process.poll() is not None: + return + if os.name != "nt": + try: + os.killpg(process.pid, signal.SIGTERM) + except ProcessLookupError: + return + except Exception: + process.terminate() + else: + process.terminate() + + try: + process.wait(timeout=10) + return + except subprocess.TimeoutExpired: + pass + + if os.name != "nt": + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + return + except Exception: + process.kill() + else: + process.kill() + + +def recent_text(path: Path, limit: int = SIMULATOR_STALL_LOG_LIMIT) -> str: + try: + with path.open("rb") as f: + try: + f.seek(-limit, os.SEEK_END) + except OSError: + f.seek(0) + return f.read().decode("utf-8", errors="replace") + except OSError: + return "" + + +def detect_nvflare_simulator_stall(sim_root: Path) -> Optional[str]: + if not sim_root.exists(): + return None + + log_paths = [ + sim_root / "server" / "log_fl.txt", + sim_root / "server" / "log.txt", + sim_root / "server" / "error_log.txt", + ] + for path in log_paths: + text = recent_text(path) + if not text: + continue + for pattern in SIMULATOR_STALL_PATTERNS: + if pattern in text: + return f"{pattern} in {path}" + return None + + +def simulator_stall_message(simulator_stall_roots: Sequence[Path]) -> Optional[str]: + for root in simulator_stall_roots: + message = detect_nvflare_simulator_stall(root) + if message: + return message + return None + + +def simulator_progress_signature(sim_root: Path) -> str: + if not sim_root.exists(): + return "" + + log_paths = [ + sim_root / "server" / "log_fl.txt", + sim_root / "server" / "log.txt", + *sorted(sim_root.glob("site-*/log_fl.txt")), + *sorted(sim_root.glob("site-*/log.txt")), + ] + markers: List[str] = [] + for path in log_paths: + text = recent_text(path, limit=SIMULATOR_PROGRESS_LOG_LIMIT) + if not text: + continue + for line in text.splitlines(): + line_lower = line.lower() + if any(pattern.lower() in line_lower for pattern in SIMULATOR_PROGRESS_PATTERNS): + markers.append(f"{path.relative_to(sim_root)}:{line}") + return "\n".join(markers[-200:]) + + +def simulator_progress_signature_for_roots(simulator_stall_roots: Sequence[Path]) -> str: + markers = [] + for root in simulator_stall_roots: + signature = simulator_progress_signature(root) + if signature: + markers.append(f"{root}:\n{signature}") + return "\n".join(markers) + + +def simulator_partial_aggregation_signature(sim_root: Path) -> str: + signatures = [] + for relative_log_path in ("server/log_fl.txt", "server/simulate_job/log_fl.txt"): + path = sim_root / relative_log_path + text = recent_text(path, SIMULATOR_PROGRESS_LOG_LIMIT) + if not text: + continue + for line in reversed(text.splitlines()): + match = SIMULATOR_AGGREGATION_RE.search(line) + if not match: + continue + received = int(match.group(1)) + expected = int(match.group(2)) + if received < expected: + signatures.append(f"{path.relative_to(sim_root)}:{line}") + break + return "\n".join(signatures) + + +def simulator_partial_aggregation_signature_for_roots(simulator_stall_roots: Sequence[Path]) -> str: + markers = [] + for root in simulator_stall_roots: + signature = simulator_partial_aggregation_signature(root) + if signature: + markers.append(f"{root}:\n{signature}") + return "\n".join(markers) + + +def run( + argv: Sequence[str], + cwd: Path, + timeout: int, + log_path: Path, + simulator_stall_roots: Sequence[Path] = (), + stall_check_interval: float = 5.0, + simulator_no_progress_timeout: int = DEFAULT_SIMULATOR_NO_PROGRESS_TIMEOUT, +) -> Tuple[int, str, float]: + started = time.monotonic() + next_stall_check = started + last_progress_check = started + last_progress_seen = started + last_progress_signature = "" + last_partial_aggregation_seen = started + last_partial_aggregation_signature = "" + log_path.parent.mkdir(parents=True, exist_ok=True) + chunks: List[str] = [] + with log_path.open("w", encoding="utf-8") as log_file: + process = subprocess.Popen( + argv, + cwd=str(cwd), + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=False, + bufsize=0, + start_new_session=os.name != "nt", + ) + assert process.stdout is not None + output_queue: queue.Queue[Optional[bytes]] = queue.Queue() + + def read_output() -> None: + try: + while True: + chunk = process.stdout.read(65536) + if not chunk: + break + output_queue.put(chunk) + finally: + output_queue.put(None) + + reader = threading.Thread(target=read_output, name="autofl-process-output", daemon=True) + reader.start() + decoder = codecs.getincrementaldecoder("utf-8")(errors="replace") + output_finished = False + timed_out = False + stall_message = "" + while not output_finished or process.poll() is None: + now = time.monotonic() + if not timed_out and timeout and now - started > timeout: + timed_out = True + terminate_process(process) + if not timed_out and not stall_message and simulator_stall_roots and now >= next_stall_check: + stall_message = simulator_stall_message(simulator_stall_roots) or "" + next_stall_check = now + stall_check_interval + if stall_message: + terminate_process(process) + elif simulator_no_progress_timeout: + partial_aggregation_signature = simulator_partial_aggregation_signature_for_roots( + simulator_stall_roots + ) + if ( + partial_aggregation_signature + and partial_aggregation_signature != last_partial_aggregation_signature + ): + last_partial_aggregation_signature = partial_aggregation_signature + last_partial_aggregation_seen = now + elif ( + last_partial_aggregation_signature + and now - last_partial_aggregation_seen > simulator_no_progress_timeout + ): + stall_message = ( + "partial simulator aggregation made no server-side progress for " + f"{int(now - last_partial_aggregation_seen)}s: {last_partial_aggregation_signature}" + ) + terminate_process(process) + progress_signature = simulator_progress_signature_for_roots(simulator_stall_roots) + if stall_message: + pass + elif progress_signature and progress_signature != last_progress_signature: + last_progress_signature = progress_signature + last_progress_seen = now + elif ( + last_progress_signature + and now - last_progress_seen > simulator_no_progress_timeout + and now - last_progress_check >= stall_check_interval + ): + stall_message = ( + f"no simulator progress markers changed for {int(now - last_progress_seen)}s " + f"across {', '.join(str(root) for root in simulator_stall_roots)}" + ) + terminate_process(process) + last_progress_check = now + try: + raw_chunk = output_queue.get(timeout=0.2) + except queue.Empty: + continue + if raw_chunk is None: + output_finished = True + continue + chunk = decoder.decode(raw_chunk) + if chunk: + chunks.append(chunk) + log_file.write(chunk) + log_file.flush() + reader.join(timeout=1) + remainder = decoder.decode(b"", final=True) + if remainder: + chunks.append(remainder) + log_file.write(remainder) + log_file.flush() + if timed_out: + timeout_msg = f"\nTIMEOUT after {timeout}s\n" + chunks.append(timeout_msg) + log_file.write(timeout_msg) + log_file.flush() + return 124, "".join(chunks), time.monotonic() - started + if stall_message: + stall_text = f"\nSIMULATOR_STALL: {stall_message}\n" + chunks.append(stall_text) + log_file.write(stall_text) + log_file.flush() + return SIMULATOR_STALL_EXIT_CODE, "".join(chunks), time.monotonic() - started + return process.returncode or 0, "".join(chunks), time.monotonic() - started + + +def run_allow_timeout( + argv: Sequence[str], + cwd: Path, + timeout: int, + log_path: Path, + simulator_stall_roots: Sequence[Path] = (), + simulator_no_progress_timeout: int = DEFAULT_SIMULATOR_NO_PROGRESS_TIMEOUT, +) -> Tuple[int, str, float]: + return run( + argv, + cwd, + timeout, + log_path, + simulator_stall_roots=simulator_stall_roots, + simulator_no_progress_timeout=simulator_no_progress_timeout, + ) + + +def read_yaml(path: Path) -> Dict[str, Any]: + if yaml is None: + raise RuntimeError("PyYAML is required to read YAML files") + try: + data = yaml.safe_load(path.read_text(encoding="utf-8")) or {} + except yaml.YAMLError as e: + raise ValueError(f"invalid YAML in {path}: {e}") from e + if not isinstance(data, dict): + raise ValueError(f"invalid YAML in {path}: expected a mapping") + return data + + +def write_yaml(path: Path, data: Dict[str, Any]) -> None: + if yaml is None: + raise RuntimeError("PyYAML is required to write autofl.yaml") + path.write_text(yaml.safe_dump(data, sort_keys=False), encoding="utf-8") + + +def write_json(path: Path, data: Dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(data, indent=2, sort_keys=True) + "\n", encoding="utf-8") + + +def atomic_write_bytes(path: Path, data: bytes) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(f".{path.name}.tmp-{uuid.uuid4().hex}") + try: + temporary.write_bytes(data) + os.replace(temporary, path) + finally: + temporary.unlink(missing_ok=True) + + +def capture_file_versions(paths: Iterable[Path]) -> Dict[Path, Optional[bytes]]: + return {path: path.read_bytes() if path.is_file() else None for path in paths} + + +def restore_file_versions(versions: Dict[Path, Optional[bytes]]) -> None: + for path, data in versions.items(): + if data is None: + path.unlink(missing_ok=True) + else: + atomic_write_bytes(path, data) + + +def read_json(path: Path) -> Dict[str, Any]: + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as e: + raise ValueError(f"failed to read JSON from {path}: {e}") from e + if not isinstance(payload, dict): + raise ValueError(f"expected a JSON object in {path}") + return payload + + +def utc_now() -> str: + return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z") + + +def sha256_bytes(data: bytes) -> str: + return hashlib.sha256(data).hexdigest() + + +def sha256_json(data: Any) -> str: + return sha256_bytes(json.dumps(data, sort_keys=True, separators=(",", ":")).encode("utf-8")) + + +def safe_relative_path(workspace: Path, value: str) -> str: + path = Path(value) + resolved = path.resolve() if path.is_absolute() else (workspace / path).resolve() + try: + relative = resolved.relative_to(workspace.resolve()) + except ValueError as e: + raise ValueError(f"path escapes the Auto-FL job workspace: {value}") from e + if not relative.parts or any(part in RESERVED_CANDIDATE_PATH_PARTS for part in relative.parts): + raise ValueError(f"path is reserved for Auto-FL or repository metadata: {value}") + return relative.as_posix() + + +def allowed_edit_paths(config: Dict[str, Any], workspace: Path) -> List[str]: + values = config.get("trust_contract", {}).get("allowed_edit_paths", []) or [] + if not isinstance(values, list): + raise ValueError("autofl.yaml trust_contract.allowed_edit_paths must be a list") + return list(dict.fromkeys(safe_relative_path(workspace, str(value)) for value in values)) + + +def allowed_create_patterns(config: Dict[str, Any]) -> List[str]: + values = config.get("trust_contract", {}).get("allowed_create_patterns", []) or [] + if not isinstance(values, list) or not all(isinstance(value, str) and value for value in values): + raise ValueError("autofl.yaml trust_contract.allowed_create_patterns must be a list of patterns") + return list(dict.fromkeys(values)) + + +def is_allowed_new_source(path: str, patterns: Sequence[str]) -> bool: + relative = Path(path) + if ( + relative.is_absolute() + or ".." in relative.parts + or any(part in RESERVED_CANDIDATE_PATH_PARTS for part in relative.parts) + ): + return False + posix_path = PurePosixPath(relative.as_posix()) + return any( + posix_path.match(pattern) or (pattern.startswith("**/") and posix_path.match(pattern[3:])) + for pattern in patterns + ) + + +def file_map(root: Path) -> Dict[str, str]: + if not root.exists(): + return {} + files: Dict[str, str] = {} + for path in sorted(root.rglob("*")): + if path.is_symlink(): + raise ValueError(f"candidate source contains a symlink: {path}") + if path.is_file(): + relative = path.relative_to(root).as_posix() + files[relative] = sha256_bytes(path.read_bytes()) + return files + + +def source_hash(files: Dict[str, str]) -> str: + return sha256_json(files) + + +def copy_relative_file(source_root: Path, destination_root: Path, relative: str) -> None: + relative = safe_relative_path(destination_root, relative) + source = source_root / relative + destination = destination_root / relative + try: + source.resolve().relative_to(source_root.resolve()) + except ValueError as e: + raise ValueError(f"source path escapes its managed root: {relative}") from e + if source.is_symlink() or not source.is_file(): + raise ValueError(f"managed source path is not a regular file: {source}") + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source, destination) + + +def stage_best_snapshot(workspace: Path, config: Dict[str, Any], snapshot_root: Path) -> Tuple[Path, Dict[str, str]]: + snapshot_root.parent.mkdir(parents=True, exist_ok=True) + staged_root = snapshot_root.parent / f".{snapshot_root.name}.staged-{uuid.uuid4().hex}" + source_root = staged_root / "source" + source_root.mkdir(parents=True, exist_ok=True) + try: + for relative in allowed_edit_paths(config, workspace): + source = workspace / relative + if source.is_symlink(): + raise ValueError(f"allowed edit path is a symlink: {source}") + if source.is_file(): + copy_relative_file(workspace, source_root, relative) + files = file_map(source_root) + write_json(staged_root / "snapshot.json", {"source_sha256": source_hash(files), "files": files}) + return staged_root, files + except BaseException: + shutil.rmtree(staged_root, ignore_errors=True) + raise + + +def activate_best_snapshot(snapshot_root: Path, staged_root: Path) -> Optional[Path]: + previous_root = None + if snapshot_root.exists(): + previous_root = snapshot_root.parent / f".{snapshot_root.name}.previous-{uuid.uuid4().hex}" + os.replace(snapshot_root, previous_root) + try: + os.replace(staged_root, snapshot_root) + except BaseException: + if previous_root is not None: + os.replace(previous_root, snapshot_root) + raise + return previous_root + + +def rollback_best_snapshot(snapshot_root: Path, previous_root: Optional[Path]) -> None: + if previous_root is None: + shutil.rmtree(snapshot_root, ignore_errors=True) + return + failed_root = snapshot_root.parent / f".{snapshot_root.name}.failed-{uuid.uuid4().hex}" + if snapshot_root.exists(): + os.replace(snapshot_root, failed_root) + try: + os.replace(previous_root, snapshot_root) + except BaseException: + if failed_root.exists(): + os.replace(failed_root, snapshot_root) + raise + shutil.rmtree(failed_root, ignore_errors=True) + + +def create_best_snapshot(workspace: Path, config: Dict[str, Any], snapshot_root: Path) -> Dict[str, str]: + staged_root, files = stage_best_snapshot(workspace, config, snapshot_root) + try: + previous_root = activate_best_snapshot(snapshot_root, staged_root) + finally: + if staged_root.exists(): + shutil.rmtree(staged_root, ignore_errors=True) + if previous_root is not None: + shutil.rmtree(previous_root, ignore_errors=True) + return files + + +def load_best_snapshot(snapshot_root: Path) -> Tuple[Path, Dict[str, str]]: + metadata = read_json(snapshot_root / "snapshot.json") + files = metadata.get("files") + if not isinstance(files, dict) or not all(isinstance(k, str) and isinstance(v, str) for k, v in files.items()): + raise ValueError("best snapshot metadata has an invalid files mapping") + source_root = snapshot_root / "source" + if source_hash(files) != metadata.get("source_sha256") or file_map(source_root) != files: + raise ValueError("best snapshot failed its integrity check") + return source_root, files + + +def workspace_matches_snapshot(workspace: Path, source_root: Path, files: Dict[str, str]) -> bool: + for relative, digest in files.items(): + path = workspace / relative + if path.is_symlink() or not path.is_file() or sha256_bytes(path.read_bytes()) != digest: + return False + return True + + +def validate_candidate_id(value: str) -> str: + if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9._-]{0,63}", value or ""): + raise ValueError("candidate name must match [A-Za-z0-9][A-Za-z0-9._-]{0,63}") + return value + + +def candidate_manifest_path(workspace: Path, candidate_id: str) -> Path: + return workspace / CANDIDATE_ROOT / candidate_id / "candidate_manifest.json" + + +def validate_candidate_manifest_identity(path: Path, manifest: Dict[str, Any]) -> str: + if manifest.get("schema_version") != CANDIDATE_MANIFEST_SCHEMA_VERSION: + raise ValueError(f"unsupported candidate manifest schema in {path}") + candidate_id = validate_candidate_id(str(manifest.get("candidate_id") or "")) + expected = candidate_manifest_path(Path(str(manifest.get("workspace_root") or "")), candidate_id).resolve() + if path.resolve() != expected: + raise ValueError("candidate manifest path does not match its workspace and candidate ID") + if manifest.get("status") not in CANDIDATE_MANIFEST_STATUSES: + raise ValueError(f"candidate {candidate_id} has an invalid manifest status") + return candidate_id + + +def load_candidate_manifest(path: Path) -> Dict[str, Any]: + manifest = read_json(path) + candidate_id = validate_candidate_manifest_identity(path, manifest) + if manifest.get("status") not in {"prepared", "ready_for_external_execution"}: + raise ValueError(f"candidate {candidate_id} is not pending evaluation") + return manifest + + +def campaign_metadata_path(workspace: Path) -> Path: + return workspace / CAMPAIGN_METADATA_PATH + + +def load_campaign_metadata(workspace: Path, job: Path) -> Dict[str, Any]: + metadata = read_json(campaign_metadata_path(workspace)) + if metadata.get("schema_version") != CAMPAIGN_METADATA_SCHEMA_VERSION: + raise ValueError("unsupported Auto-FL campaign metadata schema") + if Path(str(metadata.get("job") or "")).resolve() != job.resolve(): + raise ValueError("campaign metadata belongs to a different job.py") + return metadata + + +def fixed_budget_hash(config: Dict[str, Any]) -> str: + return sha256_json(config.get("budget", {}).get("fixed_training_budget", {}) or {}) + + +def candidate_changes( + workspace: Path, + config: Dict[str, Any], + best_source: Path, + best_files: Dict[str, str], + draft_source: Path, + *, + allow_materialized: bool = False, +) -> Tuple[List[str], List[str]]: + draft_files = file_map(draft_source) + deleted = sorted(set(best_files) - set(draft_files)) + if deleted: + raise ValueError(f"candidate deletes managed source files: {', '.join(deleted)}") + + allowed = set(allowed_edit_paths(config, workspace)) + create_patterns = allowed_create_patterns(config) + changed = sorted(path for path, digest in draft_files.items() if best_files.get(path) != digest) + created = [] + for relative in changed: + if relative in best_files: + if relative not in allowed: + raise ValueError(f"candidate modifies a path outside trust_contract.allowed_edit_paths: {relative}") + continue + if ((workspace / relative).exists() and not allow_materialized) or not is_allowed_new_source( + relative, create_patterns + ): + raise ValueError(f"candidate creates an unauthorized source path: {relative}") + created.append(relative) + return changed, created + + +def text_for_diff(path: Path) -> List[str]: + data = path.read_bytes() + if b"\0" in data: + raise ValueError(f"candidate diff does not support binary file: {path}") + return data.decode("utf-8", errors="replace").splitlines(keepends=True) + + +def render_candidate_patch(best_source: Path, draft_source: Path, changed: Sequence[str]) -> str: + chunks: List[str] = [] + for relative in changed: + before = text_for_diff(best_source / relative) if (best_source / relative).is_file() else [] + after = text_for_diff(draft_source / relative) + chunks.extend( + difflib.unified_diff( + before, + after, + fromfile=f"a/{relative}", + tofile=f"b/{relative}", + ) + ) + return "".join(chunks) + + +def apply_candidate_source(workspace: Path, draft_source: Path, changed: Sequence[str]) -> None: + for relative in changed: + copy_relative_file(draft_source, workspace, relative) + + +def restore_best_source(workspace: Path, best_source: Path, best_files: Dict[str, str], changed: Sequence[str]) -> None: + for value in changed: + relative = safe_relative_path(workspace, value) + destination = workspace / relative + if relative in best_files: + copy_relative_file(best_source, workspace, relative) + elif destination.exists() and not destination.is_dir(): + destination.unlink() + + +def validated_manifest_paths(manifest: Dict[str, Any], workspace: Path, field: str) -> List[str]: + values = manifest.get(field) or [] + if not isinstance(values, list) or not all(isinstance(value, str) for value in values): + raise ValueError(f"candidate manifest {field} must be a list of paths") + paths = [] + for value in values: + normalized = safe_relative_path(workspace, value) + if value != normalized: + raise ValueError(f"candidate manifest {field} contains a non-canonical path: {value}") + paths.append(normalized) + if len(paths) != len(set(paths)): + raise ValueError(f"candidate manifest {field} contains duplicate paths") + return sorted(paths) + + +def validate_manifest_change_set( + manifest: Dict[str, Any], workspace: Path, changed: Sequence[str], created: Sequence[str] +) -> None: + recorded_changed = validated_manifest_paths(manifest, workspace, "changed_files") + recorded_created = validated_manifest_paths(manifest, workspace, "created_files") + if recorded_changed != sorted(changed) or recorded_created != sorted(created): + raise ValueError("candidate manifest source lists do not match the deterministic candidate diff") + + +_CAMPAIGN_GUARD = None +_JOB_IMPORTER = None + + +def load_campaign_guard(): + global _CAMPAIGN_GUARD + if _CAMPAIGN_GUARD is not None: + return _CAMPAIGN_GUARD + + guard_path = Path(__file__).resolve().with_name("campaign_guard.py") + spec = importlib.util.spec_from_file_location("nvflare_autofl_skill_campaign_guard", guard_path) + if spec is None or spec.loader is None: + raise RuntimeError(f"Unable to load campaign guard from {guard_path}") + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + _CAMPAIGN_GUARD = module + return module + + +def load_job_importer(): + global _JOB_IMPORTER + if _JOB_IMPORTER is not None: + return _JOB_IMPORTER + + importer_path = Path(__file__).resolve().with_name("job_importer.py") + spec = importlib.util.spec_from_file_location("nvflare_autofl_skill_job_importer", importer_path) + if spec is None or spec.loader is None: + raise RuntimeError(f"Unable to load Auto-FL job importer from {importer_path}") + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + _JOB_IMPORTER = module + return module + + +def resolve_output_path(cwd: Path, value: str) -> Path: + path = Path(value) + if path.is_absolute(): + return path + return cwd / path + + +def resolve_stop_files(cwd: Path, values: Optional[Sequence[str]]) -> List[str]: + guard = load_campaign_guard() + stop_files = values if values is not None else list(guard.DEFAULT_STOP_FILES) + return [str(resolve_output_path(cwd, value)) for value in stop_files] + + +def extract_result_dir(output: str, cwd: Optional[Path] = None) -> Optional[Path]: + patterns = [ + r"Result can be found in\s*:\s*(?P\S+)", + r"Results:\s*(?P\S+)", + r"result_dir=(?P\S+)", + ] + for pattern in patterns: + match = re.search(pattern, output) + if match: + path = Path(match.group("path")).expanduser() + if not path.is_absolute() and cwd is not None: + path = cwd / path + return path.resolve() + return None + + +def objective_contract(config: Dict[str, Any], requested_metric: Optional[str]) -> Dict[str, Any]: + objective = config.get("objective", {}) + if not isinstance(objective, dict): + objective = {} + metric = str(objective.get("metric") or requested_metric or "accuracy") + requested = str(objective.get("requested_metric") or metric) + optimization = str(objective.get("optimization_metric") or metric) + extraction_order = objective.get("metric_extraction_order") + if not isinstance(extraction_order, list) or not extraction_order: + extraction_order = [optimization] + extraction_order = [str(item) for item in extraction_order if item] + if optimization not in extraction_order: + extraction_order.insert(0, optimization) + return { + **objective, + "metric": metric, + "requested_metric": requested, + "optimization_metric": optimization, + "metric_extraction_order": extraction_order, + } + + +def apply_metric_contract( + config: Dict[str, Any], requested_metric: Optional[str], schema: Optional[Dict[str, Any]] +) -> Dict[str, Any]: + objective = objective_contract(config, requested_metric) + schema_objective = (schema or {}).get("objective", {}) + if isinstance(schema_objective, dict): + schema_requested = schema_objective.get("requested_metric") or schema_objective.get("metric") + if not schema_requested or schema_requested == objective["requested_metric"]: + for key in ("optimization_metric", "metric_extraction_order", "metric_source"): + if key in schema_objective: + objective[key] = schema_objective[key] + config["objective"] = objective_contract({"objective": objective}, requested_metric) + return config + + +def metric_extraction_order(config: Dict[str, Any], requested_metric: Optional[str]) -> List[str]: + return list(objective_contract(config, requested_metric)["metric_extraction_order"]) + + +def optimization_metric(config: Dict[str, Any], requested_metric: Optional[str]) -> str: + return str(objective_contract(config, requested_metric)["optimization_metric"]) + + +def metric_source(config: Dict[str, Any]) -> str: + source = config.get("objective", {}).get("metric_source", "") + return str(source) if source else "NVFlare metric artifacts" + + +def normalize_metric_order(metrics: Sequence[str] | str) -> List[str]: + if isinstance(metrics, str): + return [metrics] + return [str(metric) for metric in metrics if metric] + + +def is_numeric_metric_value(value: Any) -> bool: + return isinstance(value, (int, float)) and not isinstance(value, bool) + + +def find_metric_value(payload: Any, metric_order: Sequence[str] | str) -> Optional[float]: + metric_keys = normalize_metric_order(metric_order) + if isinstance(payload, dict): + for key in ("final_aggregated_metrics", "best_metrics", "metrics"): + for metric_key in metric_keys: + value = metric_from_list(payload.get(key), metric_key) + if value is not None: + return value + for metric_key in metric_keys: + if metric_key in payload and is_numeric_metric_value(payload[metric_key]): + return float(payload[metric_key]) + for value in payload.values(): + score = find_metric_value(value, metric_keys) + if score is not None: + return score + elif isinstance(payload, list): + for metric_key in metric_keys: + value = metric_from_list(payload, metric_key) + if value is not None: + return value + for item in payload: + score = find_metric_value(item, metric_keys) + if score is not None: + return score + return None + + +def metric_from_list(items: Any, metric: str) -> Optional[float]: + if not isinstance(items, list): + return None + for item in items: + if not isinstance(item, dict): + continue + if item.get("name") == metric and is_numeric_metric_value(item.get("value")): + return float(item["value"]) + return None + + +def extract_metric_evidence(artifact_root: Path, metrics: Sequence[str] | str) -> Optional[MetricEvidence]: + metric_order = normalize_metric_order(metrics) + metric_files = sorted(artifact_root.glob("**/metrics_summary.json")) + sorted( + artifact_root.glob("**/cross_val_results.json") + ) + payloads = [] + for path in metric_files: + try: + payload = json.loads(path.read_text(encoding="utf-8")) + except Exception: + continue + payloads.append((path, payload)) + for metric in metric_order: + for path, payload in payloads: + score = find_metric_value(payload, [metric]) + if score is not None: + return MetricEvidence( + score=score, + metric_name=metric, + source=f"structured:{path.name}", + artifact=str(path.resolve()), + ) + + text_artifacts = [] + for name in ("run.log", "log.txt", "log_fl.txt", "error_log.txt"): + text_artifacts.extend(sorted(artifact_root.glob(f"**/{name}"))) + for metric in metric_order: + number = r"[+-]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?" + pattern = rf"(? Optional[float]: + evidence = extract_metric_evidence(artifact_root, metrics) + return evidence.score if evidence else None + + +def metric_search_description(artifact_root: Path, metrics: Sequence[str] | str) -> str: + names = ["metrics_summary.json", "cross_val_results.json", "run.log", "log.txt", "log_fl.txt", "error_log.txt"] + searched = [str(path) for name in names for path in sorted(artifact_root.glob(f"**/{name}"))] + return f"metrics={normalize_metric_order(metrics)!r}; searched={searched or [str(artifact_root)]!r}" + + +def is_sandbox_socket_failure(output: str) -> bool: + text = output.lower() + return ( + "permissionerror" in text + and ("operation not permitted" in text or "[errno 1]" in text) + and ("socket" in text or "sock" in text or "get_open_ports" in text or ".bind(" in text) + ) + + +def is_nvflare_simulator_stall(output: str) -> bool: + return "SIMULATOR_STALL:" in output + + +def collect_artifacts(result_dir: Optional[Path], output_root: Path, name: str, log_path: Path) -> Path: + dest = output_root / name / "simulation" + run_log = output_root / name / "run.log" + if dest.exists(): + shutil.rmtree(dest) + if result_dir and result_dir.exists(): + shutil.copytree(result_dir, dest) + else: + dest.mkdir(parents=True, exist_ok=True) + if log_path.resolve() != run_log.resolve(): + shutil.copy2(log_path, run_log) + return dest + + +def job_help(python: str, job: Path, cwd: Path, timeout: int = DEFAULT_JOB_HELP_TIMEOUT) -> str: + try: + process = subprocess.run( + [python, str(job), "--help"], + cwd=str(cwd), + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + check=False, + timeout=timeout, + ) + except subprocess.TimeoutExpired as e: + raise RuntimeError(f"job.py --help exceeded {timeout} seconds") from e + return process.stdout + + +def supported_long_flags(help_text: str) -> set[str]: + return set(re.findall(r"(? bool: + return flag in supported_long_flags(help_text) + + +def mutable_arg_specs(schema: Dict[str, Any]) -> Dict[str, Dict[str, Any]]: + specs = schema.get("mutable_args") + return specs if isinstance(specs, dict) else {} + + +def candidate_arg_values(candidate_args: Sequence[str]) -> Dict[str, Any]: + values: Dict[str, Any] = {} + idx = 0 + while idx < len(candidate_args): + raw = candidate_args[idx] + if raw.startswith("--") and not re.match(r"^--[A-Za-z][A-Za-z0-9_-]*(?:=.*)?$", raw): + raise ValueError(f"invalid canonical long option: {raw!r}") + if ( + raw.startswith("-") + and not raw.startswith("--") + and not re.match(r"^-(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?$", raw) + ): + raise ValueError(f"candidate run arguments require canonical long options, not {raw!r}") + if not raw.startswith("--"): + idx += 1 + continue + option = raw[2:] + if "=" in option: + name, value = option.split("=", 1) + values[name.replace("-", "_")] = value + idx += 1 + continue + name = option.replace("-", "_") + if idx + 1 >= len(candidate_args) or candidate_args[idx + 1].startswith("--"): + values[name] = True + idx += 1 + else: + values[name] = candidate_args[idx + 1] + idx += 2 + return values + + +def coerce_schema_value(value: Any, spec: Dict[str, Any]) -> Any: + value_type = spec.get("type") + if value_type == "int": + return int(value) + if value_type == "float": + return float(value) + if value_type == "bool": + if isinstance(value, bool): + return value + return str(value).strip().lower() in {"1", "true", "yes", "on"} + return str(value) + + +def candidate_args_allowed(candidate_args: Sequence[str], schema: Dict[str, Any]) -> Tuple[bool, str]: + specs = mutable_arg_specs(schema) + if not specs: + return True, "" + + for name, value in candidate_arg_values(candidate_args).items(): + spec = specs.get(name) + if not isinstance(spec, dict): + continue + try: + coerced = coerce_schema_value(value, spec) + except (TypeError, ValueError) as e: + return False, f"{name}={value!r} cannot be parsed as {spec.get('type')}: {e}" + + choices = spec.get("choices") + if choices is not None and coerced not in choices: + return False, f"{name}={coerced!r} is not in allowed choices {choices!r}" + minimum = spec.get("min") + if minimum is not None and coerced < minimum: + return False, f"{name}={coerced!r} is below schema min {minimum!r}" + maximum = spec.get("max") + if maximum is not None and coerced > maximum: + return False, f"{name}={coerced!r} is above schema max {maximum!r}" + + return True, "" + + +def candidate_preserves_fixed_args( + candidate_args: Sequence[str], config: Dict[str, Any], schema: Dict[str, Any] +) -> Tuple[bool, str]: + values = candidate_arg_values(candidate_args) + fixed_names = set(fixed_within_campaign(schema)) + fixed_budget = config.get("budget", {}).get("fixed_training_budget", {}) or {} + fixed_names.update(FIXED_BUDGET_TO_CLI[field] for field in fixed_budget if field in FIXED_BUDGET_TO_CLI) + normalized_fixed_names = {name.replace("-", "_") for name in fixed_names} + changed = sorted(normalized_fixed_names.intersection(values)) + if changed: + return False, f"candidate run arguments change fixed-budget fields: {', '.join(changed)}" + for raw in candidate_args: + if not raw.startswith("--"): + continue + supplied = raw[2:].split("=", 1)[0].replace("-", "_") + abbreviated = sorted(name for name in normalized_fixed_names if name.startswith(supplied) and name != supplied) + if abbreviated: + return False, ( + f"candidate run argument {raw!r} is a strict prefix of fixed-budget option(s): " + f"{', '.join('--' + name for name in abbreviated)}" + ) + return True, "" + + +def load_mutation_schema(cwd: Path) -> Dict[str, Any]: + path = cwd / "mutation_schema.yaml" + if not path.exists(): + return {} + return read_yaml(path) + + +def apply_mutation_schema_contract(config: Dict[str, Any], schema: Dict[str, Any], workspace: Path) -> Dict[str, Any]: + preferred_targets = schema.get("preferred_targets") + if preferred_targets is None: + return config + if not isinstance(preferred_targets, list) or not all(isinstance(value, str) for value in preferred_targets): + raise ValueError("mutation_schema.yaml preferred_targets must be a list of paths") + + trust_contract = config.setdefault("trust_contract", {}) + trust_paths = trust_contract.setdefault("allowed_edit_paths", []) + if not isinstance(trust_paths, list): + raise ValueError("autofl.yaml trust_contract.allowed_edit_paths must be a list") + + resolved_targets = [] + unresolved_targets = [] + for target in preferred_targets: + try: + relative = safe_relative_path(workspace, target) + except ValueError as e: + unresolved_targets.append(_schema_target_issue(target, str(e))) + continue + path = workspace / relative + if path.is_symlink(): + unresolved_targets.append(_schema_target_issue(target, "preferred target is a symlink")) + continue + if not path.is_file(): + unresolved_targets.append( + _schema_target_issue(target, "preferred target is not an existing file under the job workspace") + ) + continue + resolved_targets.append(relative) + if relative not in trust_paths: + trust_paths.append(relative) + + trust_contract["preferred_targets"] = list(resolved_targets) + if unresolved_targets: + unresolved = config.setdefault("unresolved", []) + trust_unresolved = trust_contract.setdefault("unresolved", []) + if not isinstance(unresolved, list) or not isinstance(trust_unresolved, list): + raise ValueError("autofl.yaml unresolved fields must be lists") + unresolved.extend(unresolved_targets) + trust_unresolved.extend(dict(item) for item in unresolved_targets) + return config + + +def _schema_target_issue(target: str, reason: str) -> Dict[str, str]: + return {"field": "mutation_schema.preferred_targets", "reason": f"{target}: {reason}"} + + +def comparison_budget(schema: Dict[str, Any]) -> Dict[str, Any]: + comparison = schema.get("comparison_budget_args") + if not isinstance(comparison, dict): + return {} + budget = comparison.get("default_candidate_budget") + return budget if isinstance(budget, dict) else {} + + +def fixed_within_campaign(schema: Dict[str, Any]) -> set: + values = [] + comparison = schema.get("comparison_budget_args") + if isinstance(comparison, dict): + values = comparison.get("fixed_within_campaign") or [] + return set(values) if isinstance(values, list) else set() + + +def build_comparison_budget_args(schema: Dict[str, Any], help_text: str) -> List[str]: + budget = comparison_budget(schema) + args: List[str] = [] + for budget_field, cli_name in COMPARISON_BUDGET_TO_CLI.items(): + value = budget.get(budget_field) + if value is not None and supports_flag(help_text, f"--{cli_name}"): + args.extend([f"--{cli_name}", str(value)]) + if budget.get("cross_site_eval") and supports_flag(help_text, "--cross_site_eval"): + args.append("--cross_site_eval") + return args + + +def build_fixed_args(config: Dict[str, Any], help_text: str, schema: Dict[str, Any]) -> List[str]: + fixed = config.get("budget", {}).get("fixed_training_budget", {}) or {} + budget = comparison_budget(schema) + budget_cli_names = { + cli_name for budget_field, cli_name in COMPARISON_BUDGET_TO_CLI.items() if budget.get(budget_field) is not None + } + args: List[str] = [] + for budget_field, cli_name in FIXED_BUDGET_TO_CLI.items(): + if cli_name in budget_cli_names: + continue + value = fixed.get(budget_field) + if value is not None and supports_flag(help_text, f"--{cli_name}"): + args.extend([f"--{cli_name}", str(value)]) + return args + + +def build_base_args(args: argparse.Namespace, help_text: str, schema: Dict[str, Any]) -> List[str]: + base = shlex.split(args.base_args) + budget_args = build_comparison_budget_args(schema, help_text) + if budget_args: + base.extend(budget_args) + if args.prefer_synthetic and supports_flag(help_text, "--synthetic_data"): + if "--synthetic_data" not in base: + base.append("--synthetic_data") + if supports_flag(help_text, "--train_size") and "--train_size" not in base: + base.extend(["--train_size", str(args.synthetic_train_size)]) + if supports_flag(help_text, "--test_size") and "--test_size" not in base: + base.extend(["--test_size", str(args.synthetic_test_size)]) + return base + + +def suggested_arg_defaults(config: Dict[str, Any]) -> Dict[str, Any]: + suggested = config.get("search_space", {}).get("suggested", {}) or {} + defaults = {} + for name, spec in suggested.items(): + if isinstance(spec, dict) and "default" in spec: + defaults[name] = spec["default"] + return defaults + + +def candidate_plan( + config: Dict[str, Any], + help_text: str, + max_candidates: Optional[int], + schema: Optional[Dict[str, Any]] = None, +) -> Iterable[JobRun]: + defaults = suggested_arg_defaults(config) + candidates: List[JobRun] = [] + seen_args = set() + fixed_fields = fixed_within_campaign(schema or {}) + + def can_mutate(name: str) -> bool: + return name not in fixed_fields + + def make_candidate(name: str, candidate_args: List[str], description: str) -> Optional[JobRun]: + allowed, _ = candidate_args_allowed(candidate_args, schema or {}) + if not allowed: + return None + key = tuple(candidate_args) + if key in seen_args: + return None + seen_args.add(key) + return JobRun(name=name, args=candidate_args, description=description) + + def add(name: str, candidate_args: List[str], description: str) -> None: + if max_candidates is not None and len(candidates) >= max_candidates: + return + candidate = make_candidate(name, candidate_args, description) + if candidate is not None: + candidates.append(candidate) + + if can_mutate("aggregator") and supports_flag(help_text, "--aggregator"): + for value in ["default", "fedavg", "fedavgm", "fedadam", "scaffold"]: + add(f"aggregator_{value}", ["--aggregator", value], f"aggregator={value}") + if can_mutate("fedproxloss_mu") and supports_flag(help_text, "--fedproxloss_mu"): + add( + "fedprox_mu_1e-5", + ["--aggregator", "weighted", "--fedproxloss_mu", "1e-5"], + "aggregator=weighted fedproxloss_mu=1e-5", + ) + add( + "fedprox_mu_1e-4", + ["--aggregator", "weighted", "--fedproxloss_mu", "1e-4"], + "aggregator=weighted fedproxloss_mu=1e-4", + ) + + if can_mutate("aggregation_epochs") and supports_flag(help_text, "--aggregation_epochs"): + default = int(defaults.get("aggregation_epochs") or 4) + for value in [1, 2, 3, 4, 6, 8]: + if value != default: + add(f"aggregation_epochs_{value}", ["--aggregation_epochs", str(value)], f"aggregation_epochs={value}") + + if can_mutate("local_train_steps") and supports_flag(help_text, "--local_train_steps"): + for value in [50, 100, 200, 400]: + add(f"local_train_steps_{value}", ["--local_train_steps", str(value)], f"local_train_steps={value}") + + if can_mutate("lr") and supports_flag(help_text, "--lr"): + default = float(defaults.get("lr") or 0.05) + for value in [default / 4, default / 2, default * 2, default * 4]: + value_text = f"{value:.6g}" + add(f"lr_{value_text.replace('.', 'p').replace('-', 'm')}", ["--lr", value_text], f"lr={value_text}") + + if can_mutate("momentum") and supports_flag(help_text, "--momentum"): + for value in [0.0, 0.5, 0.8, 0.95]: + add( + f"momentum_{str(value).replace('.', 'p')}", + ["--momentum", str(value)], + f"momentum={value}", + ) + + if can_mutate("weight_decay") and supports_flag(help_text, "--weight_decay"): + for value in ["1e-5", "1e-4", "5e-4"]: + add(f"weight_decay_{value.replace('-', 'm')}", ["--weight_decay", value], f"weight_decay={value}") + + if can_mutate("batch_size") and supports_flag(help_text, "--batch_size"): + default = int(defaults.get("batch_size") or 16) + values = [max(1, default // 2), default * 2, default * 4, max(1, default // 4), 24, 40, 64, 96, 128, 256] + for value in values: + if value != default: + add(f"batch_size_{value}", ["--batch_size", str(value)], f"batch_size={value}") + + if can_mutate("epochs") and supports_flag(help_text, "--epochs"): + for value in [1, 2, 3, 4, 5]: + add(f"epochs_{value}", ["--epochs", str(value)], f"epochs={value}") + + if can_mutate("num_workers") and supports_flag(help_text, "--num_workers"): + for value in [0, 1, 2, 4]: + add(f"num_workers_{value}", ["--num_workers", str(value)], f"num_workers={value}") + + if supports_flag(help_text, "--client_memory_gc_rounds"): + add("client_memory_gc_1", ["--client_memory_gc_rounds", "1"], "client_memory_gc_rounds=1") + + if not candidates: + add("rerun", [], "repeat baseline command to test campaign plumbing") + + if max_candidates is not None: + return candidates[:max_candidates] + + def uncapped() -> Iterable[JobRun]: + for template in candidates: + yield JobRun(name=template.name, args=list(template.args), description=template.description) + + idx = 1 + batch_default = int(defaults.get("batch_size") or 16) + while True: + generated = False + if can_mutate("batch_size") and supports_flag(help_text, "--batch_size"): + # Walk a broad conservative range before repeats so uncapped + # campaigns keep doing comparable, reviewable work for hours. + value = 1 + ((batch_default + idx * 7) % 512) + candidate_args = ["--batch_size", str(value)] + candidate = make_candidate( + f"batch_size_auto_{value}", + candidate_args, + f"batch_size={value}", + ) + if value != batch_default and candidate is not None: + yield candidate + generated = True + + if not generated and can_mutate("aggregation_epochs") and supports_flag(help_text, "--aggregation_epochs"): + value = 1 + ((idx - 1) % 8) + candidate_args = ["--aggregation_epochs", str(value)] + candidate = make_candidate( + f"aggregation_epochs_auto_{value}", + candidate_args, + f"aggregation_epochs={value}", + ) + if candidate is not None: + yield candidate + generated = True + + if not generated and can_mutate("lr") and supports_flag(help_text, "--lr"): + value = 10 ** (-4 + ((idx - 1) % 25) / 10) + value_text = f"{value:.6g}" + candidate_args = ["--lr", value_text] + candidate = make_candidate( + f"lr_auto_{value_text.replace('.', 'p').replace('-', 'm')}", + candidate_args, + f"lr={value_text}", + ) + if candidate is not None: + yield candidate + generated = True + + if not generated and can_mutate("epochs") and supports_flag(help_text, "--epochs"): + value = 1 + ((idx - 1) % 20) + candidate_args = ["--epochs", str(value)] + candidate = make_candidate(f"epochs_auto_{value}", candidate_args, f"epochs={value}") + if candidate is not None: + yield candidate + generated = True + + if not generated and can_mutate("num_workers") and supports_flag(help_text, "--num_workers"): + value = (idx - 1) % 9 + candidate_args = ["--num_workers", str(value)] + candidate = make_candidate(f"num_workers_auto_{value}", candidate_args, f"num_workers={value}") + if candidate is not None: + yield candidate + generated = True + + if not generated: + template = candidates[(idx - 1) % len(candidates)] + cycle = (idx - 1) // len(candidates) + 2 + yield JobRun( + name=f"{template.name}_repeat_{cycle:04d}", + args=list(template.args), + description=f"{template.description}; repeat_cycle={cycle}", + ) + idx += 1 + + return uncapped() + + +def imported_job_names(config: Dict[str, Any]) -> List[str]: + names = [] + job_config = config.get("job", {}) + if not isinstance(job_config, dict): + return names + for key in ("recipe_args", "fed_job_args"): + arguments = job_config.get(key) + if not isinstance(arguments, dict): + continue + name = arguments.get("name") + value = name.get("value") if isinstance(name, dict) else name + if isinstance(name, dict) and name.get("confidence") != "high": + continue + if isinstance(value, str) and value and value not in names: + names.append(value) + return names + + +def simulator_workspace_root(config: Dict[str, Any], cwd: Path) -> Path: + discovered = config.get("environment", {}).get("discovered", {}) + if isinstance(discovered, dict): + arguments = discovered.get("args", {}) + workspace_arg = arguments.get("workspace_root") if isinstance(arguments, dict) else None + if isinstance(workspace_arg, dict): + value = workspace_arg.get("value") + if workspace_arg.get("confidence") == "high" and isinstance(value, str) and value: + root = Path(value).expanduser() + return (root if root.is_absolute() else cwd / root).resolve() + return Path("/tmp/nvflare/simulation").resolve() + + +def expected_simulator_roots(config: Dict[str, Any], injected_name: Optional[str], cwd: Path) -> List[Path]: + names = ([injected_name] if injected_name else []) + imported_job_names(config) + roots = [] + simulator_base = simulator_workspace_root(config, cwd) + for name in names: + if Path(name).name != name or name in {".", ".."}: + raise ValueError(f"unsafe simulator job name: {name!r}") + root = (simulator_base / name).resolve() + if root.parent != simulator_base: + raise ValueError(f"unsafe simulator result root: {root}") + if root not in roots: + roots.append(root) + return roots + + +def run_job( + run_def: JobRun, + *, + python: str, + job: Path, + cwd: Path, + help_text: str, + fixed_args: List[str], + base_args: List[str], + output_root: Path, + timeout: int, + simulator_no_progress_timeout: int, + metrics: Sequence[str], + config: Dict[str, Any], +) -> RunRecord: + baseline_run = run_def.status == "baseline" + log_path = output_root / run_def.name / "run.log" + run_name = f"autofl_{run_def.name}" + name_args = ["--name", run_name] if supports_flag(help_text, "--name") else [] + simulator_roots = expected_simulator_roots(config, run_name if name_args else None, cwd) + command = [python, str(job), *fixed_args, *base_args, *name_args, *run_def.args] + run_def.command = command + for simulator_root in simulator_roots: + shutil.rmtree(simulator_root, ignore_errors=True) + rc, stdout, runtime = run_allow_timeout( + command, + cwd, + timeout, + log_path, + simulator_stall_roots=simulator_roots, + simulator_no_progress_timeout=simulator_no_progress_timeout, + ) + run_def.runtime_seconds = runtime + printed_result_dir = extract_result_dir(stdout, cwd) + existing_roots = [root.resolve() for root in simulator_roots if root.exists()] + result_dir = printed_result_dir if printed_result_dir in existing_roots else None + injected_root = (simulator_workspace_root(config, cwd) / run_name).resolve() if name_args else None + if result_dir is None and injected_root in existing_roots: + result_dir = injected_root + if result_dir is None and len(existing_roots) == 1: + result_dir = existing_roots[0] + artifact_dir = collect_artifacts(result_dir, output_root, run_def.name, log_path) + run_def.artifacts = str(artifact_dir) + + evidence = None + if rc != 0: + if is_sandbox_socket_failure(stdout): + run_def.status = INFRASTRUCTURE_RETRY + run_def.failure_reason = "sandbox/socket permission failure; rerun runner with escalated execution" + elif is_nvflare_simulator_stall(stdout): + run_def.status = "crash" + run_def.failure_reason = ( + "nvflare simulator watchdog detected a child connection/no-progress stall; " + "candidate killed and campaign continued" + ) + else: + run_def.status = "crash" + run_def.failure_reason = f"exit_code={rc}" + elif result_dir is None: + run_def.status = "crash" + run_def.failure_reason = ( + "job exited successfully but no deterministic NVFlare result directory was resolved; " + "expose literal SimEnv workspace_root and job name values" + ) + else: + artifact_root = artifact_dir.parent + evidence = extract_metric_evidence(artifact_root, metrics) + if evidence is None: + run_def.status = "crash" + run_def.failure_reason = f"matching metric not found; {metric_search_description(artifact_root, metrics)}" + else: + run_def.score = evidence.score + + record_status = "baseline" if baseline_run and run_def.status == "crash" else run_def.status + return RunRecord( + status=record_status, + name=run_def.name, + score=run_def.score, + runtime_seconds=run_def.runtime_seconds, + changed_files="none", + diff_summary=run_def.description, + run_command=shlex.join(command), + artifacts=run_def.artifacts, + failure_reason=run_def.failure_reason, + metric_name=evidence.metric_name if evidence else "", + metric_source=evidence.source if evidence else "", + metric_artifact=evidence.artifact if evidence else "", + ) + + +def write_results(path: Path, records: List[RunRecord]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + stream = io.StringIO(newline="") + writer = csv.DictWriter(stream, fieldnames=RESULT_FIELDS, delimiter="\t") + writer.writeheader() + for record in records: + writer.writerow( + { + "status": record.status, + "name": record.name, + "score": "" if record.score is None else f"{record.score:.6f}", + "runtime_seconds": f"{record.runtime_seconds:.3f}", + "changed_files": record.changed_files, + "diff_summary": record.diff_summary, + "run_command": record.run_command, + "artifacts": record.artifacts, + "failure_reason": record.failure_reason, + "candidate_manifest": record.candidate_manifest, + "base_candidate": record.base_candidate, + "patch_sha256": record.patch_sha256, + "metric_name": record.metric_name, + "metric_source": record.metric_source, + "metric_artifact": record.metric_artifact, + } + ) + atomic_write_bytes(path, stream.getvalue().encode("utf-8")) + + +def load_results(path: Path) -> List[RunRecord]: + if not path.exists(): + return [] + records = [] + with path.open("r", encoding="utf-8", newline="") as f: + for row in csv.DictReader(f, delimiter="\t"): + score_text = row.get("score", "") + records.append( + RunRecord( + status=row.get("status", ""), + name=row.get("name", ""), + score=float(score_text) if score_text else None, + runtime_seconds=float(row.get("runtime_seconds") or 0.0), + changed_files=row.get("changed_files", ""), + diff_summary=row.get("diff_summary", ""), + run_command=row.get("run_command", ""), + artifacts=row.get("artifacts", ""), + failure_reason=row.get("failure_reason", ""), + candidate_manifest=row.get("candidate_manifest", ""), + base_candidate=row.get("base_candidate", ""), + patch_sha256=row.get("patch_sha256", ""), + metric_name=row.get("metric_name", ""), + metric_source=row.get("metric_source", ""), + metric_artifact=row.get("metric_artifact", ""), + ) + ) + return records + + +def better(new_score: Optional[float], old_score: Optional[float], mode: str) -> bool: + if new_score is None: + return False + if old_score is None: + return True + return new_score > old_score if mode == "max" else new_score < old_score + + +def write_state( + path: Path, + results_path: Path, + records: List[RunRecord], + max_candidates: Optional[int], + *, + mode: str = "max", + stop_files: Optional[List[str]] = None, + plateau_threshold: int = DEFAULT_PLATEAU_THRESHOLD, + plateau_min_delta: float = DEFAULT_PLATEAU_MIN_DELTA, + hard_crash_threshold: int = DEFAULT_HARD_CRASH_THRESHOLD, + pending_manifest_count: int = 0, + persist: bool = True, +) -> Dict[str, Any]: + guard = load_campaign_guard() + attempts = len( + [ + record + for record in records + if record.status in {"candidate", "keep", "discard", "crash"} and not is_baseline_record(record) + ] + ) + if records and records[-1].status == INFRASTRUCTURE_RETRY: + state = guard.guard_state( + results_path, + max_candidates=max_candidates, + stop_files=stop_files, + plateau_threshold=plateau_threshold, + min_delta=plateau_min_delta, + hard_crash_threshold=hard_crash_threshold, + mode=mode, + pending_manifest_count=pending_manifest_count, + ) + state.update( + { + "candidate_attempts": attempts, + "decision": "retry_infrastructure", + "reason": "infrastructure_retry", + "next_action": "rerun_with_escalated_execution", + "final_response_allowed": False, + "agent_instruction": ( + "Do not produce a final answer. Rerun the same command with escalated execution or repaired " + "runtime permissions; infrastructure retries do not count against the candidate budget." + ), + } + ) + if persist: + write_json(path, state) + return state + + state = guard.guard_state( + results_path, + max_candidates=max_candidates, + stop_files=stop_files, + plateau_threshold=plateau_threshold, + min_delta=plateau_min_delta, + hard_crash_threshold=hard_crash_threshold, + mode=mode, + pending_manifest_count=pending_manifest_count, + ) + if persist: + write_json(path, state) + return state + + +def write_state_if_changed(path: Path, state: Dict[str, Any]) -> bool: + previous = read_json(path) if path.exists() else {} + previous_comparable = dict(previous) + state_comparable = dict(state) + previous_comparable.pop("updated_at", None) + state_comparable.pop("updated_at", None) + if previous_comparable == state_comparable: + if previous.get("updated_at"): + state["updated_at"] = previous["updated_at"] + return False + write_json(path, state) + return True + + +def load_progress_plotter(): + script_path = Path(__file__).with_name("plot_progress.py") + spec = importlib.util.spec_from_file_location("nvflare_autofl_plot_progress", script_path) + if spec is None or spec.loader is None: + raise RuntimeError(f"cannot load Auto-FL progress plotter from {script_path}") + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def write_progress_fallback(path: Path, records: List[RunRecord], mode: str, metric_label: str) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + try: + from PIL import Image, ImageDraw, ImageFont + except ImportError: + path.write_text("Pillow is not installed; progress image unavailable.\n", encoding="utf-8") + return + + width, height = 1000, 560 + margin = 70 + scores = [r.score for r in records if r.score is not None] + image = Image.new("RGB", (width, height), "white") + draw = ImageDraw.Draw(image) + font = ImageFont.load_default() + draw.text( + (margin, 24), + f"Auto-FL Progress ({metric_label}): {len(records)} rows, {len(scores)} scored", + fill="black", + font=font, + ) + draw.line((margin, height - margin, width - margin, height - margin), fill=(80, 80, 80), width=2) + draw.line((margin, margin, margin, height - margin), fill=(80, 80, 80), width=2) + + if scores: + lo, hi = min(scores), max(scores) + if lo == hi: + lo -= 1.0 + hi += 1.0 + span = hi - lo + running_best: Optional[float] = None + last_point: Optional[Tuple[float, float]] = None + denom = max(1, len(records) - 1) + for idx, record in enumerate(records): + if record.score is None: + continue + x = margin + (width - 2 * margin) * idx / denom + y = height - margin - (height - 2 * margin) * (record.score - lo) / span + color = (40, 160, 90) if record.status in {"baseline", "keep"} else (150, 150, 150) + draw.ellipse((x - 5, y - 5, x + 5, y + 5), fill=color, outline="black") + draw.text((x + 6, y - 14), f"{record.name}: {record.score:.3f}", fill=color, font=font) + if better(record.score, running_best, mode): + running_best = record.score + if running_best == record.score: + if last_point: + draw.line((last_point[0], last_point[1], x, y), fill=(40, 160, 90), width=2) + last_point = (x, y) + image.save(path) + + +def write_progress(path: Path, records: List[RunRecord], mode: str, metric_label: str) -> None: + plotter = load_progress_plotter() + try: + plotter.plot_progress(records, path, mode, metric_label) + except (plotter.NoScoredResultsError, plotter.PlotDependencyError): + write_progress_fallback(path, records, mode, metric_label) + + +def write_report(path: Path, config: Dict[str, Any], records: List[RunRecord], args: argparse.Namespace) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + best = None + for record in records: + if record.status in {"baseline", "keep"} and better(record.score, best.score if best else None, args.mode): + best = record + candidate_budget = ( + str(args.max_candidates) if args.max_candidates is not None else "uncapped; runs until manual interruption" + ) + objective = objective_contract(config, args.metric) + lines = [ + "# Auto-FL Report", + "", + f"Objective: optimize `{objective['optimization_metric']}` in `{args.target_env}`.", + f"Requested metric: `{objective['requested_metric']}`.", + f"Metric source: `{metric_source(config)}`.", + f"Metric extraction order: `{', '.join(objective['metric_extraction_order'])}`.", + f"Candidate budget: `{candidate_budget}`.", + f"Config: `{args.autofl_yaml}`.", + f"Fixed budget: `{json.dumps(config.get('budget', {}).get('fixed_training_budget', {}), sort_keys=True)}`.", + "", + "## Leaderboard", + "", + "| Status | Name | Score | Changed files | Manifest | Artifacts | Notes |", + "| --- | --- | ---: | --- | --- | --- | --- |", + ] + for record in records: + score = "" if record.score is None else f"{record.score:.6f}" + lines.append( + f"| {record.status} | {record.name} | {score} | `{record.changed_files}` | " + f"`{record.candidate_manifest}` | `{record.artifacts}` | {record.diff_summary} |" + ) + lines.extend( + [ + "", + "## Artifacts", + "", + f"- Ledger: `{args.results}`", + f"- Progress plot: `{args.progress}`", + f"- Campaign state: `{args.state}`", + "", + "## Outcome", + "", + ] + ) + if best: + lines.append(f"Best retained run: `{best.name}` with `{objective['optimization_metric']}={best.score:.6f}`.") + else: + lines.append("No scored run was retained.") + path.write_text("\n".join(lines) + "\n", encoding="utf-8") + + +def candidate_attempts(records: List[RunRecord]) -> int: + return len( + [ + record + for record in records + if record.status in {"candidate", "keep", "discard", "crash"} and not is_baseline_record(record) + ] + ) + + +def is_baseline_record(record: RunRecord) -> bool: + return record.status == "baseline" + + +def campaign_summary( + autofl_yaml: Path, + results: Path, + state: Path, + progress: Path, + report: Path, + records: List[RunRecord], + state_payload: Optional[Dict[str, Any]] = None, +) -> Dict[str, Any]: + payload = { + "autofl_yaml": str(autofl_yaml.resolve()), + "results": str(results.resolve()), + "state": str(state.resolve()), + "progress": str(progress.resolve()), + "report": str(report.resolve()), + "candidate_attempts": candidate_attempts(records), + } + if state_payload: + for key in [ + "decision", + "reason", + "next_action", + "final_response_allowed", + "candidate_cap", + "candidate_cap_source", + "agent_instruction", + "required_exploration", + ]: + if key in state_payload: + payload[key] = state_payload[key] + return payload + + +def campaign_paths(args: argparse.Namespace, job: Path) -> Dict[str, Path]: + workspace = job.parent + return { + "workspace": workspace, + "autofl_yaml": resolve_output_path(workspace, args.autofl_yaml), + "results": resolve_output_path(workspace, args.results), + "state": resolve_output_path(workspace, args.state), + "progress": resolve_output_path(workspace, args.progress), + "report": resolve_output_path(workspace, args.report), + "output_root": resolve_output_path(workspace, args.output_root), + "snapshot_root": workspace / BEST_SNAPSHOT_ROOT, + } + + +CAMPAIGN_SETTING_NAMES = ( + "metric", + "mode", + "target_env", + "max_candidates", + "autofl_yaml", + "results", + "state", + "progress", + "report", + "output_root", + "plateau_threshold", + "plateau_min_delta", + "hard_crash_threshold", + "stop_file", + "base_args", + "timeout", + "simulator_no_progress_timeout", + "python", + "prefer_synthetic", + "synthetic_train_size", + "synthetic_test_size", +) + +MUTABLE_CAMPAIGN_SETTING_NAMES = { + "max_candidates", + "plateau_threshold", + "plateau_min_delta", + "hard_crash_threshold", + "stop_file", + "timeout", + "simulator_no_progress_timeout", +} + + +def campaign_settings(args: argparse.Namespace) -> Dict[str, Any]: + return {name: getattr(args, name) for name in CAMPAIGN_SETTING_NAMES} + + +def restore_campaign_settings(args: argparse.Namespace, metadata: Dict[str, Any]) -> None: + settings = metadata.get("settings") + if not isinstance(settings, dict): + raise ValueError("campaign metadata is missing settings") + explicit = getattr(args, "_explicit_settings", set()) + changed = False + for name in CAMPAIGN_SETTING_NAMES: + if name not in settings: + continue + if name in explicit: + requested = getattr(args, name) + if name in MUTABLE_CAMPAIGN_SETTING_NAMES: + if settings.get(name) != requested: + settings[name] = requested + changed = True + continue + if requested != settings[name]: + raise ValueError( + f"campaign setting {name} is immutable after initialization: " + f"configured={settings[name]!r}, requested={requested!r}" + ) + setattr(args, name, settings[name]) + if changed: + metadata["updated_at"] = utc_now() + workspace_value = metadata.get("workspace_root") + if not isinstance(workspace_value, str) or not workspace_value: + raise ValueError("campaign metadata is missing workspace_root") + workspace = Path(workspace_value) + write_json(campaign_metadata_path(workspace), metadata) + + +def campaign_timeout(args: argparse.Namespace, schema: Dict[str, Any]) -> Tuple[int, int]: + budget = comparison_budget(schema) + budget_timeout = parse_timeout_setting(budget.get("run_timeout_seconds"), "run_timeout_seconds") + timeout = max(args.timeout, budget_timeout) if budget_timeout is not None else args.timeout + budget_no_progress_timeout = parse_timeout_setting( + budget.get("simulator_no_progress_timeout_seconds"), "simulator_no_progress_timeout_seconds" + ) + no_progress_timeout = ( + budget_no_progress_timeout if budget_no_progress_timeout is not None else args.simulator_no_progress_timeout + ) + return timeout, no_progress_timeout + + +def parse_timeout_setting(value: Any, name: str) -> Optional[int]: + if value is None: + return None + if isinstance(value, bool) or (isinstance(value, float) and not value.is_integer()): + raise ValueError(f"{name} must be a non-negative integer") + try: + parsed = int(value) + except (TypeError, ValueError) as e: + raise ValueError(f"{name} must be a non-negative integer") from e + if parsed < 0: + raise ValueError(f"{name} must be a non-negative integer") + return parsed + + +def import_job_config( + args: argparse.Namespace, + job: Path, + output: Path, + log_path: Path, + timeout: int, +) -> Dict[str, Any]: + del timeout + importer = load_job_importer() + config = importer.import_job_to_autofl_config( + str(job), + workspace_root=str(job.parent), + metric=args.metric, + mode=args.mode, + target_env=args.target_env, + max_candidates=args.max_candidates, + ) + atomic_write_bytes(output, importer.dump_autofl_yaml(config).encode("utf-8")) + atomic_write_bytes( + log_path, + ( + json.dumps( + { + "importer": str(Path(importer.__file__).resolve()), + "job": str(job.resolve()), + "output": str(output.resolve()), + "source_sha256": config.get("import", {}).get("source_sha256"), + }, + indent=2, + sort_keys=True, + ) + + "\n" + ).encode("utf-8"), + ) + return config + + +def campaign_admission_errors(config: Dict[str, Any]) -> List[str]: + errors = [] + support = config.get("import", {}).get("support", {}) + if not isinstance(support, dict) or support.get("status") != "supported": + errors.append("job surface is not deterministically supported") + fixed_budget = config.get("budget", {}).get("fixed_training_budget") + if not isinstance(fixed_budget, dict) or not fixed_budget: + errors.append("fixed comparison budget is unresolved") + unresolved = config.get("unresolved", []) + if isinstance(unresolved, list): + critical_fields = [] + for item in unresolved: + field = item.get("field", "") if isinstance(item, dict) else "" + if field.startswith("budget.fixed_training_budget"): + critical_fields.append(field) + if critical_fields: + errors.append(f"safety-critical fields are unresolved: {', '.join(sorted(set(critical_fields)))}") + return errors + + +def best_retained_record(records: Sequence[RunRecord], mode: str) -> Optional[RunRecord]: + best = None + for record in records: + if record.status in {"baseline", "keep"} and better(record.score, best.score if best else None, mode): + best = record + return best + + +def refresh_campaign_artifacts( + args: argparse.Namespace, + paths: Dict[str, Path], + config: Dict[str, Any], + records: List[RunRecord], + metadata: Dict[str, Any], + *, + pending_manifest: Optional[Path] = None, + next_action: Optional[str] = None, + reason: Optional[str] = None, + required_exploration: Optional[str] = None, +) -> Dict[str, Any]: + pending_manifests = pending_candidate_manifests(paths["workspace"]) + if pending_manifest is not None and pending_manifest not in pending_manifests: + pending_manifests.append(pending_manifest) + write_results(paths["results"], records) + state = write_state( + paths["state"], + paths["results"], + records, + args.max_candidates, + mode=args.mode, + stop_files=resolve_stop_files(paths["workspace"], args.stop_file), + plateau_threshold=args.plateau_threshold, + plateau_min_delta=args.plateau_min_delta, + hard_crash_threshold=args.hard_crash_threshold, + pending_manifest_count=len(pending_manifests), + ) + if state.get("next_action") == "abandon_candidate": + next_action = "abandon_candidate" + reason = state.get("reason") + elif pending_manifests and next_action is None: + pending_status = read_json(pending_manifests[0]).get("status") + next_action = "submit_candidate" if pending_status == "ready_for_external_execution" else "edit_candidate" + reason = "pending_candidates" + if not state.get("final_response_allowed") and next_action: + state["next_action"] = next_action + state["reason"] = reason or next_action + state["agent_instruction"] = f"Do not produce a final answer. Execute `{next_action}` for this campaign." + if required_exploration: + state["required_exploration"] = required_exploration + state["agent_instruction"] = ( + f"{state['agent_instruction']} Prepare and evaluate at least one source-backed server aggregation " + "candidate when compatible with the job contract; otherwise record why it is incompatible." + ) + state.update( + { + "best_candidate": metadata.get("best_candidate"), + "best_source_sha256": metadata.get("best_source_sha256"), + "pending_candidate_manifest": str(pending_manifests[0].resolve()) if pending_manifests else None, + } + ) + write_json(paths["state"], state) + write_progress(paths["progress"], records, args.mode, optimization_metric(config, args.metric)) + write_report(paths["report"], config, records, args) + return state + + +def print_campaign_result( + paths: Dict[str, Path], records: List[RunRecord], state: Dict[str, Any], **extra: Any +) -> None: + payload = campaign_summary( + paths["autofl_yaml"], + paths["results"], + paths["state"], + paths["progress"], + paths["report"], + records, + state, + ) + payload.update(extra) + print(json.dumps(payload, indent=2, sort_keys=True)) + + +def execute_sim_baseline( + args: argparse.Namespace, + job: Path, + paths: Dict[str, Path], + config: Dict[str, Any], + schema: Dict[str, Any], + name: str = "baseline", +) -> RunRecord: + timeout, no_progress_timeout = campaign_timeout(args, schema) + help_text = job_help(args.python, job, job.parent) + return run_job( + JobRun(name=name, args=[], description="baseline", status="baseline"), + python=args.python, + job=job, + cwd=job.parent, + help_text=help_text, + fixed_args=build_fixed_args(config, help_text, schema), + base_args=build_base_args(args, help_text, schema), + output_root=paths["output_root"], + timeout=timeout, + simulator_no_progress_timeout=no_progress_timeout, + metrics=metric_extraction_order(config, args.metric), + config=config, + ) + + +def initialize_campaign(args: argparse.Namespace, job: Path) -> int: + workspace = job.parent + metadata_path = campaign_metadata_path(workspace) + if metadata_path.exists(): + metadata = load_campaign_metadata(workspace, job) + restore_campaign_settings(args, metadata) + paths = campaign_paths(args, job) + records = load_results(paths["results"]) + if args.target_env == "sim" and not any( + record.status == "baseline" and record.score is not None for record in records + ): + config = read_yaml(paths["autofl_yaml"]) + retry_number = sum(is_baseline_record(record) for record in records) + 1 + baseline = execute_sim_baseline( + args, + job, + paths, + config, + load_mutation_schema(workspace), + name=f"baseline_retry_{retry_number}", + ) + records.append(baseline) + metadata["best_score"] = baseline.score + if baseline.score is not None: + metadata["best_candidate"] = baseline.name + metadata["updated_at"] = utc_now() + write_json(metadata_path, metadata) + next_action = "propose_candidate" if baseline.score is not None else "repair_baseline" + state = refresh_campaign_artifacts( + args, paths, config, records, metadata, next_action=next_action, reason="baseline_retried" + ) + print_campaign_result(paths, records, state, initialized=False, baseline_retried=True) + if baseline.status == INFRASTRUCTURE_RETRY: + return 75 + return 0 if baseline.score is not None else 1 + print_campaign_result(paths, records, read_json(paths["state"]), initialized=False) + return 0 + + paths = campaign_paths(args, job) + paths["output_root"].mkdir(parents=True, exist_ok=True) + schema = load_mutation_schema(workspace) + timeout, _ = campaign_timeout(args, schema) + config = import_job_config(args, job, paths["autofl_yaml"], paths["output_root"] / "import.log", timeout) + config = apply_metric_contract(config, args.metric, schema) + args.metric = str(config.get("objective", {}).get("requested_metric") or config["objective"]["metric"]) + config = apply_mutation_schema_contract(config, schema, workspace) + config.setdefault("trust_contract", {})["allowed_create_patterns"] = list(ALLOWED_CREATE_PATTERNS) + write_yaml(paths["autofl_yaml"], config) + admission_errors = campaign_admission_errors(config) + if admission_errors: + raise ValueError( + f"autofl.yaml was generated at {paths['autofl_yaml']}, but baseline execution is unsafe: " + f"{'; '.join(admission_errors)}. Resolve these fields and initialize again." + ) + snapshot_files = create_best_snapshot(workspace, config, paths["snapshot_root"]) + metadata = { + "schema_version": CAMPAIGN_METADATA_SCHEMA_VERSION, + "created_at": utc_now(), + "updated_at": utc_now(), + "job": str(job.resolve()), + "workspace_root": str(workspace.resolve()), + "settings": campaign_settings(args), + "best_candidate": "baseline", + "best_score": None, + "best_source_sha256": source_hash(snapshot_files), + "fixed_budget_sha256": fixed_budget_hash(config), + } + write_json(metadata_path, metadata) + + records: List[RunRecord] = [] + if args.target_env == "sim": + baseline = execute_sim_baseline(args, job, paths, config, schema) + records.append(baseline) + metadata["best_score"] = baseline.score + metadata["updated_at"] = utc_now() + write_json(metadata_path, metadata) + next_action = "propose_candidate" if baseline.score is not None else "repair_baseline" + else: + next_action = "submit_baseline" + + state = refresh_campaign_artifacts( + args, paths, config, records, metadata, next_action=next_action, reason="campaign_initialized" + ) + print_campaign_result(paths, records, state, initialized=True) + if records and records[0].status == INFRASTRUCTURE_RETRY: + return 75 + if args.target_env == "sim" and (not records or records[0].score is None): + return 1 + return 0 + + +def pending_candidate_manifests(workspace: Path) -> List[Path]: + root = workspace / CANDIDATE_ROOT + pending = [] + if not root.exists(): + return pending + for path in sorted(root.glob("*/candidate_manifest.json")): + manifest = read_json(path) + validate_candidate_manifest_identity(path, manifest) + status = manifest.get("status") + if status in {"prepared", "ready_for_external_execution"}: + pending.append(path) + return pending + + +def prepare_candidate(args: argparse.Namespace, job: Path) -> int: + workspace = job.parent + metadata = load_campaign_metadata(workspace, job) + restore_campaign_settings(args, metadata) + paths = campaign_paths(args, job) + records = load_results(paths["results"]) + if not any(record.status == "baseline" and record.score is not None for record in records): + raise ValueError("a scored baseline is required before preparing candidates") + state = read_json(paths["state"]) + if state.get("final_response_allowed"): + raise ValueError(f"campaign is already final: {state.get('reason')}") + pending = pending_candidate_manifests(workspace) + if pending: + raise ValueError(f"campaign already has a pending candidate: {pending[0]}") + candidate_id = validate_candidate_id(args.name or "") + if not args.hypothesis: + raise ValueError("--hypothesis is required for prepare") + manifest_path = candidate_manifest_path(workspace, candidate_id) + candidate_dir = manifest_path.parent + if candidate_dir.exists(): + raise ValueError(f"candidate already exists: {candidate_id}") + + best_source, best_files = load_best_snapshot(paths["snapshot_root"]) + if source_hash(best_files) != metadata.get("best_source_sha256"): + raise ValueError("campaign best-source hash is stale") + if not workspace_matches_snapshot(workspace, best_source, best_files): + raise ValueError("job workspace differs from the recorded best candidate; reconcile edits before preparing") + + draft_source = candidate_dir / "source" + shutil.copytree(best_source, draft_source) + manifest = { + "schema_version": CANDIDATE_MANIFEST_SCHEMA_VERSION, + "candidate_id": candidate_id, + "name": candidate_id, + "hypothesis": args.hypothesis, + "status": "prepared", + "created_at": utc_now(), + "updated_at": utc_now(), + "workspace_root": str(workspace.resolve()), + "base_candidate": metadata.get("best_candidate"), + "base_source_sha256": source_hash(best_files), + "fixed_budget_sha256": metadata.get("fixed_budget_sha256"), + "objective": {"metric": args.metric, "mode": args.mode}, + "environment": args.target_env, + "run_args": shlex.split(args.run_args), + "changed_files": [], + "patch_sha256": "", + "candidate_source_sha256": source_hash(best_files), + "provenance": { + "job": str(job.resolve()), + "autofl_yaml": str(paths["autofl_yaml"].resolve()), + "import_source_sha256": read_yaml(paths["autofl_yaml"]).get("import", {}).get("source_sha256"), + }, + "artifacts": {}, + "result": {}, + } + write_json(manifest_path, manifest) + config = read_yaml(paths["autofl_yaml"]) + state = refresh_campaign_artifacts( + args, + paths, + config, + records, + metadata, + pending_manifest=manifest_path, + next_action="edit_candidate", + reason="candidate_prepared", + ) + print_campaign_result( + paths, + records, + state, + candidate_manifest=str(manifest_path.resolve()), + candidate_source=str(draft_source.resolve()), + ) + return 0 + + +def validate_candidate_for_evaluation( + args: argparse.Namespace, + job: Path, + metadata: Dict[str, Any], + paths: Dict[str, Path], + manifest_path: Path, +) -> Tuple[Dict[str, Any], Dict[str, Any], Path, Dict[str, str], List[str], List[str], str]: + manifest = load_candidate_manifest(manifest_path) + if Path(str(manifest.get("workspace_root") or "")).resolve() != job.parent.resolve(): + raise ValueError("candidate manifest belongs to a different job workspace") + config = read_yaml(paths["autofl_yaml"]) + best_source, best_files = load_best_snapshot(paths["snapshot_root"]) + best_hash = source_hash(best_files) + if manifest.get("base_source_sha256") != best_hash or metadata.get("best_source_sha256") != best_hash: + raise ValueError("candidate was prepared from a stale best candidate") + if manifest.get("fixed_budget_sha256") != metadata.get("fixed_budget_sha256"): + raise ValueError("candidate fixed-budget provenance is stale") + if not workspace_matches_snapshot(job.parent, best_source, best_files): + raise ValueError("job workspace differs from the recorded best candidate") + draft_source = manifest_path.parent / "source" + changed, created = candidate_changes(job.parent, config, best_source, best_files, draft_source) + recorded_changed = validated_manifest_paths(manifest, job.parent, "changed_files") + recorded_created = validated_manifest_paths(manifest, job.parent, "created_files") + if (recorded_changed or recorded_created) and (recorded_changed != changed or recorded_created != created): + raise ValueError("candidate manifest source lists do not match the deterministic candidate diff") + run_args = manifest.get("run_args") + if not isinstance(run_args, list) or not all(isinstance(item, str) for item in run_args): + raise ValueError("candidate manifest run_args must be a list of strings") + allowed, reason = candidate_args_allowed(run_args, load_mutation_schema(job.parent)) + if not allowed: + raise ValueError(reason) + allowed, reason = candidate_preserves_fixed_args(run_args, config, load_mutation_schema(job.parent)) + if not allowed: + raise ValueError(reason) + if not changed and not run_args: + raise ValueError("candidate has no source changes or run arguments") + patch = render_candidate_patch(best_source, draft_source, changed) + return manifest, config, best_source, best_files, changed, created, patch + + +def update_config_for_kept_sources(config: Dict[str, Any], created: Sequence[str]) -> None: + if not created: + return + trust_paths = config.setdefault("trust_contract", {}).setdefault("allowed_edit_paths", []) + for relative in created: + if relative not in trust_paths: + trust_paths.append(relative) + + +def candidate_campaign_config( + candidate_config: Dict[str, Any], current_config: Dict[str, Any], args: argparse.Namespace, schema: Dict[str, Any] +) -> Dict[str, Any]: + candidate_config = apply_metric_contract(candidate_config, args.metric, schema) + candidate_config["objective"] = dict(current_config.get("objective", {})) + current_paths = current_config.get("trust_contract", {}).get("allowed_edit_paths", []) or [] + trust_paths = candidate_config.setdefault("trust_contract", {}).setdefault("allowed_edit_paths", []) + for path in current_paths: + if path not in trust_paths: + trust_paths.append(path) + candidate_config["trust_contract"]["allowed_create_patterns"] = allowed_create_patterns(current_config) + return candidate_config + + +def finalize_candidate_result( + args: argparse.Namespace, + job: Path, + metadata: Dict[str, Any], + paths: Dict[str, Path], + config: Dict[str, Any], + manifest_path: Path, + manifest: Dict[str, Any], + best_source: Path, + best_files: Dict[str, str], + changed: List[str], + created: List[str], + patch: str, + record: RunRecord, +) -> Tuple[List[RunRecord], Dict[str, Any]]: + rollback_files: Dict[Path, Optional[bytes]] = {} + staged_snapshot = None + previous_snapshot = None + try: + records = load_results(paths["results"]) + previous_best = best_retained_record(records, args.mode) + if record.status == "candidate": + record.status = ( + "keep" if better(record.score, previous_best.score if previous_best else None, args.mode) else "discard" + ) + patch_path = manifest_path.parent / "candidate.patch" + rollback_files = capture_file_versions( + [ + paths["autofl_yaml"], + manifest_path, + campaign_metadata_path(job.parent), + paths["results"], + paths["state"], + paths["progress"], + paths["report"], + patch_path, + ] + ) + patch_path.write_text(patch, encoding="utf-8") + patch_sha256 = sha256_bytes(patch.encode("utf-8")) + record.changed_files = ",".join(changed) if changed else "none" + record.diff_summary = str(manifest.get("hypothesis") or "candidate") + record.candidate_manifest = str(manifest_path.resolve()) + record.base_candidate = str(manifest.get("base_candidate") or "") + record.patch_sha256 = patch_sha256 + + if record.status == "keep": + update_config_for_kept_sources(config, created) + staged_snapshot, snapshot_files = stage_best_snapshot(job.parent, config, paths["snapshot_root"]) + write_yaml(paths["autofl_yaml"], config) + previous_snapshot = activate_best_snapshot(paths["snapshot_root"], staged_snapshot) + staged_snapshot = None + metadata.update( + { + "best_candidate": record.name, + "best_score": record.score, + "best_source_sha256": source_hash(snapshot_files), + "updated_at": utc_now(), + } + ) + else: + restore_best_source(job.parent, best_source, best_files, changed) + + manifest.update( + { + "status": record.status, + "updated_at": utc_now(), + "changed_files": changed, + "created_files": created, + "patch_sha256": patch_sha256, + "artifacts": {"patch": str(patch_path.resolve()), "run": record.artifacts}, + "result": { + "score": record.score, + "metric_name": record.metric_name, + "metric_source": record.metric_source, + "metric_artifact": record.metric_artifact, + "runtime_seconds": record.runtime_seconds, + "run_command": record.run_command, + "failure_reason": record.failure_reason, + }, + } + ) + write_json(manifest_path, manifest) + write_json(campaign_metadata_path(job.parent), metadata) + records.append(record) + state = refresh_campaign_artifacts(args, paths, config, records, metadata) + except BaseException as error: + try: + if previous_snapshot is not None: + rollback_best_snapshot(paths["snapshot_root"], previous_snapshot) + previous_snapshot = None + restore_best_source(job.parent, paths["snapshot_root"] / "source", best_files, changed) + if rollback_files: + restore_file_versions(rollback_files) + except BaseException as rollback_error: + raise RuntimeError( + f"candidate finalization failed ({error}); automatic workspace rollback also failed ({rollback_error})" + ) from rollback_error + raise + finally: + if staged_snapshot is not None: + shutil.rmtree(staged_snapshot, ignore_errors=True) + + if previous_snapshot is not None: + shutil.rmtree(previous_snapshot, ignore_errors=True) + return records, state + + +def evaluate_candidate(args: argparse.Namespace, job: Path) -> int: + workspace = job.parent + metadata = load_campaign_metadata(workspace, job) + restore_campaign_settings(args, metadata) + paths = campaign_paths(args, job) + manifest_path = Path(args.manifest).resolve() if args.manifest else None + if manifest_path is None: + pending = pending_candidate_manifests(workspace) + if len(pending) != 1: + raise ValueError("--manifest is required when there is not exactly one pending candidate") + manifest_path = pending[0] + manifest, config, best_source, best_files, changed, created, patch = validate_candidate_for_evaluation( + args, job, metadata, paths, manifest_path + ) + patch_path = manifest_path.parent / "candidate.patch" + patch_path.write_text(patch, encoding="utf-8") + manifest.update( + { + "updated_at": utc_now(), + "changed_files": changed, + "created_files": created, + "patch_sha256": sha256_bytes(patch.encode("utf-8")), + "candidate_source_sha256": source_hash(file_map(manifest_path.parent / "source")), + } + ) + write_json(manifest_path, manifest) + schema = load_mutation_schema(workspace) + timeout, no_progress_timeout = campaign_timeout(args, schema) + candidate_config_path = manifest_path.parent / "candidate_autofl.yaml" + try: + apply_candidate_source(workspace, manifest_path.parent / "source", changed) + candidate_config = import_job_config( + args, + job, + candidate_config_path, + manifest_path.parent / "import.log", + timeout, + ) + if fixed_budget_hash(candidate_config) != metadata.get("fixed_budget_sha256"): + raise ValueError("candidate changes budget.fixed_training_budget") + candidate_config = candidate_campaign_config(candidate_config, config, args, schema) + except BaseException: + restore_best_source(workspace, best_source, best_files, changed) + raise + + if args.target_env != "sim": + manifest["status"] = "ready_for_external_execution" + manifest["updated_at"] = utc_now() + write_json(manifest_path, manifest) + records = load_results(paths["results"]) + state = refresh_campaign_artifacts( + args, + paths, + config, + records, + metadata, + pending_manifest=manifest_path, + next_action="submit_candidate", + reason="candidate_validated", + ) + print_campaign_result( + paths, + records, + state, + candidate_manifest=str(manifest_path.resolve()), + job=str(job.resolve()), + ) + return 0 + + try: + help_text = job_help(args.python, job, workspace) + run_record = run_job( + JobRun( + name=str(manifest["candidate_id"]), + args=list(manifest.get("run_args") or []), + description=str(manifest.get("hypothesis") or "candidate"), + ), + python=args.python, + job=job, + cwd=workspace, + help_text=help_text, + fixed_args=build_fixed_args(config, help_text, schema), + base_args=build_base_args(args, help_text, schema), + output_root=paths["output_root"], + timeout=timeout, + simulator_no_progress_timeout=no_progress_timeout, + metrics=metric_extraction_order(config, args.metric), + config=config, + ) + except BaseException: + restore_best_source(workspace, best_source, best_files, changed) + raise + if run_record.status == INFRASTRUCTURE_RETRY: + restore_best_source(workspace, best_source, best_files, changed) + manifest["status"] = "prepared" + manifest["result"] = {"failure_reason": run_record.failure_reason} + manifest["updated_at"] = utc_now() + write_json(manifest_path, manifest) + records = load_results(paths["results"]) + state = refresh_campaign_artifacts( + args, + paths, + config, + records, + metadata, + pending_manifest=manifest_path, + next_action="rerun_with_escalated_execution", + reason="infrastructure_retry", + ) + print_campaign_result(paths, records, state, candidate_manifest=str(manifest_path.resolve())) + return 75 + + records, state = finalize_candidate_result( + args, + job, + metadata, + paths, + candidate_config, + manifest_path, + manifest, + best_source, + best_files, + changed, + created, + patch, + run_record, + ) + print_campaign_result(paths, records, state, candidate_manifest=str(manifest_path.resolve())) + return 0 + + +def abandon_candidate(args: argparse.Namespace, job: Path) -> int: + workspace = job.parent + metadata = load_campaign_metadata(workspace, job) + restore_campaign_settings(args, metadata) + paths = campaign_paths(args, job) + manifest_path = Path(args.manifest).resolve() if args.manifest else None + if manifest_path is None: + pending = pending_candidate_manifests(workspace) + if len(pending) != 1: + raise ValueError("--manifest is required when there is not exactly one pending candidate") + manifest_path = pending[0] + manifest = load_candidate_manifest(manifest_path) + best_source, best_files = load_best_snapshot(paths["snapshot_root"]) + config = read_yaml(paths["autofl_yaml"]) + draft_source = manifest_path.parent / "source" + changed, created = candidate_changes( + workspace, + config, + best_source, + best_files, + draft_source, + allow_materialized=manifest.get("status") == "ready_for_external_execution", + ) + recorded_changed = validated_manifest_paths(manifest, workspace, "changed_files") + recorded_created = validated_manifest_paths(manifest, workspace, "created_files") + if (recorded_changed or recorded_created) and (recorded_changed != changed or recorded_created != created): + raise ValueError("candidate manifest source lists do not match the deterministic candidate diff") + if manifest.get("status") == "ready_for_external_execution": + if not workspace_matches_snapshot(workspace, draft_source, file_map(draft_source)): + raise ValueError("materialized candidate source changed after validation") + restore_best_source(workspace, best_source, best_files, changed) + elif not workspace_matches_snapshot(workspace, best_source, best_files): + raise ValueError("job workspace differs from the recorded best candidate") + manifest["status"] = "abandoned" + manifest["updated_at"] = utc_now() + write_json(manifest_path, manifest) + records = load_results(paths["results"]) + state = refresh_campaign_artifacts( + args, + paths, + config, + records, + metadata, + next_action="propose_candidate", + reason="candidate_abandoned", + ) + print_campaign_result(paths, records, state, candidate_manifest=str(manifest_path.resolve())) + return 0 + + +def suggest_candidates(args: argparse.Namespace, job: Path) -> int: + metadata = load_campaign_metadata(job.parent, job) + restore_campaign_settings(args, metadata) + if args.limit < 1: + raise ValueError("--limit must be positive") + config = read_yaml(campaign_paths(args, job)["autofl_yaml"]) + help_text = job_help(args.python, job, job.parent) + suggestions = [ + {"name": candidate.name, "run_args": candidate.args, "hypothesis": candidate.description} + for candidate in candidate_plan( + config, + help_text, + args.limit, + load_mutation_schema(job.parent), + ) + ] + print(json.dumps({"suggestions": suggestions}, indent=2, sort_keys=True)) + return 0 + + +def record_external_result(args: argparse.Namespace, job: Path) -> int: + workspace = job.parent + metadata = load_campaign_metadata(workspace, job) + restore_campaign_settings(args, metadata) + paths = campaign_paths(args, job) + config = read_yaml(paths["autofl_yaml"]) + artifact_path = Path(args.external_artifacts).resolve() if args.external_artifacts else None + score = args.score + evidence = None + if score is None and artifact_path: + evidence = extract_metric_evidence(artifact_path, metric_extraction_order(config, args.metric)) + score = evidence.score if evidence else None + elif score is not None: + evidence = MetricEvidence( + score=score, + metric_name=optimization_metric(config, args.metric), + source="provided", + artifact=str(artifact_path or ""), + ) + if args.literature: + records = load_results(paths["results"]) + name = f"literature_review_{sum(record.status == 'literature' for record in records) + 1}" + records.append( + RunRecord( + status="literature", + name=name, + score=None, + runtime_seconds=0.0, + changed_files="none", + diff_summary=args.hypothesis or "source-backed literature review", + run_command="agent literature review", + artifacts=str(artifact_path or ""), + ) + ) + state = refresh_campaign_artifacts( + args, + paths, + config, + records, + metadata, + next_action="propose_candidate", + reason="literature_review_recorded", + required_exploration="source_backed_server_aggregation", + ) + print_campaign_result(paths, records, state, literature_event=name) + return 0 + if args.baseline: + records = load_results(paths["results"]) + if any( + record.status == "baseline" + or (record.status in {"candidate", "keep", "discard", "crash"} and not is_baseline_record(record)) + for record in records + ): + raise ValueError("external baseline can only be recorded before a baseline or campaign candidate") + if score is None: + raise ValueError("a score or extractable --artifacts path is required for the baseline") + record = RunRecord( + status="baseline", + name="baseline", + score=score, + runtime_seconds=0.0, + changed_files="none", + diff_summary="external baseline", + run_command=f"nvflare job id={args.job_id or 'unreported'}", + artifacts=str(artifact_path or ""), + metric_name=evidence.metric_name if evidence else "", + metric_source=evidence.source if evidence else "", + metric_artifact=evidence.artifact if evidence else "", + ) + records.append(record) + metadata["best_score"] = score + metadata["updated_at"] = utc_now() + write_json(campaign_metadata_path(workspace), metadata) + state = refresh_campaign_artifacts( + args, + paths, + config, + records, + metadata, + next_action="propose_candidate", + reason="baseline_recorded", + ) + print_campaign_result(paths, records, state, job_id=args.job_id) + return 0 + + manifest_path = Path(args.manifest).resolve() if args.manifest else None + if manifest_path is None: + pending = pending_candidate_manifests(workspace) + if len(pending) != 1: + raise ValueError("--manifest is required when there is not exactly one pending candidate") + manifest_path = pending[0] + manifest = load_candidate_manifest(manifest_path) + if manifest.get("status") != "ready_for_external_execution": + raise ValueError("external candidate must be validated before its result is recorded") + draft_source = manifest_path.parent / "source" + draft_files = file_map(draft_source) + if source_hash(draft_files) != manifest.get("candidate_source_sha256") or not workspace_matches_snapshot( + workspace, draft_source, draft_files + ): + raise ValueError("materialized candidate source changed after validation") + best_source, best_files = load_best_snapshot(paths["snapshot_root"]) + changed, created = candidate_changes( + workspace, config, best_source, best_files, draft_source, allow_materialized=True + ) + validate_manifest_change_set(manifest, workspace, changed, created) + patch_path = manifest_path.parent / "candidate.patch" + patch = patch_path.read_text(encoding="utf-8") if patch_path.exists() else "" + candidate_config_path = manifest_path.parent / "candidate_autofl.yaml" + candidate_config = read_yaml(candidate_config_path) if candidate_config_path.exists() else config + candidate_config = candidate_campaign_config(candidate_config, config, args, load_mutation_schema(workspace)) + status = "crash" if args.failure_reason or score is None else "candidate" + record = RunRecord( + status=status, + name=str(manifest["candidate_id"]), + score=score, + runtime_seconds=0.0, + changed_files="none", + diff_summary=str(manifest.get("hypothesis") or "candidate"), + run_command=f"nvflare job id={args.job_id or 'unreported'}", + artifacts=str(artifact_path or ""), + failure_reason=args.failure_reason or ("metric not found" if score is None else ""), + metric_name=evidence.metric_name if evidence else "", + metric_source=evidence.source if evidence else "", + metric_artifact=evidence.artifact if evidence else "", + ) + records, state = finalize_candidate_result( + args, + job, + metadata, + paths, + candidate_config, + manifest_path, + manifest, + best_source, + best_files, + changed, + created, + patch, + record, + ) + updated_manifest = read_json(manifest_path) + updated_manifest.setdefault("artifacts", {})["job_id"] = args.job_id + write_json(manifest_path, updated_manifest) + print_campaign_result(paths, records, state, candidate_manifest=str(manifest_path.resolve()), job_id=args.job_id) + return 0 + + +def show_campaign_status(args: argparse.Namespace, job: Path) -> int: + metadata = load_campaign_metadata(job.parent, job) + restore_campaign_settings(args, metadata) + paths = campaign_paths(args, job) + records = load_results(paths["results"]) + pending = pending_candidate_manifests(job.parent) + state = write_state( + paths["state"], + paths["results"], + records, + args.max_candidates, + mode=args.mode, + stop_files=resolve_stop_files(job.parent, args.stop_file), + plateau_threshold=args.plateau_threshold, + plateau_min_delta=args.plateau_min_delta, + hard_crash_threshold=args.hard_crash_threshold, + pending_manifest_count=len(pending), + persist=False, + ) + state.update( + { + "best_candidate": metadata.get("best_candidate"), + "best_source_sha256": metadata.get("best_source_sha256"), + "pending_candidate_manifest": str(pending[0].resolve()) if pending else None, + } + ) + write_state_if_changed(paths["state"], state) + print_campaign_result(paths, records, state, campaign=metadata) + return 0 + + +def validate_args(args: argparse.Namespace) -> None: + if args.max_candidates is not None and args.max_candidates < 1: + raise ValueError("--max-candidates must be positive when provided") + if args.plateau_threshold < 1: + raise ValueError("--plateau-threshold must be positive") + if args.plateau_min_delta < 0: + raise ValueError("--plateau-min-delta must be non-negative") + if args.hard_crash_threshold < 0: + raise ValueError("--hard-crash-threshold must be non-negative") + + +def main(argv: Optional[Sequence[str]] = None) -> int: + args = parse_args(argv) + try: + validate_args(args) + job = Path(args.job).resolve() + if not job.is_file(): + raise ValueError(f"job.py does not exist: {job}") + actions = { + "initialize": initialize_campaign, + "prepare": prepare_candidate, + "evaluate": evaluate_candidate, + "abandon": abandon_candidate, + "suggest": suggest_candidates, + "record": record_external_result, + "status": show_campaign_status, + } + return actions[args.action](args, job) + except (OSError, RuntimeError, ValueError) as e: + print(f"Auto-FL {args.action} failed: {e}", file=sys.stderr) + return 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/skills/nvflare-autofl/tests/helper_scripts.md b/skills/nvflare-autofl/tests/helper_scripts.md new file mode 100644 index 0000000000..f7a33e6e37 --- /dev/null +++ b/skills/nvflare-autofl/tests/helper_scripts.md @@ -0,0 +1,8 @@ +# Helper Script Coverage + +The bundled campaign helpers are covered by focused repository tests: + +- `scripts/campaign_guard.py`: + `tests/unit_test/tool/autofl_skill_campaign_guard_test.py` +- `scripts/run_job_campaign.py`: + `tests/unit_test/tool/autofl_skill_runner_test.py` diff --git a/tests/unit_test/tool/agent/seed_skill_packaging_test.py b/tests/unit_test/tool/agent/seed_skill_packaging_test.py index c80e5baef8..aed3291fa0 100644 --- a/tests/unit_test/tool/agent/seed_skill_packaging_test.py +++ b/tests/unit_test/tool/agent/seed_skill_packaging_test.py @@ -27,6 +27,7 @@ from nvflare.tool.agent.skill_manifest import build_skill_manifest, copy_released_skills_to_bundle SEED_SKILLS = { + "nvflare-autofl", "nvflare-orient", "nvflare-convert-pytorch", "nvflare-diagnose-job", @@ -96,6 +97,8 @@ def test_seed_skills_install_into_codex_and_claude_temp_targets(tmp_path): _assert_diagnose_runtime_payload(claude_target / "nvflare-diagnose-job") _assert_analysis_payload_present(codex_target / "nvflare-diagnose-job") _assert_analysis_payload_present(claude_target / "nvflare-diagnose-job") + _assert_autofl_runtime_payload(codex_target / "nvflare-autofl", tmp_path / "codex-job") + _assert_autofl_runtime_payload(claude_target / "nvflare-autofl", tmp_path / "claude-job") codex_list = list_skills(agent="codex", target_dir=codex_target, source=source) claude_list = list_skills(agent="claude", target_dir=claude_target, source=source) @@ -287,6 +290,23 @@ def _assert_diagnose_runtime_payload(skill_dir: Path) -> None: assert skill_dir.joinpath("references", "failure-patterns.md").is_file() +def _assert_autofl_runtime_payload(skill_dir: Path, unrelated_job_dir: Path) -> None: + runner = skill_dir / "scripts" / "run_job_campaign.py" + importer = skill_dir / "scripts" / "job_importer.py" + assert skill_dir.joinpath("SKILL.md").is_file() + assert runner.is_file() + assert importer.is_file() + unrelated_job_dir.mkdir() + result = subprocess.run( + [sys.executable, str(runner), "--help"], + cwd=unrelated_job_dir, + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 0, result.stdout + result.stderr + + def _assert_analysis_payload_present(skill_dir: Path) -> None: assert skill_dir.joinpath("evals", "evals.json").is_file() assert skill_dir.joinpath("evals", "files", "poc_component_not_authorized.log").is_file() diff --git a/tests/unit_test/tool/agent_skill_checks/frontmatter_test.py b/tests/unit_test/tool/agent_skill_checks/frontmatter_test.py index 61b942482b..484468d25b 100644 --- a/tests/unit_test/tool/agent_skill_checks/frontmatter_test.py +++ b/tests/unit_test/tool/agent_skill_checks/frontmatter_test.py @@ -66,6 +66,55 @@ def test_validate_skill_dir_accepts_fixture_skill(): assert result.issues == () +def test_validate_skill_dir_accepts_spec_metadata_extension(tmp_path): + skill_dir = tmp_path / "nvflare-spec-skill" + skill_dir.mkdir() + skill_dir.joinpath("SKILL.md").write_text( + "---\n" + "name: nvflare-spec-skill\n" + "description: Test specification-compatible metadata.\n" + "license: Apache-2.0\n" + "compatibility: Requires NVFLARE.\n" + "metadata:\n" + ' min_flare_version: "2.8.0"\n' + " blast_radius: runs_simulator\n" + ' skill_version: "0.1.0"\n' + "---\n\n" + "# Test Skill\n", + encoding="utf-8", + ) + + result = validate_skill_dir(skill_dir) + + assert result.ok + assert result.metadata["min_flare_version"] == "2.8.0" + assert result.metadata["blast_radius"] == "runs_simulator" + assert result.metadata["skill_version"] == "0.1.0" + + +def test_validate_skill_dir_rejects_conflicting_legacy_and_nested_metadata(tmp_path): + skill_dir = tmp_path / "nvflare-conflicting-skill" + skill_dir.mkdir() + skill_dir.joinpath("SKILL.md").write_text( + "---\n" + "name: nvflare-conflicting-skill\n" + "description: Test conflicting metadata.\n" + 'min_flare_version: "2.8.0"\n' + "blast_radius: read_only\n" + "metadata:\n" + " blast_radius: runs_simulator\n" + "---\n\n" + "# Test Skill\n", + encoding="utf-8", + ) + + result = validate_skill_dir(skill_dir) + + assert not result.ok + assert _issue_codes(result) == {"skill-frontmatter-invalid"} + assert "conflicting values" in result.issues[0].message + + def test_validate_skill_dir_requires_directory_name_to_match_frontmatter(tmp_path): skill_dir = _write_skill(tmp_path, "nvflare-other-name", name="nvflare-example-skill") diff --git a/tests/unit_test/tool/autofl_skill_campaign_guard_test.py b/tests/unit_test/tool/autofl_skill_campaign_guard_test.py new file mode 100644 index 0000000000..8576c2efc1 --- /dev/null +++ b/tests/unit_test/tool/autofl_skill_campaign_guard_test.py @@ -0,0 +1,280 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import csv +import importlib.util +import json +import subprocess +import sys +from pathlib import Path + +RESULT_FIELDS = [ + "status", + "name", + "score", + "runtime_seconds", + "changed_files", + "diff_summary", + "run_command", + "artifacts", + "failure_reason", + "candidate_manifest", + "base_candidate", + "patch_sha256", +] + + +def _load_guard(): + repo_root = Path(__file__).parents[3] + guard_path = repo_root / "skills" / "nvflare-autofl" / "scripts" / "campaign_guard.py" + spec = importlib.util.spec_from_file_location("nvflare_autofl_skill_campaign_guard_test", guard_path) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def _write_results(path, rows): + with path.open("w", encoding="utf-8", newline="") as f: + writer = csv.DictWriter(f, fieldnames=RESULT_FIELDS, delimiter="\t") + writer.writeheader() + for row in rows: + writer.writerow(row) + + +def _row(status, name, score="", diff_summary="candidate", run_command="python job.py"): + return { + "status": status, + "name": name, + "score": score, + "runtime_seconds": "10.0", + "changed_files": "none", + "diff_summary": diff_summary, + "run_command": run_command, + "artifacts": "/tmp/run", + "failure_reason": "", + "candidate_manifest": "", + "base_candidate": "", + "patch_sha256": "", + } + + +def test_guard_continues_uncapped_before_plateau(): + guard = _load_guard() + rows = [ + _row("baseline", "baseline", "0.85"), + _row("discard", "batch_size_64", "0.851"), + _row("discard", "lr_0p01", "0.850"), + ] + + state = guard.guard_state_for_rows(rows, plateau_threshold=4) + + assert state["schema_version"] == "nvflare.autofl.campaign_state.v1" + assert state["decision"] == "continue" + assert state["next_action"] == "propose_candidate" + assert state["final_response_allowed"] is False + assert state["candidate_cap"] is None + assert state["candidate_cap_source"] == "uncapped" + assert state["candidate_attempts"] == 2 + assert state["best_score"] == 0.85 + + +def test_guard_routes_plateau_to_literature_without_finalizing(): + guard = _load_guard() + rows = [_row("baseline", "baseline", "0.85")] + rows.extend(_row("discard", f"candidate_{idx}", "0.840") for idx in range(4)) + + state = guard.guard_state_for_rows(rows, plateau_threshold=4) + + assert state["decision"] == "continue" + assert state["reason"] == "plateau_literature" + assert state["next_action"] == "run_literature_loop" + assert state["final_response_allowed"] is False + assert state["best_score"] == 0.85 + assert state["required_exploration"] == "source_backed_server_aggregation" + assert "server aggregation candidate" in state["agent_instruction"] + assert "Do not produce a final answer" in state["agent_instruction"] + + +def test_guard_literature_event_resets_plateau_clock(): + guard = _load_guard() + rows = [_row("baseline", "baseline", "0.85")] + rows.extend(_row("discard", f"before_lit_{idx}", "0.840") for idx in range(4)) + rows.append(_row("literature", "literature_review", "", diff_summary="literature review")) + rows.extend(_row("discard", f"after_lit_{idx}", "0.841") for idx in range(2)) + + state = guard.guard_state_for_rows(rows, plateau_threshold=4) + + assert state["reason"] == "continue" + assert state["next_action"] == "propose_candidate" + assert state["plateau"]["last_literature_event_index"] == 5 + assert state["plateau"]["scored_since_reset"] == 2 + + +def test_guard_counts_candidate_with_baseline_in_description(): + guard = _load_guard() + rows = [ + _row("baseline", "baseline", "0.85"), + _row("discard", "weighted_rerun", "0.8503", diff_summary="weighted baseline escalated rerun"), + ] + + state = guard.guard_state_for_rows(rows, max_candidates=1) + + assert state["candidate_attempts"] == 1 + assert state["decision"] == "stop" + assert state["reason"] == "candidate_cap_exhausted" + assert state["candidate_cap_source"] == "explicit" + + +def test_guard_ignores_ambient_candidate_cap(monkeypatch): + guard = _load_guard() + monkeypatch.setenv("AUTOFL_MAX_CANDIDATES", "1") + rows = [ + _row("baseline", "baseline", "0.85"), + _row("discard", "candidate_1", "0.84"), + ] + + state = guard.guard_state_for_rows(rows) + + assert state["decision"] == "continue" + assert state["reason"] == "continue" + assert state["candidate_cap"] is None + assert state["candidate_cap_source"] == "uncapped" + assert state["final_response_allowed"] is False + + +def test_guard_cli_is_diagnostic_only(tmp_path): + repo_root = Path(__file__).parents[3] + guard_path = repo_root / "skills" / "nvflare-autofl" / "scripts" / "campaign_guard.py" + results_path = tmp_path / "results.tsv" + state_path = tmp_path / "state.json" + state_path.write_text('{"authoritative": true}\n', encoding="utf-8") + _write_results( + results_path, + [ + _row("baseline", "baseline", "0.85"), + _row("discard", "candidate_1", "0.840"), + _row("discard", "candidate_2", "0.841"), + ], + ) + + proc = subprocess.run( + [ + sys.executable, + str(guard_path), + str(results_path), + "--plateau-threshold", + "2", + "--format", + "json", + ], + text=True, + capture_output=True, + check=True, + ) + + payload = json.loads(proc.stdout) + assert json.loads(state_path.read_text(encoding="utf-8")) == {"authoritative": True} + assert payload["next_action"] == "run_literature_loop" + + +def test_guard_resolves_default_and_custom_stop_files_from_results_directory(tmp_path): + repo_root = Path(__file__).parents[3] + guard_path = repo_root / "skills" / "nvflare-autofl" / "scripts" / "campaign_guard.py" + results_path = tmp_path / "results.tsv" + _write_results(results_path, [_row("baseline", "baseline", "0.85")]) + tmp_path.joinpath("STOP_AUTOFL").touch() + + default_proc = subprocess.run( + [sys.executable, str(guard_path), str(results_path), "--format", "json"], + cwd=repo_root, + text=True, + capture_output=True, + check=True, + ) + custom_proc = subprocess.run( + [ + sys.executable, + str(guard_path), + str(results_path), + "--stop-file", + "CUSTOM_STOP", + "--format", + "json", + ], + cwd=repo_root, + text=True, + capture_output=True, + check=True, + ) + + assert json.loads(default_proc.stdout)["reason"] == "manual_stop_file" + assert json.loads(custom_proc.stdout)["reason"] == "continue" + + +def test_numeric_crashes_and_pending_candidates_do_not_change_retained_best_or_plateau(): + guard = _load_guard() + rows = [ + _row("baseline", "baseline", "0.85"), + _row("crash", "crashed_candidate", "0.99"), + _row("candidate", "pending_candidate", "0.98"), + _row("discard", "discarded_candidate", "0.84"), + ] + + state = guard.guard_state_for_rows(rows, plateau_threshold=1) + + assert state["best_score"] == 0.85 + assert state["scored_attempts"] == 1 + assert state["plateau"]["scored_since_reset"] == 1 + assert state["reason"] == "pending_candidates" + assert state["next_action"] == "edit_candidate" + + +def test_stop_file_requires_pending_candidate_abandonment_before_final_report(tmp_path): + guard = _load_guard() + stop_file = tmp_path / "STOP_AUTOFL" + stop_file.touch() + rows = [_row("baseline", "baseline", "0.85"), _row("discard", "candidate_1", "0.84")] + + state = guard.guard_state_for_rows( + rows, + max_candidates=1, + stop_files=[str(stop_file)], + pending_manifest_count=1, + ) + + assert state["decision"] == "stop" + assert state["reason"] == "manual_stop_pending_candidate" + assert state["next_action"] == "abandon_candidate" + assert state["final_response_allowed"] is False + + +def test_continuous_campaign_reference_documents_only_emitted_actions(): + repo_root = Path(__file__).parents[3] + reference = repo_root / "skills" / "nvflare-autofl" / "references" / "continuous-campaigns.md" + text = reference.read_text(encoding="utf-8") + actions = { + "repair_baseline", + "edit_candidate", + "abandon_candidate", + "propose_candidate", + "submit_baseline", + "submit_candidate", + "rerun_with_escalated_execution", + "run_literature_loop", + "final_report", + } + + assert all(f"`{action}`" in text for action in actions) + assert "`evaluate_candidate`" not in text diff --git a/tests/unit_test/tool/autofl_skill_job_importer_test.py b/tests/unit_test/tool/autofl_skill_job_importer_test.py new file mode 100644 index 0000000000..1c39a5b51f --- /dev/null +++ b/tests/unit_test/tool/autofl_skill_job_importer_test.py @@ -0,0 +1,584 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib.util +import sys +from pathlib import Path + +import pytest +import yaml + + +def _load_importer(): + repo_root = Path(__file__).parents[3] + importer_path = repo_root / "skills" / "nvflare-autofl" / "scripts" / "job_importer.py" + spec = importlib.util.spec_from_file_location("nvflare_autofl_skill_job_importer", importer_path) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +job_importer = _load_importer() +AUTOFL_CONFIG_SCHEMA_VERSION = job_importer.AUTOFL_CONFIG_SCHEMA_VERSION +DeterministicJobImporter = job_importer.DeterministicJobImporter +dump_autofl_yaml = job_importer.dump_autofl_yaml +import_job_to_autofl_config = job_importer.import_job_to_autofl_config + + +def _objective(metric, source="user_request"): + return { + "metric": metric, + "requested_metric": metric, + "optimization_metric": metric, + "metric_extraction_order": [metric], + "mode": "max", + "metric_contract_source": source, + } + + +def _write_recipe_job(root): + (root / "model.py").write_text( + """ +class SimpleNetwork: + pass +""", + encoding="utf-8", + ) + (root / "client.py").write_text( + """ +import argparse + + +def build_parser(): + parser = argparse.ArgumentParser() + parser.add_argument("--lr", type=float, default=0.01) + parser.add_argument("--batch_size", type=int, default=64) + parser.add_argument("--weight_decay", type=float, default=0.001) + return parser +""", + encoding="utf-8", + ) + (root / "job.py").write_text( + """ +import argparse + +from model import SimpleNetwork +from nvflare.app_opt.pt.recipes.fedavg import FedAvgRecipe +from nvflare.recipe import SimEnv + + +def define_parser(): + parser = argparse.ArgumentParser() + parser.add_argument("--n_clients", type=int, default=3) + parser.add_argument("--num_rounds", type=int, default=5) + parser.add_argument("--train_script", type=str, default="client.py") + parser.add_argument("--key_metric", type=str, default="accuracy") + return parser.parse_args() + + +def main(): + args = define_parser() + recipe = FedAvgRecipe( + name="demo", + min_clients=args.n_clients, + num_rounds=args.num_rounds, + model=SimpleNetwork(), + train_script=args.train_script, + key_metric=args.key_metric, + ) + env = SimEnv(num_clients=args.n_clients) + recipe.execute(env) + + +if __name__ == "__main__": + main() +""", + encoding="utf-8", + ) + return root / "job.py" + + +def test_import_recipe_job_extracts_trust_contract_without_executing_code(tmp_path): + job_path = _write_recipe_job(tmp_path) + + config = import_job_to_autofl_config( + str(job_path), + workspace_root=str(tmp_path), + metric="AUC", + target_env="prod", + max_candidates=8, + ) + + assert config["schema_version"] == AUTOFL_CONFIG_SCHEMA_VERSION + assert config["import"]["support"]["patterns"] == ["recipe:FedAvgRecipe", "env:SimEnv"] + assert config["import"]["confidence"] == "high" + assert config["job"]["surface"] == "recipe" + assert config["job"]["recipe"] == "FedAvgRecipe" + assert config["job"]["train_script"] == "client.py" + assert config["objective"] == _objective("AUC") + assert config["budget"]["max_candidates"] == 8 + assert config["budget"]["fixed_training_budget"] == { + "num_rounds": 5, + "min_clients": 3, + "num_clients": 3, + } + assert config["environment"]["requested"] == "prod" + assert config["environment"]["profiles"]["sim"] == {"num_clients": 3} + assert config["search_space"]["suggested"]["lr"]["default"] == 0.01 + assert config["search_space"]["suggested"]["batch_size"]["type"] == "int" + assert config["trust_contract"]["allowed_edit_paths"] == ["job.py", "client.py", "model.py"] + assert config["trust_contract"]["allowed_create_patterns"] == ["**/*.py"] + assert "allowed_edit_paths" not in config["job"] + assert "allowed_create_patterns" not in config["job"] + assert config["trust_contract"]["agent_controls"]["must_not_edit_outside_allowed_paths"] is True + assert config["unresolved"] == [] + + +def test_import_is_repeatable_and_yaml_round_trips(tmp_path): + job_path = _write_recipe_job(tmp_path) + importer = DeterministicJobImporter(workspace_root=str(tmp_path)) + + first = importer.import_job(str(job_path), max_candidates=4) + second = importer.import_job(str(job_path), max_candidates=4) + yaml_text = dump_autofl_yaml(first) + + assert first == second + assert yaml.safe_load(yaml_text) == first + assert "&id" not in yaml_text + assert first["trust_contract"]["unresolved"] is not first["unresolved"] + + +def test_import_marks_dynamic_argparse_defaults_unresolved(tmp_path): + (tmp_path / "client.py").write_text( + """ +import argparse + + +def build_parser(): + parser = argparse.ArgumentParser() + parser.add_argument("--model_arch", type=str, default=DEFAULT_MODEL_ARCH) + return parser +""", + encoding="utf-8", + ) + job_path = tmp_path / "job.py" + job_path.write_text( + """ +import argparse + +from nvflare.app_opt.pt.recipes.fedavg import FedAvgRecipe +from nvflare.recipe import SimEnv + + +def define_parser(): + parser = argparse.ArgumentParser() + parser.add_argument("--n_clients", type=int, default=2) + parser.add_argument("--num_rounds", type=int, default=3) + parser.add_argument("--train_script", type=str, default="client.py") + return parser.parse_args() + + +def main(): + args = define_parser() + recipe = FedAvgRecipe( + name="demo", + min_clients=args.n_clients, + num_rounds=args.num_rounds, + train_script=args.train_script, + ) + recipe.execute(SimEnv(num_clients=args.n_clients)) +""", + encoding="utf-8", + ) + + config = import_job_to_autofl_config(str(job_path), workspace_root=str(tmp_path)) + + model_arch = config["search_space"]["suggested"]["model_arch"] + assert model_arch["default"] == "DEFAULT_MODEL_ARCH" + assert model_arch["confidence"] == "low" + assert model_arch["unresolved"] is True + assert config["import"]["confidence"] == "medium" + assert { + "field": "search_space.suggested.model_arch.default", + "reason": "default is dynamic expression: DEFAULT_MODEL_ARCH", + } in config["unresolved"] + + +def test_import_marks_dynamic_train_script_unresolved_without_client_fallback(tmp_path): + (tmp_path / "client.py").write_text( + """ +import argparse + + +def build_parser(): + parser = argparse.ArgumentParser() + parser.add_argument("--lr", type=float, default=0.01) + return parser +""", + encoding="utf-8", + ) + job_path = tmp_path / "job.py" + job_path.write_text( + """ +from nvflare.app_opt.pt.recipes.fedavg import FedAvgRecipe +from nvflare.recipe import SimEnv + + +def get_script(): + return "client.py" + + +def main(): + recipe = FedAvgRecipe( + name="demo", + min_clients=2, + num_rounds=3, + train_script=get_script(), + ) + recipe.execute(SimEnv(num_clients=2)) +""", + encoding="utf-8", + ) + + config = import_job_to_autofl_config(str(job_path), workspace_root=str(tmp_path)) + + assert "train_script" not in config["job"] + assert "client.py" not in config["trust_contract"]["allowed_edit_paths"] + assert {"field": "job.train_script", "reason": "no train_script was found or resolved"} in config["unresolved"] + + +def test_import_marks_imported_budget_and_metric_constants_unresolved(tmp_path): + (tmp_path / "client.py").write_text( + """ +def train(): + pass +""", + encoding="utf-8", + ) + job_path = tmp_path / "job.py" + job_path.write_text( + """ +from config import KEY_METRIC, NUM_ROUNDS +from nvflare.app_opt.pt.recipes.fedavg import FedAvgRecipe +from nvflare.recipe import SimEnv + + +def main(): + recipe = FedAvgRecipe( + name="demo", + min_clients=2, + num_rounds=NUM_ROUNDS, + train_script="client.py", + key_metric=KEY_METRIC, + ) + recipe.execute(SimEnv(num_clients=2)) +""", + encoding="utf-8", + ) + + config = import_job_to_autofl_config(str(job_path), workspace_root=str(tmp_path)) + + assert config["objective"] == _objective("accuracy", source="default") + assert config["budget"]["fixed_training_budget"] == {"min_clients": 2, "num_clients": 2} + assert { + "field": "budget.fixed_training_budget.num_rounds", + "reason": "name:NUM_ROUNDS", + } in config["unresolved"] + assert {"field": "objective.metric", "reason": "name:KEY_METRIC"} in config["unresolved"] + assert {"field": "job.FedAvgRecipe.key_metric", "reason": "name:KEY_METRIC"} in config["unresolved"] + assert {"field": "job.FedAvgRecipe.num_rounds", "reason": "name:NUM_ROUNDS"} in config["unresolved"] + + +def test_import_marks_call_expression_budget_and_metric_unresolved(tmp_path): + (tmp_path / "client.py").write_text( + """ +def train(): + pass +""", + encoding="utf-8", + ) + job_path = tmp_path / "job.py" + job_path.write_text( + """ +from nvflare.app_opt.pt.recipes.fedavg import FedAvgRecipe +from nvflare.recipe import SimEnv + + +def get_metric(): + return "accuracy" + + +def get_rounds(): + return 5 + + +def main(): + recipe = FedAvgRecipe( + name="demo", + min_clients=2, + num_rounds=get_rounds(), + train_script="client.py", + key_metric=get_metric(), + ) + recipe.execute(SimEnv(num_clients=2)) +""", + encoding="utf-8", + ) + + config = import_job_to_autofl_config(str(job_path), workspace_root=str(tmp_path)) + + assert config["objective"] == _objective("accuracy", source="default") + assert config["budget"]["fixed_training_budget"] == {"min_clients": 2, "num_clients": 2} + assert { + "field": "budget.fixed_training_budget.num_rounds", + "reason": "call:get_rounds", + } in config["unresolved"] + assert {"field": "objective.metric", "reason": "call:get_metric"} in config["unresolved"] + assert {"field": "job.FedAvgRecipe.key_metric", "reason": "call:get_metric"} in config["unresolved"] + assert {"field": "job.FedAvgRecipe.num_rounds", "reason": "call:get_rounds"} in config["unresolved"] + assert config["job"]["recipe_args"]["num_rounds"] == { + "value": "get_rounds()", + "source": "call:get_rounds", + "confidence": "low", + } + + +def test_import_marks_unsupported_custom_job_as_partial(tmp_path): + job_path = tmp_path / "job.py" + job_path.write_text( + """ +def main(): + run_custom_workflow() + + +if __name__ == "__main__": + main() +""", + encoding="utf-8", + ) + + config = import_job_to_autofl_config(str(job_path), workspace_root=str(tmp_path)) + + assert config["import"]["support"]["status"] == "partial" + assert config["job"]["surface"] == "unknown" + unresolved_fields = {item["field"] for item in config["unresolved"]} + assert "job.surface" in unresolved_fields + assert "job.train_script" in unresolved_fields + assert "budget.fixed_training_budget" in unresolved_fields + + +def test_import_resolves_nvflare_fed_job_alias_and_script_runner(tmp_path): + (tmp_path / "train.py").write_text("print('train')\n", encoding="utf-8") + job_path = tmp_path / "job.py" + job_path.write_text( + """ +from nvflare.app_common.workflows.fedavg import FedAvgJob as ImportedJob +from nvflare.app_common.executors.script_runner import ScriptRunner as Runner + + +def main(): + job = ImportedJob(name="fedavg-alias", n_clients=8, min_clients=4, num_rounds=10, key_metric="AUC") + runner = Runner(script="train.py") + return job, runner +""".lstrip(), + encoding="utf-8", + ) + + config = import_job_to_autofl_config(str(job_path), workspace_root=str(tmp_path)) + + assert config["import"]["support"]["status"] == "supported" + assert config["job"]["fed_job"] == "FedAvgJob" + assert config["job"]["fed_job_class"] == "nvflare.app_common.workflows.fedavg.FedAvgJob" + assert config["job"]["train_script"] == "train.py" + assert config["objective"] == _objective("AUC", source="literal") + assert config["budget"]["fixed_training_budget"] == { + "num_rounds": 10, + "min_clients": 4, + "num_clients": 8, + } + + +def test_import_resolves_module_aliases_for_nvflare_job_subclasses(tmp_path): + (tmp_path / "train.py").write_text("print('train')\n", encoding="utf-8") + job_path = tmp_path / "job.py" + job_path.write_text( + """ +import nvflare.app_common.workflows as workflows +import nvflare.app_common.executors.script_runner as runner_module + + +def main(): + job = workflows.CCWFJob(name="ccwf", n_clients=3, min_clients=2, num_rounds=4) + runner = runner_module.ScriptRunner(script="train.py") + return job, runner +""".lstrip(), + encoding="utf-8", + ) + + config = import_job_to_autofl_config(str(job_path), workspace_root=str(tmp_path)) + + assert config["job"]["fed_job"] == "CCWFJob" + assert config["job"]["fed_job_class"] == "nvflare.app_common.workflows.CCWFJob" + assert config["job"]["train_script"] == "train.py" + assert config["budget"]["fixed_training_budget"]["num_clients"] == 3 + + +def test_import_recognizes_future_nvflare_job_subclasses_but_not_generic_or_local_jobs(tmp_path): + (tmp_path / "train.py").write_text("print('train')\n", encoding="utf-8") + stats_job = tmp_path / "stats_job.py" + stats_job.write_text( + """ +from nvflare.app_common.workflows import StatsJob +from nvflare.app_common.executors.script_runner import ScriptRunner + +StatsJob(name="stats", n_clients=2, min_clients=2, num_rounds=1) +ScriptRunner(script="train.py") +""".lstrip(), + encoding="utf-8", + ) + assert import_job_to_autofl_config(str(stats_job), workspace_root=str(tmp_path))["job"]["fed_job"] == "StatsJob" + + local_job = tmp_path / "local_job.py" + local_job.write_text( + """ +class CustomJob: + pass + +CustomJob() +""".lstrip(), + encoding="utf-8", + ) + local_config = import_job_to_autofl_config(str(local_job), workspace_root=str(tmp_path)) + assert local_config["import"]["support"]["status"] == "partial" + assert "local or non-NVFlare Job subclass" in next( + item["reason"] for item in local_config["unresolved"] if item["field"] == "job.surface" + ) + + generic_job = tmp_path / "generic_job.py" + generic_job.write_text("from nvflare.apis.job_def import Job\nJob()\n", encoding="utf-8") + generic_config = import_job_to_autofl_config(str(generic_job), workspace_root=str(tmp_path)) + assert generic_config["import"]["support"]["status"] == "partial" + + +def test_import_leaves_multiple_script_runners_unresolved(tmp_path): + for name in ("train_a.py", "train_b.py"): + tmp_path.joinpath(name).write_text("print('train')\n", encoding="utf-8") + job_path = tmp_path / "job.py" + job_path.write_text( + """ +from nvflare.app_common.workflows import EdgeJob +from nvflare.app_common.executors.script_runner import ScriptRunner + +EdgeJob(name="edge", n_clients=2, num_rounds=1) +ScriptRunner(script="train_a.py") +ScriptRunner(script="train_b.py") +""".lstrip(), + encoding="utf-8", + ) + + config = import_job_to_autofl_config(str(job_path), workspace_root=str(tmp_path)) + + assert "train_script" not in config["job"] + assert {"field": "job.train_script", "reason": "no train_script was found or resolved"} in config["unresolved"] + + +def test_import_keeps_async_function_assignments_out_of_module_scope(tmp_path): + job_path = tmp_path / "job.py" + job_path.write_text( + """ +from nvflare.app_opt.pt.recipes.fedavg import FedAvgRecipe + +NUM_ROUNDS = 3 + +async def helper(): + NUM_ROUNDS = 99 + return NUM_ROUNDS + +FedAvgRecipe(name="demo", model=object(), num_rounds=NUM_ROUNDS, min_clients=2) +""".lstrip(), + encoding="utf-8", + ) + + config = import_job_to_autofl_config(str(job_path), workspace_root=str(tmp_path)) + + assert config["budget"]["fixed_training_budget"]["num_rounds"] == 3 + + +def test_import_marks_augmented_budget_assignment_unresolved(tmp_path): + job_path = tmp_path / "job.py" + job_path.write_text( + """ +from nvflare.app_opt.pt.recipes.fedavg import FedAvgRecipe + +NUM_ROUNDS = 3 +NUM_ROUNDS += 97 +FedAvgRecipe(name="demo", model=object(), num_rounds=NUM_ROUNDS, min_clients=2) +""".lstrip(), + encoding="utf-8", + ) + + config = import_job_to_autofl_config(str(job_path), workspace_root=str(tmp_path)) + + assert "num_rounds" not in config["budget"]["fixed_training_budget"] + assert any(item["field"] == "budget.fixed_training_budget.num_rounds" for item in config["unresolved"]) + + +def test_import_ignores_add_argument_on_non_argparse_objects(tmp_path): + job_path = tmp_path / "job.py" + job_path.write_text( + """ +import argparse +from nvflare.app_opt.pt.recipes.fedavg import FedAvgRecipe + +class Registry: + def add_argument(self, *args, **kwargs): + pass + +registry = Registry() +registry.add_argument("--num_rounds", default=99) +parser = argparse.ArgumentParser() +parser.add_argument("--num_rounds", type=int, default=4) +args = parser.parse_args() +FedAvgRecipe(name="demo", model=object(), num_rounds=args.num_rounds, min_clients=2) +""".lstrip(), + encoding="utf-8", + ) + + config = import_job_to_autofl_config(str(job_path), workspace_root=str(tmp_path)) + + assert config["budget"]["fixed_training_budget"]["num_rounds"] == 4 + + +def test_import_does_not_admit_local_recipe_classes(tmp_path): + job_path = tmp_path / "job.py" + job_path.write_text( + """ +class MyRecipe: + pass + +MyRecipe(name="local", num_rounds=1, min_clients=2) +""".lstrip(), + encoding="utf-8", + ) + + config = import_job_to_autofl_config(str(job_path), workspace_root=str(tmp_path)) + + assert config["import"]["support"]["status"] == "partial" + assert config["job"]["surface"] == "unknown" + + +def test_import_returns_clean_error_for_missing_job(tmp_path): + with pytest.raises(job_importer.JobImportError, match="job.py not found"): + import_job_to_autofl_config(str(tmp_path / "missing.py"), workspace_root=str(tmp_path)) diff --git a/tests/unit_test/tool/autofl_skill_plot_progress_test.py b/tests/unit_test/tool/autofl_skill_plot_progress_test.py new file mode 100644 index 0000000000..eefaa47987 --- /dev/null +++ b/tests/unit_test/tool/autofl_skill_plot_progress_test.py @@ -0,0 +1,137 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import csv +import importlib.util +import sys +from pathlib import Path + +import pytest + + +def _load_plotter(): + repo_root = Path(__file__).parents[3] + script_path = repo_root / "skills" / "nvflare-autofl" / "scripts" / "plot_progress.py" + spec = importlib.util.spec_from_file_location("nvflare_autofl_skill_plot_progress", script_path) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def _record(plotter, index, score, status="discard", name=None, runtime=300.0): + return plotter.ProgressRecord( + index=index, + status=status, + name=name or f"candidate_{index}", + score=score, + runtime_seconds=runtime, + description=f"candidate {index}", + ) + + +def test_robust_y_limits_focus_on_improvement_region_for_maximize(): + plotter = _load_plotter() + scores = [0.50, 0.55, 0.687] + [0.70 + index * 0.002 for index in range(20)] + + lower, upper = plotter.default_y_limits(scores, baseline=0.687, mode="max") + + assert lower > 0.60 + assert lower < 0.687 + assert upper > max(scores) + + +def test_robust_y_limits_focus_on_improvement_region_for_minimize(): + plotter = _load_plotter() + scores = [2.0, 1.8, 0.90] + [0.80 - index * 0.01 for index in range(20)] + + lower, upper = plotter.default_y_limits(scores, baseline=0.90, mode="min") + + assert lower < min(scores) + assert upper > 0.90 + assert upper < 1.5 + + +@pytest.mark.parametrize( + "mode,scores,expected_best", + [ + ("max", [0.5, 0.6, 0.55, 0.7, 0.69, 0.8], 0.8), + ("min", [0.8, 0.7, 0.75, 0.6, 0.62, 0.5], 0.5), + ], +) +def test_milestone_selection_supports_both_objective_directions(mode, scores, expected_best): + plotter = _load_plotter() + records = [ + _record(plotter, index, score, status="baseline" if index == 0 else "keep") + for index, score in enumerate(scores) + ] + + milestones = plotter.select_observed_milestones(records, mode=mode, max_labels=3) + + assert len(milestones) <= 3 + assert milestones[-1][1].score == expected_best + + +def test_load_results_uses_productized_ledger_fields(tmp_path): + plotter = _load_plotter() + ledger = tmp_path / "results.tsv" + with ledger.open("w", encoding="utf-8", newline="") as f: + writer = csv.DictWriter( + f, + fieldnames=["status", "name", "score", "runtime_seconds", "diff_summary"], + delimiter="\t", + ) + writer.writeheader() + writer.writerow( + { + "status": "keep", + "name": "fedavgm", + "score": "0.75", + "runtime_seconds": "123.5", + "diff_summary": "server momentum", + } + ) + + records = plotter.load_results(ledger) + + assert records == [plotter.ProgressRecord(0, "keep", "fedavgm", 0.75, 123.5, "server momentum")] + assert plotter.normalize_records(records) == records + + +def test_zero_score_label_uses_the_actual_score(): + plotter = _load_plotter() + record = _record(plotter, 1, 0.0) + + offset, _, vertical_alignment = plotter.label_placement(0, record, (0.0, 2.0), (-1.0, 0.1)) + + assert offset[1] < 0 + assert vertical_alignment == "top" + + +def test_rich_progress_plot_renders_png(tmp_path): + pytest.importorskip("matplotlib") + plotter = _load_plotter() + records = [_record(plotter, 0, 0.687, status="baseline", name="baseline")] + for index in range(1, 21): + status = "keep" if index in {3, 8, 15} else "discard" + records.append(_record(plotter, index, 0.69 + index * 0.002, status=status)) + records.insert(10, _record(plotter, 10, None, status="literature", name="literature_review_1", runtime=45.0)) + output = tmp_path / "progress.png" + + baseline, best = plotter.plot_progress(records, output, mode="max", metric_label="test_accuracy") + + assert baseline == pytest.approx(0.687) + assert best == pytest.approx(0.73) + assert output.read_bytes().startswith(b"\x89PNG\r\n\x1a\n") + assert output.stat().st_size > 20_000 diff --git a/tests/unit_test/tool/autofl_skill_runner_test.py b/tests/unit_test/tool/autofl_skill_runner_test.py new file mode 100644 index 0000000000..2f878f6f43 --- /dev/null +++ b/tests/unit_test/tool/autofl_skill_runner_test.py @@ -0,0 +1,1523 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib.util +import json +import os +import subprocess +import sys +from copy import deepcopy +from pathlib import Path +from types import SimpleNamespace + +import pytest + + +def _load_runner(): + repo_root = Path(__file__).parents[3] + runner_path = repo_root / "skills" / "nvflare-autofl" / "scripts" / "run_job_campaign.py" + spec = importlib.util.spec_from_file_location("nvflare_autofl_skill_runner", runner_path) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def _campaign_config(): + return { + "schema_version": "nvflare.autofl.config.v1", + "import": {"source_sha256": "a" * 64, "support": {"status": "supported"}}, + "job": {}, + "objective": { + "metric": "accuracy", + "requested_metric": "accuracy", + "optimization_metric": "accuracy", + "metric_extraction_order": ["accuracy"], + }, + "budget": {"fixed_training_budget": {"num_clients": 2, "num_rounds": 1}}, + "environment": {"requested": "sim"}, + "search_space": {"suggested": {"lr": {"default": 0.1}}}, + "trust_contract": { + "allowed_edit_paths": ["job.py", "client.py"], + "allowed_create_patterns": ["**/*.py"], + "unresolved": [], + }, + } + + +def _initialize_fake_campaign(runner, tmp_path, monkeypatch, *, target_env="sim", baseline_score=0.5): + job = tmp_path / "job.py" + client = tmp_path / "client.py" + job.write_text("print('job')\n", encoding="utf-8") + client.write_text("ALGORITHM = 'baseline'\n", encoding="utf-8") + config = _campaign_config() + config["environment"]["requested"] = target_env + monkeypatch.setattr(runner, "import_job_config", lambda *args, **kwargs: deepcopy(config)) + monkeypatch.setattr(runner, "job_help", lambda *args, **kwargs: "") + monkeypatch.setattr(runner, "write_progress", lambda path, *args: path.write_bytes(b"progress")) + + def fake_run(run_def, **kwargs): + return runner.RunRecord( + run_def.status, + run_def.name, + baseline_score, + 1.0, + "none", + run_def.description, + "python job.py", + str(tmp_path / "artifacts" / run_def.name), + ) + + monkeypatch.setattr(runner, "run_job", fake_run) + argv = ["initialize", str(job), "--env", target_env, "--no-prefer-synthetic"] + assert runner.main(argv) == 0 + return job, client, config + + +def test_candidate_args_respect_mutation_schema_bounds(): + runner = _load_runner() + schema = {"mutable_args": {"lr": {"type": "float", "min": 0.0001, "max": 0.1}}} + + assert runner.candidate_args_allowed(["--lr", "0.1"], schema) == (True, "") + + allowed, reason = runner.candidate_args_allowed(["--lr", "0.2"], schema) + assert not allowed + assert "above schema max" in reason + + +def test_candidate_run_args_cannot_override_fixed_budget(): + runner = _load_runner() + config = {"budget": {"fixed_training_budget": {"num_clients": 8, "num_rounds": 20}}} + + assert runner.candidate_preserves_fixed_args(["--lr", "0.01"], config, {}) == (True, "") + assert runner.candidate_preserves_fixed_args(["--num_rounds=2"], config, {})[0] is False + assert runner.candidate_preserves_fixed_args(["--n-clients", "4"], config, {})[0] is False + assert runner.candidate_preserves_fixed_args(["--num_round", "2"], config, {})[0] is False + with pytest.raises(ValueError, match="canonical long options"): + runner.candidate_preserves_fixed_args(["-r", "2"], config, {}) + assert runner.candidate_preserves_fixed_args(["--temperature", "-1.5"], config, {}) == (True, "") + assert runner.candidate_preserves_fixed_args(["--new-algorithm", "fednova"], config, {}) == (True, "") + + +def test_supported_flags_are_exact_help_tokens(): + runner = _load_runner() + help_text = "usage: job.py [--num_rounds NUM_ROUNDS] [--name NAME]" + + assert runner.supports_flag(help_text, "--num_rounds") + assert runner.supports_flag(help_text, "--name") + assert not runner.supports_flag(help_text, "--num_round") + assert not runner.supports_flag(help_text, "--round") + + +def test_candidate_plan_skips_out_of_bounds_learning_rate(): + runner = _load_runner() + config = {"search_space": {"suggested": {"lr": {"default": 0.05}}}} + schema = {"mutable_args": {"lr": {"type": "float", "min": 0.0001, "max": 0.1}}} + help_text = "--lr" + + candidates = list(runner.candidate_plan(config, help_text, max_candidates=10, schema=schema)) + candidate_commands = [" ".join(candidate.args) for candidate in candidates] + + assert "--lr 0.1" in candidate_commands + assert "--lr 0.2" not in candidate_commands + + +def test_initialize_merges_existing_mutation_schema_preferred_targets_into_autofl(tmp_path, monkeypatch): + runner = _load_runner() + job = tmp_path / "job.py" + job.write_text("print('job')\n", encoding="utf-8") + tmp_path.joinpath("client.py").write_text("print('client')\n", encoding="utf-8") + tmp_path.joinpath("custom_aggregators.py").write_text("class CustomAggregator:\n pass\n", encoding="utf-8") + tmp_path.joinpath("mutation_schema.yaml").write_text( + "preferred_targets:\n - client.py\n - custom_aggregators.py\n", + encoding="utf-8", + ) + monkeypatch.setattr(runner, "import_job_config", lambda *args, **kwargs: deepcopy(_campaign_config())) + monkeypatch.setattr(runner, "job_help", lambda *args, **kwargs: "") + monkeypatch.setattr(runner, "write_progress", lambda path, *args: path.write_bytes(b"progress")) + monkeypatch.setattr( + runner, + "run_job", + lambda run_def, **kwargs: runner.RunRecord( + run_def.status, + run_def.name, + 0.5, + 1.0, + "none", + run_def.description, + "python job.py", + "/tmp/baseline", + ), + ) + + assert runner.main(["initialize", str(job), "--no-prefer-synthetic"]) == 0 + + config = runner.read_yaml(tmp_path / "autofl.yaml") + assert "custom_aggregators.py" in config["trust_contract"]["allowed_edit_paths"] + assert config["trust_contract"]["preferred_targets"] == ["client.py", "custom_aggregators.py"] + + assert ( + runner.main( + [ + "prepare", + str(job), + "--name", + "new_server_aggregator", + "--hypothesis", + "implement source-backed server aggregation", + ] + ) + == 0 + ) + draft = tmp_path / ".nvflare/autofl/candidates/new_server_aggregator/source/custom_aggregators.py" + draft.write_text("class CustomAggregator:\n improved = True\n", encoding="utf-8") + monkeypatch.setattr( + runner, + "run_job", + lambda run_def, **kwargs: runner.RunRecord( + "candidate", + run_def.name, + 0.8, + 1.0, + "none", + run_def.description, + "python job.py", + "/tmp/candidate", + ), + ) + + assert runner.main(["evaluate", str(job)]) == 0 + assert "improved = True" in tmp_path.joinpath("custom_aggregators.py").read_text(encoding="utf-8") + + +def test_missing_or_escaping_preferred_targets_remain_unresolved(tmp_path): + runner = _load_runner() + config = deepcopy(_campaign_config()) + schema = {"preferred_targets": ["missing_aggregator.py", "../shared/custom_aggregators.py"]} + + updated = runner.apply_mutation_schema_contract(config, schema, tmp_path) + + assert "missing_aggregator.py" not in updated["trust_contract"]["allowed_edit_paths"] + assert "../shared/custom_aggregators.py" not in updated["trust_contract"]["allowed_edit_paths"] + unresolved_reasons = [ + item["reason"] for item in updated["unresolved"] if item["field"] == "mutation_schema.preferred_targets" + ] + assert any(reason.startswith("missing_aggregator.py:") for reason in unresolved_reasons) + assert any(reason.startswith("../shared/custom_aggregators.py:") for reason in unresolved_reasons) + + +def test_runner_prefers_explicit_test_accuracy_alias(tmp_path): + runner = _load_runner() + result_path = tmp_path / "cross_val_results.json" + result_path.write_text( + json.dumps( + { + "site-1": { + "SRV_FL_global_model.pt": { + "accuracy": 0.5, + "test_accuracy": 0.8, + } + } + } + ), + encoding="utf-8", + ) + + assert runner.extract_score(tmp_path, ["test_accuracy", "accuracy"]) == 0.8 + + +def test_metric_order_precedes_artifact_file_order(tmp_path): + runner = _load_runner() + tmp_path.joinpath("metrics_summary.json").write_text(json.dumps({"accuracy": 0.7}), encoding="utf-8") + tmp_path.joinpath("cross_val_results.json").write_text(json.dumps({"test_accuracy": 0.8}), encoding="utf-8") + + assert runner.extract_score(tmp_path, ["test_accuracy", "accuracy"]) == 0.8 + + +def test_structured_metric_extraction_rejects_boolean_values(): + runner = _load_runner() + + assert runner.find_metric_value({"accuracy": True}, ["accuracy"]) is None + assert runner.find_metric_value({"metrics": [{"name": "accuracy", "value": False}]}, ["accuracy"]) is None + + +def test_text_metric_extraction_supports_scientific_notation(tmp_path): + runner = _load_runner() + tmp_path.joinpath("log_fl.txt").write_text("validation_loss: 1.25e-4\n", encoding="utf-8") + + assert runner.extract_score(tmp_path, ["validation_loss"]) == pytest.approx(0.000125) + + +def test_text_metric_extraction_requires_exact_metric_token_and_records_provenance(tmp_path): + runner = _load_runner() + log_path = tmp_path / "log.txt" + log_path.write_text("val_accuracy: 0.9\naccuracy: 0.7\n", encoding="utf-8") + + evidence = runner.extract_metric_evidence(tmp_path, ["accuracy"]) + + assert evidence.score == 0.7 + assert evidence.metric_name == "accuracy" + assert evidence.source == "text:log.txt" + assert evidence.artifact == str(log_path.resolve()) + + +def test_runner_applies_schema_metric_contract(): + runner = _load_runner() + config = { + "objective": { + "metric": "accuracy", + "requested_metric": "accuracy", + "optimization_metric": "accuracy", + "metric_extraction_order": ["accuracy"], + } + } + schema = { + "objective": { + "requested_metric": "accuracy", + "optimization_metric": "test_accuracy", + "metric_extraction_order": ["test_accuracy", "accuracy"], + "metric_source": "held-out CIFAR-10 test set", + } + } + + updated = runner.apply_metric_contract(config, "accuracy", schema) + + assert updated["objective"]["metric"] == "accuracy" + assert updated["objective"]["requested_metric"] == "accuracy" + assert updated["objective"]["optimization_metric"] == "test_accuracy" + assert updated["objective"]["metric_extraction_order"] == ["test_accuracy", "accuracy"] + assert updated["objective"]["metric_source"] == "held-out CIFAR-10 test set" + + +def test_comparison_budget_suppresses_duplicate_imported_fixed_budget_args(): + runner = _load_runner() + config = {"budget": {"fixed_training_budget": {"num_clients": 8, "num_rounds": 10}}} + schema = { + "comparison_budget_args": { + "default_candidate_budget": { + "n_clients": 8, + "num_rounds": 20, + } + } + } + help_text = "--n_clients --num_rounds" + + assert runner.build_fixed_args(config, help_text, schema) == [] + + args = SimpleNamespace(base_args="", prefer_synthetic=False, synthetic_train_size=1, synthetic_test_size=1) + assert runner.build_base_args(args, help_text, schema) == ["--n_clients", "8", "--num_rounds", "20"] + + +@pytest.mark.parametrize( + "field,value", + [ + ("run_timeout_seconds", [120]), + ("run_timeout_seconds", {"seconds": 120}), + ("simulator_no_progress_timeout_seconds", [120]), + ("simulator_no_progress_timeout_seconds", {"seconds": 120}), + ("run_timeout_seconds", -1), + ("run_timeout_seconds", True), + ("run_timeout_seconds", 1.5), + ], +) +def test_campaign_timeout_rejects_malformed_schema_values(field, value): + runner = _load_runner() + args = SimpleNamespace(timeout=900, simulator_no_progress_timeout=240) + schema = {"comparison_budget_args": {"default_candidate_budget": {field: value}}} + + with pytest.raises(ValueError, match=f"{field} must be a non-negative integer"): + runner.campaign_timeout(args, schema) + + +def test_run_streams_output_before_timeout(tmp_path): + runner = _load_runner() + log_path = tmp_path / "run.log" + + rc, output, _ = runner.run( + [ + sys.executable, + "-c", + "import time; print('started', flush=True); time.sleep(2); print('finished', flush=True)", + ], + tmp_path, + timeout=1, + log_path=log_path, + ) + + log_text = log_path.read_text(encoding="utf-8") + assert rc == 124 + assert "started" in output + assert "started" in log_text + assert "TIMEOUT after 1s" in log_text + + +def test_run_streams_partial_line_before_timeout(tmp_path): + runner = _load_runner() + log_path = tmp_path / "run.log" + + rc, output, runtime = runner.run( + [ + sys.executable, + "-c", + "import sys, time; sys.stdout.write('partial output'); sys.stdout.flush(); time.sleep(30)", + ], + tmp_path, + timeout=1, + log_path=log_path, + ) + + assert rc == 124 + assert runtime < 5 + assert "partial output" in output + assert "partial output" in log_path.read_text(encoding="utf-8") + + +def test_job_help_is_bounded(tmp_path): + runner = _load_runner() + job = tmp_path / "job.py" + job.write_text("import time\ntime.sleep(30)\n", encoding="utf-8") + + with pytest.raises(RuntimeError, match="job.py --help exceeded 1 seconds"): + runner.job_help(sys.executable, job, tmp_path, timeout=1) + + +def test_metric_extraction_prefers_structured_artifacts_then_known_text_logs(tmp_path): + runner = _load_runner() + nested = tmp_path / "server" + nested.mkdir() + nested.joinpath("metrics_summary.json").write_text(json.dumps({"accuracy": 0.7}), encoding="utf-8") + nested.joinpath("cross_val_results.json").write_text(json.dumps({"accuracy": 0.8}), encoding="utf-8") + nested.joinpath("log_fl.txt").write_text("accuracy: 0.9\n", encoding="utf-8") + tmp_path.joinpath("run.log").write_text("accuracy: 0.95\n", encoding="utf-8") + + assert runner.extract_score(tmp_path, ["accuracy"]) == 0.7 + nested.joinpath("metrics_summary.json").unlink() + assert runner.extract_score(tmp_path, ["accuracy"]) == 0.8 + nested.joinpath("cross_val_results.json").unlink() + assert runner.extract_score(tmp_path, ["accuracy"]) == 0.95 + + +def test_result_paths_and_static_job_names_are_resolved_deterministically(tmp_path): + runner = _load_runner() + relative = tmp_path / "relative-result" + relative.mkdir() + config = { + "job": { + "recipe_args": {"name": {"value": "recipe-name", "confidence": "high"}}, + "fed_job_args": {"name": {"value": "fed-job-name", "confidence": "high"}}, + } + } + + assert runner.extract_result_dir("Result can be found in : relative-result", tmp_path) == relative.resolve() + assert runner.imported_job_names(config) == ["recipe-name", "fed-job-name"] + assert runner.expected_simulator_roots(config, "injected", tmp_path) == [ + Path("/tmp/nvflare/simulation/injected").resolve(), + Path("/tmp/nvflare/simulation/recipe-name").resolve(), + Path("/tmp/nvflare/simulation/fed-job-name").resolve(), + ] + config["job"]["recipe_args"]["name"]["confidence"] = "low" + assert runner.imported_job_names(config) == ["fed-job-name"] + with pytest.raises(ValueError, match="unsafe simulator job name"): + runner.expected_simulator_roots({}, "../other-run", tmp_path) + + +def test_success_without_deterministic_result_root_is_actionable_failure(tmp_path, monkeypatch): + runner = _load_runner() + job = tmp_path / "job.py" + job.write_text("print('done')\n", encoding="utf-8") + monkeypatch.setattr(runner, "run_allow_timeout", lambda *args, **kwargs: (0, "done\n", 0.1)) + + record = runner.run_job( + runner.JobRun("candidate", [], "candidate"), + python=sys.executable, + job=job, + cwd=tmp_path, + help_text="", + fixed_args=[], + base_args=[], + output_root=tmp_path / "runs", + timeout=10, + simulator_no_progress_timeout=0, + metrics=["accuracy"], + config={"job": {}}, + ) + + assert record.status == "crash" + assert "no deterministic NVFlare result directory" in record.failure_reason + + +def test_run_job_collects_configured_sim_result_and_standard_nvflare_text_metric(tmp_path): + runner = _load_runner() + job = tmp_path / "job.py" + job.write_text( + f""" +import argparse +from pathlib import Path + +parser = argparse.ArgumentParser() +parser.add_argument("--name", default="run") +args = parser.parse_args() +result = Path({str(tmp_path / 'simulation')!r}) / args.name +server = result / "server" +server.mkdir(parents=True, exist_ok=True) +server.joinpath("log.txt").write_text("accuracy: 0.76\\n") +print(f"Result can be found in : {{result}}") +""".lstrip(), + encoding="utf-8", + ) + + record = runner.run_job( + runner.JobRun("candidate", [], "candidate"), + python=sys.executable, + job=job, + cwd=tmp_path, + help_text="--name NAME", + fixed_args=[], + base_args=[], + output_root=tmp_path / "runs", + timeout=10, + simulator_no_progress_timeout=0, + metrics=["accuracy"], + config={ + "job": {}, + "environment": { + "discovered": { + "name": "SimEnv", + "args": { + "workspace_root": { + "value": str(tmp_path / "simulation"), + "confidence": "high", + } + }, + } + }, + }, + ) + + assert record.status == "candidate" + assert record.score == pytest.approx(0.76) + assert record.metric_source == "text:log.txt" + assert tmp_path.joinpath("runs/candidate/simulation/server/log.txt").is_file() + assert tmp_path.joinpath("runs/candidate/run.log").is_file() + assert tmp_path.joinpath("simulation/autofl_candidate").is_dir() + + +def test_run_stops_on_nvflare_simulator_stall_log(tmp_path): + runner = _load_runner() + log_path = tmp_path / "run.log" + sim_root = tmp_path / "simulation" / "autofl_candidate" + server_log = sim_root / "server" / "log_fl.txt" + server_log.parent.mkdir(parents=True) + server_log.write_text( + "SimulatorClientRunner - ERROR - run_client_thread error: RuntimeError: " + "Failed to create connection to the child process in SimulatorClientRunner, timeout: 60.0\n", + encoding="utf-8", + ) + + rc, output, runtime = runner.run( + [ + sys.executable, + "-c", + "import time; print('started', flush=True); time.sleep(30)", + ], + tmp_path, + timeout=30, + log_path=log_path, + simulator_stall_roots=[sim_root], + stall_check_interval=0.01, + ) + + log_text = log_path.read_text(encoding="utf-8") + assert rc == runner.SIMULATOR_STALL_EXIT_CODE + assert runtime < 5 + assert "SIMULATOR_STALL:" in output + assert "SIMULATOR_STALL:" in log_text + + +def test_run_stops_on_nvflare_simulator_no_progress_log(tmp_path): + runner = _load_runner() + log_path = tmp_path / "run.log" + sim_root = tmp_path / "simulation" / "autofl_candidate" + server_log = sim_root / "server" / "log_fl.txt" + server_log.parent.mkdir(parents=True) + server_log.write_text("Round 0 started\n", encoding="utf-8") + + rc, output, runtime = runner.run( + [ + sys.executable, + "-c", + "import time; print('started', flush=True); time.sleep(30)", + ], + tmp_path, + timeout=30, + log_path=log_path, + simulator_stall_roots=[sim_root], + stall_check_interval=0.01, + simulator_no_progress_timeout=1, + ) + + log_text = log_path.read_text(encoding="utf-8") + assert rc == runner.SIMULATOR_STALL_EXIT_CODE + assert runtime < 5 + assert "SIMULATOR_STALL: no simulator progress markers changed" in output + assert "SIMULATOR_STALL: no simulator progress markers changed" in log_text + + +def test_run_stops_on_stale_partial_simulator_aggregation(tmp_path): + runner = _load_runner() + log_path = tmp_path / "run.log" + sim_root = tmp_path / "simulation" / "autofl_candidate" + server_log = sim_root / "server" / "log_fl.txt" + site_log = sim_root / "site-1" / "log_fl.txt" + server_log.parent.mkdir(parents=True) + site_log.parent.mkdir(parents=True) + server_log.write_text( + "Round 0 started\n" "2026-06-25 06:32:33 - FedAvg - INFO - Aggregated 1/8 results\n", + encoding="utf-8", + ) + site_log.write_text("[site=site-1] round=0\n", encoding="utf-8") + + rc, output, runtime = runner.run( + [ + sys.executable, + "-c", + ( + "import pathlib, time; " + f"path = pathlib.Path({str(site_log)!r}); " + "time.sleep(0.2); " + "path.write_text('[site=site-1] round=0\\n[site=site-2] round=0\\n'); " + "time.sleep(30)" + ), + ], + tmp_path, + timeout=30, + log_path=log_path, + simulator_stall_roots=[sim_root], + stall_check_interval=0.01, + simulator_no_progress_timeout=1, + ) + + log_text = log_path.read_text(encoding="utf-8") + assert rc == runner.SIMULATOR_STALL_EXIT_CODE + assert runtime < 5 + assert "SIMULATOR_STALL: partial simulator aggregation made no server-side progress" in output + assert "Aggregated 1/8 results" in log_text + + +def test_runner_state_routes_plateau_to_literature_checkpoint(tmp_path): + runner = _load_runner() + records = [ + runner.RunRecord("baseline", "baseline", 0.85, 1.0, "none", "baseline", "python job.py", "/tmp/baseline"), + runner.RunRecord("discard", "candidate_1", 0.84, 1.0, "none", "candidate", "python job.py", "/tmp/c1"), + runner.RunRecord("discard", "candidate_2", 0.84, 1.0, "none", "candidate", "python job.py", "/tmp/c2"), + ] + results_path = tmp_path / "results.tsv" + state_path = tmp_path / "state.json" + runner.write_results(results_path, records) + + state = runner.write_state( + state_path, + results_path, + records, + None, + plateau_threshold=2, + ) + + assert state["schema_version"] == "nvflare.autofl.campaign_state.v1" + assert state["reason"] == "plateau_literature" + assert state["next_action"] == "run_literature_loop" + assert state["final_response_allowed"] is False + assert state == json.loads(state_path.read_text(encoding="utf-8")) + + +def test_small_retained_improvement_does_not_reset_plateau_clock(tmp_path): + runner = _load_runner() + records = [ + runner.RunRecord("baseline", "baseline", 0.85, 1.0, "none", "baseline", "python job.py", "/tmp/baseline"), + runner.RunRecord("keep", "candidate_1", 0.8503, 1.0, "client.py", "candidate", "python job.py", "/tmp/c1"), + ] + results_path = tmp_path / "results.tsv" + runner.write_results(results_path, records) + + state = runner.write_state( + tmp_path / "state.json", + results_path, + records, + None, + plateau_threshold=1, + plateau_min_delta=0.0005, + ) + + assert runner.best_retained_record(records, "max").name == "candidate_1" + assert state["best_score"] == pytest.approx(0.8503) + assert state["plateau"]["best_score"] == pytest.approx(0.85) + assert state["reason"] == "plateau_literature" + assert state["next_action"] == "run_literature_loop" + + +def test_runner_state_finalizes_after_explicit_candidate_cap(tmp_path): + runner = _load_runner() + records = [ + runner.RunRecord("baseline", "baseline", 0.85, 1.0, "none", "baseline", "python job.py", "/tmp/baseline"), + runner.RunRecord("discard", "candidate_1", 0.84, 1.0, "none", "candidate", "python job.py", "/tmp/c1"), + ] + results_path = tmp_path / "results.tsv" + state_path = tmp_path / "state.json" + runner.write_results(results_path, records) + + state = runner.write_state(state_path, results_path, records, 1) + + assert state["decision"] == "stop" + assert state["reason"] == "candidate_cap_exhausted" + assert state["next_action"] == "final_report" + assert state["final_response_allowed"] is True + assert state["candidate_cap_source"] == "explicit" + + +def test_runner_state_ignores_ambient_candidate_cap(tmp_path, monkeypatch): + runner = _load_runner() + monkeypatch.setenv("AUTOFL_MAX_CANDIDATES", "1") + records = [ + runner.RunRecord("baseline", "baseline", 0.85, 1.0, "none", "baseline", "python job.py", "/tmp/baseline"), + runner.RunRecord("discard", "candidate_1", 0.84, 1.0, "none", "candidate", "python job.py", "/tmp/c1"), + ] + results_path = tmp_path / "results.tsv" + state_path = tmp_path / "state.json" + runner.write_results(results_path, records) + + state = runner.write_state(state_path, results_path, records, None) + + assert state["decision"] == "continue" + assert state["reason"] == "continue" + assert state["candidate_cap"] is None + assert state["candidate_cap_source"] == "uncapped" + assert state["final_response_allowed"] is False + + +def test_runner_state_marks_infrastructure_retry_non_final(tmp_path): + runner = _load_runner() + records = [ + runner.RunRecord( + runner.INFRASTRUCTURE_RETRY, + "baseline", + None, + 1.0, + "none", + "baseline", + "python job.py", + "/tmp/baseline", + ) + ] + results_path = tmp_path / "results.tsv" + state_path = tmp_path / "state.json" + runner.write_results(results_path, records) + + state = runner.write_state(state_path, results_path, records, None) + + assert state["decision"] == "retry_infrastructure" + assert state["reason"] == "infrastructure_retry" + assert state["next_action"] == "rerun_with_escalated_execution" + assert state["final_response_allowed"] is False + + +def test_successful_retry_clears_historical_infrastructure_decision(tmp_path): + runner = _load_runner() + records = [ + runner.RunRecord( + runner.INFRASTRUCTURE_RETRY, + "baseline", + None, + 1.0, + "none", + "baseline", + "python job.py", + "/tmp/failed", + ), + runner.RunRecord( + "baseline", + "baseline_retry_2", + 0.5, + 1.0, + "none", + "baseline", + "python job.py", + "/tmp/success", + ), + ] + results_path = tmp_path / "results.tsv" + runner.write_results(results_path, records) + + state = runner.write_state(tmp_path / "state.json", results_path, records, max_candidates=None) + + assert state["decision"] == "continue" + assert state["next_action"] == "propose_candidate" + + +def test_baseline_crash_is_not_counted_as_candidate_attempt(): + runner = _load_runner() + records = [ + runner.RunRecord("baseline", "baseline", None, 1.0, "none", "baseline", "python job.py", "/tmp/baseline") + ] + + assert runner.candidate_attempts(records) == 0 + assert runner.is_sandbox_socket_failure( + "PermissionError: [Errno 1] Operation not permitted in get_open_ports while calling s.bind(('', 0))" + ) + + +def test_code_candidate_keeps_improvement_and_restores_discard_without_git(tmp_path, monkeypatch): + runner = _load_runner() + job, client, _ = _initialize_fake_campaign(runner, tmp_path, monkeypatch) + + assert runner.main(["prepare", str(job), "--name", "new_algo", "--hypothesis", "add a new algorithm"]) == 0 + draft = tmp_path / ".nvflare" / "autofl" / "candidates" / "new_algo" / "source" + draft.joinpath("client.py").write_text("from new_algorithm import VALUE\n", encoding="utf-8") + draft.joinpath("new_algorithm.py").write_text("VALUE = 'improved'\n", encoding="utf-8") + + def improved_run(run_def, **kwargs): + return runner.RunRecord( + "candidate", run_def.name, 0.7, 2.0, "none", run_def.description, "python job.py", "/tmp/new_algo" + ) + + monkeypatch.setattr(runner, "run_job", improved_run) + assert runner.main(["evaluate", str(job)]) == 0 + assert client.read_text(encoding="utf-8") == "from new_algorithm import VALUE\n" + assert tmp_path.joinpath("new_algorithm.py").read_text(encoding="utf-8") == "VALUE = 'improved'\n" + + kept_manifest = json.loads( + tmp_path.joinpath(".nvflare/autofl/candidates/new_algo/candidate_manifest.json").read_text(encoding="utf-8") + ) + assert kept_manifest["status"] == "keep" + assert kept_manifest["changed_files"] == ["client.py", "new_algorithm.py"] + assert kept_manifest["patch_sha256"] + + assert runner.main(["prepare", str(job), "--name", "bad_algo", "--hypothesis", "try a regression"]) == 0 + bad_draft = tmp_path / ".nvflare" / "autofl" / "candidates" / "bad_algo" / "source" + bad_draft.joinpath("client.py").write_text("ALGORITHM = 'regression'\n", encoding="utf-8") + + def regressed_run(run_def, **kwargs): + return runner.RunRecord( + "candidate", run_def.name, 0.3, 2.0, "none", run_def.description, "python job.py", "/tmp/bad_algo" + ) + + monkeypatch.setattr(runner, "run_job", regressed_run) + assert runner.main(["evaluate", str(job)]) == 0 + assert client.read_text(encoding="utf-8") == "from new_algorithm import VALUE\n" + records = runner.load_results(tmp_path / "results.tsv") + assert [record.status for record in records] == ["baseline", "keep", "discard"] + assert records[1].changed_files == "client.py,new_algorithm.py" + assert records[1].candidate_manifest.endswith("candidate_manifest.json") + + +def test_candidate_rejects_unauthorized_existing_source_and_symlink(tmp_path, monkeypatch): + runner = _load_runner() + job, _, _ = _initialize_fake_campaign(runner, tmp_path, monkeypatch) + tmp_path.joinpath("secret.py").write_text("SECRET = True\n", encoding="utf-8") + assert runner.main(["prepare", str(job), "--name", "unsafe", "--hypothesis", "touch secret"]) == 0 + draft = tmp_path / ".nvflare" / "autofl" / "candidates" / "unsafe" / "source" + draft.joinpath("secret.py").write_text("SECRET = False\n", encoding="utf-8") + assert runner.main(["evaluate", str(job)]) == 2 + + with pytest.raises(ValueError, match="escapes"): + runner.safe_relative_path(tmp_path, "../outside.py") + assert tmp_path.joinpath("secret.py").read_text(encoding="utf-8") == "SECRET = True\n" + + draft.joinpath("secret.py").unlink() + link = draft / "linked.py" + try: + link.symlink_to(tmp_path / "secret.py") + except OSError: + pytest.skip("symlinks are unavailable on this platform") + assert runner.main(["evaluate", str(job)]) == 2 + + +def test_candidate_creation_uses_only_trust_contract_patterns(tmp_path, monkeypatch): + runner = _load_runner() + job, _, _ = _initialize_fake_campaign(runner, tmp_path, monkeypatch) + config_path = tmp_path / "autofl.yaml" + config = runner.read_yaml(config_path) + config["trust_contract"]["allowed_create_patterns"] = ["algorithms/*.py"] + runner.write_yaml(config_path, config) + + assert runner.main(["prepare", str(job), "--name", "new_module", "--hypothesis", "add algorithm"]) == 0 + draft = tmp_path / ".nvflare/autofl/candidates/new_module/source" + draft.joinpath("new_module.py").write_text("VALUE = 1\n", encoding="utf-8") + assert runner.main(["evaluate", str(job)]) == 2 + + draft.joinpath("new_module.py").unlink() + draft.joinpath("algorithms").mkdir() + draft.joinpath("algorithms/new_module.py").write_text("VALUE = 1\n", encoding="utf-8") + monkeypatch.setattr( + runner, + "run_job", + lambda run_def, **kwargs: runner.RunRecord( + "candidate", run_def.name, 0.7, 1.0, "none", run_def.description, "python job.py", "/tmp/candidate" + ), + ) + assert runner.main(["evaluate", str(job)]) == 0 + assert tmp_path.joinpath("algorithms/new_module.py").is_file() + + +def test_missing_create_patterns_deny_new_source(tmp_path, monkeypatch): + runner = _load_runner() + job, _, _ = _initialize_fake_campaign(runner, tmp_path, monkeypatch) + config_path = tmp_path / "autofl.yaml" + config = runner.read_yaml(config_path) + config["trust_contract"].pop("allowed_create_patterns") + runner.write_yaml(config_path, config) + + assert runner.main(["prepare", str(job), "--name", "denied_module", "--hypothesis", "add algorithm"]) == 0 + draft = tmp_path / ".nvflare/autofl/candidates/denied_module/source" + draft.joinpath("new_module.py").write_text("VALUE = 1\n", encoding="utf-8") + + assert runner.main(["evaluate", str(job)]) == 2 + + +def test_candidate_rejects_stale_manifest_and_budget_drift(tmp_path, monkeypatch): + runner = _load_runner() + job, client, config = _initialize_fake_campaign(runner, tmp_path, monkeypatch) + assert runner.main(["prepare", str(job), "--name", "stale", "--hypothesis", "change code"]) == 0 + candidate_dir = tmp_path / ".nvflare" / "autofl" / "candidates" / "stale" + candidate_dir.joinpath("source/client.py").write_text("ALGORITHM = 'candidate'\n", encoding="utf-8") + manifest_path = candidate_dir / "candidate_manifest.json" + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + manifest["base_source_sha256"] = "0" * 64 + manifest_path.write_text(json.dumps(manifest), encoding="utf-8") + assert runner.main(["evaluate", str(job)]) == 2 + + manifest["base_source_sha256"] = json.loads( + tmp_path.joinpath(".nvflare/autofl/campaign.json").read_text(encoding="utf-8") + )["best_source_sha256"] + manifest_path.write_text(json.dumps(manifest), encoding="utf-8") + drifted = deepcopy(config) + drifted["budget"]["fixed_training_budget"]["num_rounds"] = 2 + monkeypatch.setattr(runner, "import_job_config", lambda *args, **kwargs: deepcopy(drifted)) + assert runner.main(["evaluate", str(job)]) == 2 + assert client.read_text(encoding="utf-8") == "ALGORITHM = 'baseline'\n" + + +def test_candidate_schema_failure_does_not_modify_workspace(tmp_path, monkeypatch): + runner = _load_runner() + job, client, _ = _initialize_fake_campaign(runner, tmp_path, monkeypatch) + assert runner.main(["prepare", str(job), "--name", "bad_schema", "--hypothesis", "change code"]) == 0 + candidate_dir = tmp_path / ".nvflare" / "autofl" / "candidates" / "bad_schema" + candidate_dir.joinpath("source/client.py").write_text("ALGORITHM = 'candidate'\n", encoding="utf-8") + tmp_path.joinpath("mutation_schema.yaml").write_text( + "comparison_budget_args:\n default_candidate_budget:\n run_timeout_seconds: fast\n", + encoding="utf-8", + ) + + assert runner.main(["evaluate", str(job)]) == 2 + assert client.read_text(encoding="utf-8") == "ALGORITHM = 'baseline'\n" + + +def test_candidate_partial_apply_failure_restores_workspace(tmp_path, monkeypatch): + runner = _load_runner() + job, client, _ = _initialize_fake_campaign(runner, tmp_path, monkeypatch) + baseline_job = job.read_text(encoding="utf-8") + baseline_client = client.read_text(encoding="utf-8") + assert runner.main(["prepare", str(job), "--name", "partial", "--hypothesis", "change two files"]) == 0 + draft = tmp_path / ".nvflare" / "autofl" / "candidates" / "partial" / "source" + draft.joinpath("client.py").write_text("ALGORITHM = 'candidate'\n", encoding="utf-8") + draft.joinpath("job.py").write_text("print('candidate')\n", encoding="utf-8") + original_copy = runner.copy_relative_file + + def fail_second_candidate_copy(source_root, destination_root, relative): + if source_root == draft and relative == "job.py": + raise OSError("simulated candidate copy failure") + original_copy(source_root, destination_root, relative) + + monkeypatch.setattr(runner, "copy_relative_file", fail_second_candidate_copy) + + assert runner.main(["evaluate", str(job)]) == 2 + assert job.read_text(encoding="utf-8") == baseline_job + assert client.read_text(encoding="utf-8") == baseline_client + + +def test_candidate_job_help_failure_restores_workspace(tmp_path, monkeypatch): + runner = _load_runner() + job, client, _ = _initialize_fake_campaign(runner, tmp_path, monkeypatch) + assert runner.main(["prepare", str(job), "--name", "help_failure", "--hypothesis", "change code"]) == 0 + draft = tmp_path / ".nvflare" / "autofl" / "candidates" / "help_failure" / "source" + draft.joinpath("client.py").write_text("ALGORITHM = 'candidate'\n", encoding="utf-8") + + def fail_job_help(*args, **kwargs): + raise OSError("simulated missing Python executable") + + monkeypatch.setattr(runner, "job_help", fail_job_help) + + assert runner.main(["evaluate", str(job)]) == 2 + assert client.read_text(encoding="utf-8") == "ALGORITHM = 'baseline'\n" + + +def test_keyboard_interrupt_during_candidate_import_restores_workspace(tmp_path, monkeypatch): + runner = _load_runner() + job, client, _ = _initialize_fake_campaign(runner, tmp_path, monkeypatch) + assert runner.main(["prepare", str(job), "--name", "interrupt", "--hypothesis", "change code"]) == 0 + draft = tmp_path / ".nvflare" / "autofl" / "candidates" / "interrupt" / "source" + draft.joinpath("client.py").write_text("ALGORITHM = 'candidate'\n", encoding="utf-8") + + def interrupt_import(*args, **kwargs): + raise KeyboardInterrupt + + monkeypatch.setattr(runner, "import_job_config", interrupt_import) + args = runner.parse_args(["evaluate", str(job)]) + with pytest.raises(KeyboardInterrupt): + runner.evaluate_candidate(args, job) + + assert client.read_text(encoding="utf-8") == "ALGORITHM = 'baseline'\n" + + +def test_candidate_finalization_failure_rolls_back_workspace_and_campaign_files(tmp_path, monkeypatch): + runner = _load_runner() + job, client, _ = _initialize_fake_campaign(runner, tmp_path, monkeypatch) + assert runner.main(["prepare", str(job), "--name", "late_failure", "--hypothesis", "improve code"]) == 0 + candidate_dir = tmp_path / ".nvflare" / "autofl" / "candidates" / "late_failure" + candidate_dir.joinpath("source/client.py").write_text("ALGORITHM = 'candidate'\n", encoding="utf-8") + original_autofl = tmp_path.joinpath("autofl.yaml").read_bytes() + original_results = tmp_path.joinpath("results.tsv").read_bytes() + original_state = tmp_path.joinpath(".nvflare/autofl/campaign_state.json").read_bytes() + original_progress = tmp_path.joinpath("progress.png").read_bytes() + original_report = tmp_path.joinpath("autofl_report.md").read_bytes() + + def improved_run(run_def, **kwargs): + return runner.RunRecord( + "candidate", run_def.name, 0.7, 2.0, "none", run_def.description, "python job.py", "/tmp/late_failure" + ) + + original_refresh = runner.refresh_campaign_artifacts + + def fail_artifact_refresh(*args, **kwargs): + original_refresh(*args, **kwargs) + raise OSError("simulated report write failure") + + monkeypatch.setattr(runner, "run_job", improved_run) + monkeypatch.setattr(runner, "refresh_campaign_artifacts", fail_artifact_refresh) + + assert runner.main(["evaluate", str(job)]) == 2 + assert client.read_text(encoding="utf-8") == "ALGORITHM = 'baseline'\n" + assert tmp_path.joinpath("autofl.yaml").read_bytes() == original_autofl + assert tmp_path.joinpath("results.tsv").read_bytes() == original_results + assert tmp_path.joinpath(".nvflare/autofl/campaign_state.json").read_bytes() == original_state + assert tmp_path.joinpath("progress.png").read_bytes() == original_progress + assert tmp_path.joinpath("autofl_report.md").read_bytes() == original_report + best_source, best_files = runner.load_best_snapshot(tmp_path / ".nvflare" / "autofl" / "snapshots" / "best") + assert runner.workspace_matches_snapshot(tmp_path, best_source, best_files) + metadata = json.loads(tmp_path.joinpath(".nvflare/autofl/campaign.json").read_text(encoding="utf-8")) + assert metadata["best_candidate"] == "baseline" + manifest = json.loads(candidate_dir.joinpath("candidate_manifest.json").read_text(encoding="utf-8")) + assert manifest["status"] == "prepared" + + +def test_candidate_snapshot_stage_failure_preserves_previous_best(tmp_path, monkeypatch): + runner = _load_runner() + job, client, _ = _initialize_fake_campaign(runner, tmp_path, monkeypatch) + assert runner.main(["prepare", str(job), "--name", "snapshot_failure", "--hypothesis", "improve code"]) == 0 + draft = tmp_path / ".nvflare" / "autofl" / "candidates" / "snapshot_failure" / "source" + draft.joinpath("client.py").write_text("ALGORITHM = 'candidate'\n", encoding="utf-8") + + def improved_run(run_def, **kwargs): + return runner.RunRecord( + "candidate", run_def.name, 0.7, 2.0, "none", run_def.description, "python job.py", "/tmp/snapshot" + ) + + def fail_snapshot_stage(*args, **kwargs): + raise OSError("simulated snapshot copy failure") + + monkeypatch.setattr(runner, "run_job", improved_run) + monkeypatch.setattr(runner, "stage_best_snapshot", fail_snapshot_stage) + + assert runner.main(["evaluate", str(job)]) == 2 + assert client.read_text(encoding="utf-8") == "ALGORITHM = 'baseline'\n" + best_source, best_files = runner.load_best_snapshot(tmp_path / ".nvflare" / "autofl" / "snapshots" / "best") + assert runner.workspace_matches_snapshot(tmp_path, best_source, best_files) + + +def test_candidate_discard_restore_failure_retries_rollback(tmp_path, monkeypatch): + runner = _load_runner() + job, client, _ = _initialize_fake_campaign(runner, tmp_path, monkeypatch) + assert runner.main(["prepare", str(job), "--name", "discard_failure", "--hypothesis", "regress code"]) == 0 + draft = tmp_path / ".nvflare" / "autofl" / "candidates" / "discard_failure" / "source" + draft.joinpath("client.py").write_text("ALGORITHM = 'candidate'\n", encoding="utf-8") + original_restore = runner.restore_best_source + restore_calls = 0 + + def fail_first_restore(*args, **kwargs): + nonlocal restore_calls + restore_calls += 1 + if restore_calls == 1: + raise OSError("simulated restore failure") + original_restore(*args, **kwargs) + + def regressed_run(run_def, **kwargs): + return runner.RunRecord( + "candidate", run_def.name, 0.3, 2.0, "none", run_def.description, "python job.py", "/tmp/discard_failure" + ) + + monkeypatch.setattr(runner, "restore_best_source", fail_first_restore) + monkeypatch.setattr(runner, "run_job", regressed_run) + + assert runner.main(["evaluate", str(job)]) == 2 + assert restore_calls == 2 + assert client.read_text(encoding="utf-8") == "ALGORITHM = 'baseline'\n" + + +def test_malformed_yaml_returns_clean_cli_errors(tmp_path, monkeypatch, capsys): + runner = _load_runner() + job, _, config = _initialize_fake_campaign(runner, tmp_path, monkeypatch) + capsys.readouterr() + autofl_yaml = tmp_path / "autofl.yaml" + autofl_yaml.write_text("job: [\n", encoding="utf-8") + + assert runner.main(["suggest", str(job)]) == 2 + stderr = capsys.readouterr().err + assert f"Auto-FL suggest failed: invalid YAML in {autofl_yaml}" in stderr + assert "Traceback" not in stderr + + runner.write_yaml(autofl_yaml, config) + mutation_schema = tmp_path / "mutation_schema.yaml" + mutation_schema.write_text("comparison_budget_args: [\n", encoding="utf-8") + + assert runner.main(["suggest", str(job)]) == 2 + stderr = capsys.readouterr().err + assert f"Auto-FL suggest failed: invalid YAML in {mutation_schema}" in stderr + assert "Traceback" not in stderr + + +def test_abandon_candidate_clears_pending_draft_without_touching_best(tmp_path, monkeypatch): + runner = _load_runner() + job, client, _ = _initialize_fake_campaign(runner, tmp_path, monkeypatch) + assert runner.main(["prepare", str(job), "--name", "abandoned", "--hypothesis", "temporary idea"]) == 0 + draft = tmp_path / ".nvflare" / "autofl" / "candidates" / "abandoned" / "source" / "client.py" + draft.write_text("ALGORITHM = 'temporary'\n", encoding="utf-8") + + assert runner.main(["abandon", str(job)]) == 0 + assert client.read_text(encoding="utf-8") == "ALGORITHM = 'baseline'\n" + manifest = json.loads( + tmp_path.joinpath(".nvflare/autofl/candidates/abandoned/candidate_manifest.json").read_text(encoding="utf-8") + ) + assert manifest["status"] == "abandoned" + state = json.loads(tmp_path.joinpath(".nvflare/autofl/campaign_state.json").read_text(encoding="utf-8")) + assert state["next_action"] == "propose_candidate" + assert state["pending_candidate_manifest"] is None + + +def test_abandon_rejects_agent_modified_manifest_paths(tmp_path, monkeypatch): + runner = _load_runner() + job, client, _ = _initialize_fake_campaign(runner, tmp_path, monkeypatch) + assert runner.main(["prepare", str(job), "--name", "tampered", "--hypothesis", "temporary idea"]) == 0 + draft = tmp_path / ".nvflare/autofl/candidates/tampered/source/client.py" + draft.write_text("ALGORITHM = 'temporary'\n", encoding="utf-8") + manifest_path = tmp_path / ".nvflare/autofl/candidates/tampered/candidate_manifest.json" + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + manifest["changed_files"] = ["../client.py"] + manifest_path.write_text(json.dumps(manifest), encoding="utf-8") + + assert runner.main(["abandon", str(job)]) == 2 + assert client.read_text(encoding="utf-8") == "ALGORITHM = 'baseline'\n" + + +def test_status_rescans_pending_manifests_before_writing_state(tmp_path, monkeypatch): + runner = _load_runner() + job, _, _ = _initialize_fake_campaign(runner, tmp_path, monkeypatch) + assert runner.main(["prepare", str(job), "--name", "pending", "--hypothesis", "change code"]) == 0 + state_path = tmp_path / ".nvflare" / "autofl" / "campaign_state.json" + state_path.write_text('{"final_response_allowed": true}\n', encoding="utf-8") + + assert runner.main(["status", str(job)]) == 0 + + state = json.loads(state_path.read_text(encoding="utf-8")) + assert state["reason"] == "pending_candidates" + assert state["next_action"] == "edit_candidate" + assert state["final_response_allowed"] is False + assert state["pending_candidate_manifest"].endswith("pending/candidate_manifest.json") + + +def test_unchanged_status_does_not_regenerate_campaign_artifacts(tmp_path, monkeypatch): + runner = _load_runner() + job, _, _ = _initialize_fake_campaign(runner, tmp_path, monkeypatch) + artifact_paths = [tmp_path / "results.tsv", tmp_path / "progress.png", tmp_path / "autofl_report.md"] + before = {path: (path.read_bytes(), path.stat().st_mtime_ns) for path in artifact_paths} + + assert runner.main(["status", str(job)]) == 0 + first_state = tmp_path.joinpath(".nvflare/autofl/campaign_state.json").read_bytes() + first_state_mtime = tmp_path.joinpath(".nvflare/autofl/campaign_state.json").stat().st_mtime_ns + assert runner.main(["status", str(job)]) == 0 + + assert {path: (path.read_bytes(), path.stat().st_mtime_ns) for path in artifact_paths} == before + assert tmp_path.joinpath(".nvflare/autofl/campaign_state.json").read_bytes() == first_state + assert tmp_path.joinpath(".nvflare/autofl/campaign_state.json").stat().st_mtime_ns == first_state_mtime + + +def test_status_refuses_malformed_candidate_manifest(tmp_path, monkeypatch): + runner = _load_runner() + job, _, _ = _initialize_fake_campaign(runner, tmp_path, monkeypatch) + assert runner.main(["prepare", str(job), "--name", "malformed", "--hypothesis", "change code"]) == 0 + manifest = tmp_path / ".nvflare" / "autofl" / "candidates" / "malformed" / "candidate_manifest.json" + manifest.write_text("not json\n", encoding="utf-8") + + assert runner.main(["status", str(job)]) == 2 + + +def test_status_uses_persisted_custom_stop_file_in_job_directory(tmp_path, monkeypatch): + runner = _load_runner() + job, _, _ = _initialize_fake_campaign(runner, tmp_path, monkeypatch) + metadata_path = tmp_path / ".nvflare" / "autofl" / "campaign.json" + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + metadata["settings"]["stop_file"] = ["CUSTOM_STOP"] + metadata_path.write_text(json.dumps(metadata), encoding="utf-8") + tmp_path.joinpath("STOP_AUTOFL").touch() + + assert runner.main(["status", str(job)]) == 0 + default_state = json.loads(tmp_path.joinpath(".nvflare/autofl/campaign_state.json").read_text(encoding="utf-8")) + assert default_state["reason"] == "continue" + + tmp_path.joinpath("CUSTOM_STOP").touch() + assert runner.main(["status", str(job)]) == 0 + custom_state = json.loads(tmp_path.joinpath(".nvflare/autofl/campaign_state.json").read_text(encoding="utf-8")) + assert custom_state["reason"] == "manual_stop_file" + assert custom_state["final_response_allowed"] is True + + +def test_initialize_retries_an_unscored_baseline(tmp_path, monkeypatch): + runner = _load_runner() + job = tmp_path / "job.py" + job.write_text("print('job')\n", encoding="utf-8") + tmp_path.joinpath("client.py").write_text("ALGORITHM = 'baseline'\n", encoding="utf-8") + monkeypatch.setattr(runner, "import_job_config", lambda *args, **kwargs: deepcopy(_campaign_config())) + monkeypatch.setattr(runner, "job_help", lambda *args, **kwargs: "") + scores = iter([None, 0.5]) + + def fake_run(run_def, **kwargs): + return runner.RunRecord( + "baseline", + run_def.name, + next(scores), + 1.0, + "none", + "baseline", + "python job.py", + "/tmp/baseline", + ) + + monkeypatch.setattr(runner, "run_job", fake_run) + command = ["initialize", str(job), "--no-prefer-synthetic"] + assert runner.main(command) == 1 + assert runner.main(command) == 0 + records = runner.load_results(tmp_path / "results.tsv") + assert [(record.status, record.score) for record in records] == [("baseline", None), ("baseline", 0.5)] + assert [record.name for record in records] == ["baseline", "baseline_retry_2"] + metadata = json.loads(tmp_path.joinpath(".nvflare/autofl/campaign.json").read_text(encoding="utf-8")) + assert metadata["best_candidate"] == "baseline_retry_2" + + +def test_record_literature_checkpoint_returns_to_agent_proposal(tmp_path, monkeypatch): + runner = _load_runner() + job, _, _ = _initialize_fake_campaign(runner, tmp_path, monkeypatch) + + assert ( + runner.main( + [ + "record", + str(job), + "--literature", + "--hypothesis", + "reviewed adaptive federated optimization", + ] + ) + == 0 + ) + records = runner.load_results(tmp_path / "results.tsv") + assert records[-1].status == "literature" + assert records[-1].diff_summary == "reviewed adaptive federated optimization" + state = json.loads(tmp_path.joinpath(".nvflare/autofl/campaign_state.json").read_text(encoding="utf-8")) + assert state["next_action"] == "propose_candidate" + assert state["required_exploration"] == "source_backed_server_aggregation" + assert "server aggregation candidate" in state["agent_instruction"] + + +def test_external_baseline_may_follow_literature_event(tmp_path, monkeypatch): + runner = _load_runner() + job, _, _ = _initialize_fake_campaign(runner, tmp_path, monkeypatch, target_env="prod") + + assert runner.main(["record", str(job), "--literature", "--hypothesis", "reviewed FedOpt"]) == 0 + assert runner.main(["record", str(job), "--baseline", "--score", "0.5", "--job-id", "baseline-job"]) == 0 + + records = runner.load_results(tmp_path / "results.tsv") + assert [record.status for record in records] == ["literature", "baseline"] + + +def test_omitted_metric_uses_imported_job_metric(tmp_path, monkeypatch): + runner = _load_runner() + job = tmp_path / "job.py" + job.write_text("print('job')\n", encoding="utf-8") + tmp_path.joinpath("client.py").write_text("ALGORITHM = 'baseline'\n", encoding="utf-8") + config = _campaign_config() + config["objective"] = { + "metric": "auc", + "requested_metric": "auc", + "optimization_metric": "auc", + "metric_extraction_order": ["auc"], + "metric_contract_source": "arg:key_metric", + } + monkeypatch.setattr(runner, "import_job_config", lambda *args, **kwargs: deepcopy(config)) + monkeypatch.setattr(runner, "job_help", lambda *args, **kwargs: "") + monkeypatch.setattr(runner, "write_progress", lambda path, *args: path.write_bytes(b"progress")) + + def fake_run(run_def, **kwargs): + assert kwargs["metrics"] == ["auc"] + return runner.RunRecord( + "baseline", run_def.name, 0.6, 1.0, "none", "baseline", "python job.py", "/tmp/baseline" + ) + + monkeypatch.setattr(runner, "run_job", fake_run) + + assert runner.main(["initialize", str(job), "--no-prefer-synthetic"]) == 0 + metadata = json.loads(tmp_path.joinpath(".nvflare/autofl/campaign.json").read_text(encoding="utf-8")) + assert metadata["settings"]["metric"] == "auc" + + +def test_explicit_mutable_campaign_settings_persist_and_uncapped_removes_cap(tmp_path, monkeypatch): + runner = _load_runner() + job, _, _ = _initialize_fake_campaign(runner, tmp_path, monkeypatch) + + assert runner.main(["status", str(job), "--max-candidates", "7", "--timeout", "123"]) == 0 + metadata_path = tmp_path / ".nvflare/autofl/campaign.json" + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + assert metadata["settings"]["max_candidates"] == 7 + assert metadata["settings"]["timeout"] == 123 + + assert runner.main(["status", str(job), "--uncapped"]) == 0 + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + assert metadata["settings"]["max_candidates"] is None + + +def test_explicit_immutable_campaign_setting_change_is_rejected(tmp_path, monkeypatch): + runner = _load_runner() + job, _, _ = _initialize_fake_campaign(runner, tmp_path, monkeypatch) + + assert runner.main(["status", str(job), "--metric", "loss"]) == 2 + + +def test_external_candidate_uses_standard_job_result_recording(tmp_path, monkeypatch): + runner = _load_runner() + job, client, _ = _initialize_fake_campaign(runner, tmp_path, monkeypatch, target_env="prod") + assert runner.main(["record", str(job), "--baseline", "--score", "0.5", "--job-id", "job-baseline"]) == 0 + assert runner.main(["prepare", str(job), "--name", "prod_algo", "--hypothesis", "production algorithm"]) == 0 + draft = tmp_path / ".nvflare" / "autofl" / "candidates" / "prod_algo" / "source" + draft.joinpath("client.py").write_text("ALGORITHM = 'production'\n", encoding="utf-8") + assert runner.main(["evaluate", str(job)]) == 0 + assert client.read_text(encoding="utf-8") == "ALGORITHM = 'production'\n" + + manifest_path = tmp_path / ".nvflare" / "autofl" / "candidates" / "prod_algo" / "candidate_manifest.json" + assert json.loads(manifest_path.read_text(encoding="utf-8"))["status"] == "ready_for_external_execution" + assert ( + json.loads(tmp_path.joinpath(".nvflare/autofl/campaign_state.json").read_text(encoding="utf-8"))["next_action"] + == "submit_candidate" + ) + assert ( + runner.main( + [ + "record", + str(job), + "--manifest", + str(manifest_path), + "--score", + "0.8", + "--job-id", + "job-candidate", + ] + ) + == 0 + ) + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + assert manifest["status"] == "keep" + assert manifest["artifacts"]["job_id"] == "job-candidate" + + +def test_suggest_returns_fallbacks_without_executing_them(tmp_path, monkeypatch, capsys): + runner = _load_runner() + job, _, _ = _initialize_fake_campaign(runner, tmp_path, monkeypatch) + capsys.readouterr() + monkeypatch.setattr(runner, "job_help", lambda *args, **kwargs: "--lr") + monkeypatch.setattr( + runner, + "run_job", + lambda *args, **kwargs: pytest.fail("suggest must not execute a candidate"), + ) + + assert runner.main(["suggest", str(job), "--limit", "2"]) == 0 + payload = json.loads(capsys.readouterr().out) + assert len(payload["suggestions"]) == 2 + assert all(item["run_args"] for item in payload["suggestions"]) + + +def test_import_job_config_forwards_minimization_mode(tmp_path, monkeypatch): + runner = _load_runner() + job = tmp_path / "job.py" + job.write_text("print('job')\n", encoding="utf-8") + output = tmp_path / "autofl.yaml" + captured = {} + + class FakeImporter: + __file__ = __file__ + + @staticmethod + def import_job_to_autofl_config(*args, **kwargs): + captured.update(kwargs) + return _campaign_config() + + @staticmethod + def dump_autofl_yaml(config): + return runner.yaml.safe_dump(config) + + monkeypatch.setattr(runner, "load_job_importer", lambda: FakeImporter) + args = runner.parse_args(["initialize", str(job), "--mode", "min"]) + runner.import_job_config(args, job, output, tmp_path / "import.log", 10) + + assert captured["mode"] == "min" + + +@pytest.mark.parametrize( + "config,expected", + [ + ( + { + "import": {"support": {"status": "partial"}}, + "budget": {"fixed_training_budget": {"num_rounds": 1}}, + }, + "job surface", + ), + ( + {"import": {"support": {"status": "supported"}}, "budget": {}}, + "fixed comparison budget", + ), + ( + { + "import": {"support": {"status": "supported"}}, + "budget": {"fixed_training_budget": {"num_rounds": 1}}, + "unresolved": [{"field": "budget.fixed_training_budget.num_clients", "reason": "dynamic"}], + }, + "safety-critical fields", + ), + ], +) +def test_campaign_admission_rejects_unresolved_safety_contract(config, expected): + runner = _load_runner() + + assert expected in "; ".join(runner.campaign_admission_errors(config)) + + +def test_cli_lifecycle_runs_agent_code_candidate_end_to_end(tmp_path): + repo_root = Path(__file__).parents[3] + runner_path = repo_root / "skills" / "nvflare-autofl" / "scripts" / "run_job_campaign.py" + job = tmp_path / "job.py" + simulation_root = tmp_path / "simulation" + job.write_text( + f""" +import argparse +import json +from pathlib import Path +from nvflare.app_opt.pt.recipes.fedavg import FedAvgRecipe as ImportedFedAvgRecipe +from nvflare.recipe import SimEnv as ImportedSimEnv + +SCORE = 0.5 + +class FakeFedAvgRecipe: + def __init__(self, **kwargs): + self.kwargs = kwargs + +class FakeSimEnv: + def __init__(self, **kwargs): + self.kwargs = kwargs + +ImportedFedAvgRecipe = FakeFedAvgRecipe +ImportedSimEnv = FakeSimEnv + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--name", default="run") + parser.add_argument("--num_rounds", type=int, default=1) + parser.add_argument("--n_clients", type=int, default=2) + args = parser.parse_args() + ImportedFedAvgRecipe(model=object(), num_rounds=args.num_rounds, min_clients=args.n_clients) + ImportedSimEnv(num_clients=args.n_clients, workspace_root={str(simulation_root)!r}) + result = Path({str(simulation_root)!r}) / args.name + result.mkdir(parents=True, exist_ok=True) + result.joinpath("metrics_summary.json").write_text(json.dumps({{"accuracy": SCORE}})) + print(f"Result can be found in : {{result.resolve()}}") + +if __name__ == "__main__": + main() +""".lstrip(), + encoding="utf-8", + ) + env = os.environ.copy() + env["PYTHONPATH"] = os.pathsep.join(filter(None, [str(repo_root), env.get("PYTHONPATH")])) + + subprocess.run( + [ + sys.executable, + str(runner_path), + "initialize", + str(job), + "--metric", + "accuracy", + "--no-prefer-synthetic", + ], + cwd=tmp_path, + env=env, + check=True, + capture_output=True, + text=True, + ) + subprocess.run( + [ + sys.executable, + str(runner_path), + "prepare", + str(job), + "--name", + "code_candidate", + "--hypothesis", + "raise the reported score through a source change", + ], + cwd=tmp_path, + env=env, + check=True, + capture_output=True, + text=True, + ) + draft_job = tmp_path / ".nvflare" / "autofl" / "candidates" / "code_candidate" / "source" / "job.py" + draft_job.write_text(draft_job.read_text(encoding="utf-8").replace("SCORE = 0.5", "SCORE = 0.8"), encoding="utf-8") + subprocess.run( + [sys.executable, str(runner_path), "evaluate", str(job)], + cwd=tmp_path, + env=env, + check=True, + capture_output=True, + text=True, + ) + + runner = _load_runner() + records = runner.load_results(tmp_path / "results.tsv") + assert [(record.status, record.score) for record in records] == [("baseline", 0.5), ("keep", 0.8)] + assert "SCORE = 0.8" in job.read_text(encoding="utf-8") + manifest = json.loads( + tmp_path.joinpath(".nvflare/autofl/candidates/code_candidate/candidate_manifest.json").read_text( + encoding="utf-8" + ) + ) + assert manifest["status"] == "keep" + assert manifest["changed_files"] == ["job.py"]