diff --git a/.gitignore b/.gitignore index 8c9916c6c..cf098093c 100644 --- a/.gitignore +++ b/.gitignore @@ -151,3 +151,10 @@ pyrightconfig.json /tests/test_data/.data_paths/ /logs/ poetry.toml + +# Project specific +/runs/ +/.venv-play/ +/controller.local.toml +/localworkers_config.json +/status.txt diff --git a/documentation/flower_phase1/DESIGN_DECISIONS.md b/documentation/flower_phase1/DESIGN_DECISIONS.md new file mode 100644 index 000000000..97fd57148 --- /dev/null +++ b/documentation/flower_phase1/DESIGN_DECISIONS.md @@ -0,0 +1,39 @@ +# Design Decisions - Phase 1 (Simulation-First) + +## Scope +The Phase 1 implementation provides federated training orchestration through Exaflow + Flower using simulation runtime. +It does not include production hospital deployment, SuperLink/SuperNodes, or compliance hardening. + +## Core Architecture Split +1. Exaflow = control plane. +2. Flower = training/aggregation engine. +3. Flower App = algorithm-specific training logic (flowertune_llm_medical). + +## Source of Truth Policy +1. For Phase 1, source of truth for training data/runtime knobs is `parameters.dataset`. +2. `inputdata` remains for compatibility and metadata. + +## Integration Strategy +1. Do not vendor the full `adap/flower` repository. +2. Integrate only the Flower app code as an Exaflow algorithm package. +3. Keep Flower framework as a dependency. + +## Runtime Policy +1. Only `runtime = simulation` is supported in Phase 1. +2. The same API shape will extend to Phase 2 deployment runtime. +3. `dataset` block semantics may change under `runtime=deployment` (local connectors instead of simulation partitioning). + +## Exaflow Boundary +1. Exaflow does not manage model internals (LoRA weights, state dicts). +2. Exaflow handles config, lifecycle status, metrics, and artifact references. + +## Contract Versioning +1. `schema_version = "1.1"` for RunConfig and MetricsOut in this freeze. +2. Backward-incompatible changes require a new major schema version. +3. Backward-compatible additions may remain in the same major series. + +## Serialization Policy +1. Boolean env values: `"true"` / `"false"`. +2. List env values: JSON string. +3. Optional `None`: env var omitted. +4. `${request_id}` interpolation is resolved by Exaflow before dispatch. diff --git a/documentation/flower_phase1/PHASE1_ACCEPTANCE.md b/documentation/flower_phase1/PHASE1_ACCEPTANCE.md new file mode 100644 index 000000000..035e6a03b --- /dev/null +++ b/documentation/flower_phase1/PHASE1_ACCEPTANCE.md @@ -0,0 +1,37 @@ +# Phase 1 Acceptance Criteria + +## A. Smoke +1. Job starts from Exaflow endpoint. +2. Simulation completes at least 2 rounds without crash. +3. Terminal status is produced. + +## B. Metrics +1. At least one valid `round_event` per completed round. +2. `requested_metrics` is always present. +3. `reported_metrics` follows policy: + - RUNNING rounds: equals keys of `metrics`. + - failures: empty allowed. +4. One `final_summary` is always emitted. +5. Timestamps are serialized in UTC with `Z` suffix. + +## C. Artifacts +1. `artifact_dir` resolves and is writable. +2. Checkpoints are saved when configured. +3. On `COMPLETED`, `artifacts.final_checkpoint_ref` exists. + +## D. Failure/Timeout/Cancel +1. On `FAILED/CANCELLED/TIMEOUT`, `error` exists in final summary. +2. `artifacts` is optional for non-completed final states. +3. Partial/empty metrics payloads are accepted for failed terminal states. +4. `CANCELLED` is best-effort stop and still emits terminal summary. + +## E. Contract Strictness +1. Unknown RunConfig fields are rejected. +2. Invalid combinations (e.g., `local_steps` + `epochs`) are rejected. +3. Invalid `runtime != simulation` is rejected in Phase 1. +4. Duplicate `request_id` follows idempotency policy (no duplicate run creation). + +## F. Non-Goals +1. No real hospital connectors. +2. No SuperLink/SuperNodes deployment runtime. +3. No production security/compliance hardening in Phase 1. diff --git a/documentation/flower_phase1/RUNTIME_CONTRACT.md b/documentation/flower_phase1/RUNTIME_CONTRACT.md new file mode 100644 index 000000000..b2945f266 --- /dev/null +++ b/documentation/flower_phase1/RUNTIME_CONTRACT.md @@ -0,0 +1,65 @@ +# Runtime Contract - Phase 1 + +## Trigger +`POST /algorithms/flowertune_llm_medical` + +## Request Shape +1. Exaflow standard `AlgorithmRequestDTO`. +2. RunConfig is provided inside `parameters`. + +## RunConfig -> Env Mapping Rules +1. Flatten nested config into env vars. +2. booleans -> `"true"` / `"false"`. +3. lists -> JSON strings. +4. `None` -> omit env var. +5. Canonical full naming for env vars (e.g., `MIN_EVALUATE_CLIENTS`). + +## Lifecycle +1. `SUBMITTED` +2. `RUNNING` +3. Terminal: `COMPLETED` | `FAILED` | `CANCELLED` | `TIMEOUT` + +## Concurrency Scope (Phase 1) +1. Single active Flower job is supported at a time. +2. Flower helper endpoints are request-scoped via `request_id` query parameter. + +## Event Flow +1. Exaflow receives and validates request. +2. Exaflow creates job id and resolves artifact path. +3. Exaflow dispatches Flower simulation run. +4. Flower emits `round_event` payloads. +5. Flower emits one terminal `final_summary` payload. +6. Exaflow persists status, metrics, and artifacts references. + +## MetricsOut Status Semantics +1. `round_event.status` is only `RUNNING` or `FAILED`. +2. `COMPLETED` appears only in `final_summary.status`. +3. `round_event.artifacts` object is always present; `checkpoint_ref` is optional. + +## Observability Fields +1. `job_id` is required and the primary correlation key. +2. `request_id` is optional but recommended. +3. `correlation_id` is optional. +4. `algorithm_name`, `runtime`, `timestamp` are required. +5. Timestamps must be UTC with `Z` suffix. + +## Error Contract +1. Error payload shape: `code`, `message`, optional `details`. +2. `details` is intentionally open for runtime-specific diagnostics. +3. Canonical Phase 1 error codes: + - `VALIDATION_ERROR` + - `ARTIFACT_PATH_ERROR` + - `MODEL_LOAD_ERROR` + - `DATASET_LOAD_ERROR` + - `QUANTIZATION_ERROR` + - `PEFT_CONFIG_ERROR` + - `ROUND_TIMEOUT` + - `JOB_TIMEOUT` + - `RUNTIME_ERROR` + - `CANCELLED_BY_USER` + +## Defaults Note +`default` values in JSON Schema are informational only. Defaults are applied by producer/service logic, not by JSON Schema validators. + +## Cancellation Semantics +Cancellation is best-effort stop for Phase 1 simulation. A terminal `final_summary` with `status=CANCELLED` is required. diff --git a/documentation/flower_phase1/VALIDATION_SPEC.md b/documentation/flower_phase1/VALIDATION_SPEC.md new file mode 100644 index 000000000..fb6e913f2 --- /dev/null +++ b/documentation/flower_phase1/VALIDATION_SPEC.md @@ -0,0 +1,52 @@ +# Validation Spec - Phase 1 + +## Validation Layers +1. Schema Validation (Pydantic). +2. Service Validation (Exaflow business/runtime checks). +3. Runtime Validation (Flower app startup/runtime checks). + +## A. Pydantic Validation (Request `parameters`) +1. `extra = "forbid"` on all models. +2. `runtime` must be `simulation`. +3. `schema_version` must be `1.1`. +4. `federation.num_rounds >= 1`, `num_clients >= 1`. +5. `min_fit_clients <= num_clients`. +6. `min_evaluate_clients <= num_clients`. +7. `min_available_clients <= num_clients`. +8. If `evaluation.enabled = false`: `fraction_evaluate = 0`, `min_evaluate_clients = 0`. +9. `local_steps XOR epochs` (exactly one set). +10. `fp16` and `bf16` cannot both be `true`. +11. `seed >= 0`. +12. `job_timeout_sec >= round_timeout_sec`. +13. For `partitioner = iid`: `dataset.num_partitions = federation.num_clients`. +14. If `peft.enabled = true`: `method`, `r`, `alpha` are required. +15. Metrics allowlist for Phase 1: `loss`, `perplexity`. + +## B. Service Validation (Exaflow) +1. Resolve `artifact_dir` template (`${request_id}`). +2. Verify writable artifact path. +3. Enforce capacity policy: + - max allowed `num_clients` for simulation, + - max allowed `model.max_seq_length`, + - allowed quantization modes, + - timeout upper bounds, + - optional GPU requirement policy. +4. Normalize env map with canonical serialization rules. +5. Enforce idempotency policy for duplicate `request_id` (return existing `job_id`). +6. Enforce event invariants at ingestion: + - `round <= rounds_total`, + - `rounds_completed <= rounds_total`, + - `clients.participated <= clients.expected`, + - `clients.failed <= clients.expected`, + - `clients.participated + clients.failed <= clients.expected`, + - if present: `best_round <= rounds_completed`. +7. Enforce metrics consistency policy: + - RUNNING round event: `reported_metrics == keys(metrics)`. + - FAILED/CANCELLED/TIMEOUT: empty metrics/report sets allowed. + +## C. Runtime Validation (Flower app) +1. Model/tokenizer load success. +2. Dataset load and partitioning success. +3. Quantization compatibility. +4. LoRA target modules validity. +5. Early failure produces structured `ErrorInfo`. diff --git a/exaflow/algorithms/federated/docs/glmm_binary.md b/exaflow/algorithms/federated/docs/glmm_binary.md new file mode 100644 index 000000000..b470d88e0 --- /dev/null +++ b/exaflow/algorithms/federated/docs/glmm_binary.md @@ -0,0 +1,74 @@ +## Binary GLMM (FederatedGLMMBinary) + +### Name + +**Binary GLMM (FederatedGLMMBinary)** + +### Type + +**Statistical Model** (binary outcome, random-intercept GLMM with Laplace approximation) + +### Goal (Why we need it) + +Estimate fixed effects and random-intercept variance for clustered binary outcomes in a federated setup, without exchanging row-level patient data. + +### When to use + +Use Binary GLMM when: + +* outcome is binary (`0/1`) +* data are clustered by center/site +* logistic regression without random effects is too optimistic +* you need center-level heterogeneity modeling (`sigma_u2`) + +### When NOT to use + +Avoid / be careful when: + +* outcome is ordinal (use `FederatedGLMMOrdinal`) +* no meaningful clustering exists +* you need random slopes (not currently supported) +* events are extremely rare with tiny per-center samples + +--- + +### Inputs / Outputs + +| Item | Description | +| ----------------- | ---------------------------------------------------- | +| **X** | Local feature matrix, shape `(n_local, p)` | +| **y** | Local binary labels in `{0, 1}` | +| **center_ids** | Local cluster ids, shape `(n_local,)` | +| **fit_intercept** | Whether to add intercept term (default: `True`) | +| **max_iters** | Maximum optimization iterations | +| **tol_theta** | Parameter-step tolerance | +| **tol_score** | Score-norm tolerance | +| **agg_client** | Federated aggregation client | + +**Outputs (FederatedGLMMBinaryResults)** + +* `theta`: full vector `[beta..., log_sigma_u2]` +* `params`: fixed effects `beta` +* `sigma_u2`: random-intercept variance +* `nobs`, `n_groups` +* `converged`, `n_iter` +* `predict(X)`: class probabilities +* `history` (optional): optimization diagnostics + +### Key Differences from centralized GLMM + +| Aspect | Centralized GLMM | Exaflow FederatedGLMMBinary | +| -------------------- | ------------------------ | --------------------------- | +| Data access | Full row-level data | Local-only + aggregated stats | +| Likelihood handling | Direct centralized solve | Federated Newton updates with Laplace terms | +| Random effects | Broad options | Random intercept | +| Compute topology | Single node | Multi-worker federation | + +### Approximation vs Exactness + +| Component | Status in FederatedGLMMBinary | +| --------------------------- | --------------------------------- | +| Fixed effects | Iterative estimate | +| Random-intercept variance | Iterative estimate (`sigma_u2`) | +| Likelihood treatment | Laplace approximation | +| Random slopes | Not supported | diff --git a/exaflow/algorithms/federated/docs/glmm_ordinal.md b/exaflow/algorithms/federated/docs/glmm_ordinal.md new file mode 100644 index 000000000..2a793b272 --- /dev/null +++ b/exaflow/algorithms/federated/docs/glmm_ordinal.md @@ -0,0 +1,78 @@ +## Ordinal GLMM (FederatedGLMMOrdinal) + +### Name + +**Ordinal GLMM (FederatedGLMMOrdinal)** + +### Type + +**Statistical Model** (ordered categorical outcome, random-intercept GLMM with Laplace approximation) + +### Goal (Why we need it) + +Model ordered outcomes (e.g., severity stages) with fixed effects, ordered cutpoints, and center-level random intercepts in a federated environment. + +### When to use + +Use Ordinal GLMM when: + +* outcome has natural order and `K >= 2` categories +* data are clustered by center/site +* proportional-odds style cumulative logit is appropriate +* you need both fixed effects and random-intercept variability + +### When NOT to use + +Avoid / be careful when: + +* outcome is binary only (use `FederatedGLMMBinary`) +* outcome is continuous (use `FederatedLMM`/`FederatedOLS`) +* category order is not meaningful +* you need random slopes (not currently supported) + +--- + +### Inputs / Outputs + +| Item | Description | +| ----------------- | ---------------------------------------------------- | +| **X** | Local feature matrix, shape `(n_local, p)` | +| **y** | Local labels in `{0, ..., K-1}` | +| **center_ids** | Local cluster ids, shape `(n_local,)` | +| **K** | Number of ordered categories | +| **fit_intercept** | Whether to add intercept term (default: `True`) | +| **max_iters** | Maximum optimization iterations | +| **tol_theta** | Parameter-step tolerance | +| **tol_score** | Score-norm tolerance | +| **agg_client** | Federated aggregation client | + +**Outputs (FederatedGLMMOrdinalResults)** + +* `theta`: full parameter vector +* `params`: fixed effects `beta` +* `cutpoints`: ordered category thresholds +* `sigma_u2`: random-intercept variance +* `nobs`, `n_groups` +* `converged`, `n_iter` +* `predict(X)`: predicted class label +* `predict_proba(X)`: class probabilities +* `history` (optional): optimization diagnostics + +### Key Differences from centralized ordinal mixed models + +| Aspect | Centralized model | Exaflow FederatedGLMMOrdinal | +| -------------------- | ------------------------ | ---------------------------- | +| Data access | Full row-level data | Local-only + aggregated terms | +| Ordinal structure | Cumulative logit | Cumulative logit | +| Random effects | Broad options | Random intercept | +| Compute topology | Single node | Multi-worker federation | + +### Approximation vs Exactness + +| Component | Status in FederatedGLMMOrdinal | +| --------------------------- | -------------------------------- | +| Fixed effects | Iterative estimate | +| Cutpoints | Iterative estimate (ordered) | +| Random-intercept variance | Iterative estimate (`sigma_u2`) | +| Likelihood treatment | Laplace approximation | +| Random slopes | Not supported | diff --git a/exaflow/algorithms/federated/docs/lmm.md b/exaflow/algorithms/federated/docs/lmm.md new file mode 100644 index 000000000..fe88a0179 --- /dev/null +++ b/exaflow/algorithms/federated/docs/lmm.md @@ -0,0 +1,78 @@ +## Linear Mixed Model (FederatedLMM) + +### Name + +**Linear Mixed Model (FederatedLMM)** + +### Type + +**Statistical Model** (continuous outcome, random-intercept mixed model with REML) + +### Goal (Why we need it) + +Estimate fixed effects and variance components (`sigma2`, `sigma_u2`) for clustered continuous outcomes while keeping row-level data local to each worker. + +### When to use + +Use LMM when: + +* outcome is continuous +* observations are grouped (e.g., center/hospital/site) +* you need fixed-effect inference and cluster variance estimation +* random-intercept structure is clinically/statistically reasonable + +### When NOT to use + +Avoid / be careful when: + +* outcome is binary/ordinal (use GLMM variants) +* cluster structure is not present +* you need random slopes (current implementation is random intercept only) +* cluster count is extremely small (variance estimates may be unstable) + +--- + +### Inputs / Outputs + +| Item | Description | +| ----------------- | ---------------------------------------------------- | +| **X** | Local feature matrix, shape `(n_local, p)` | +| **y** | Local continuous target, shape `(n_local,)` | +| **center_ids** | Local cluster ids, shape `(n_local,)` | +| **fit_intercept** | Whether to add intercept term (default: `True`) | +| **max_iter** | Maximum REML iterations (default: `80`) | +| **tol** | Convergence tolerance (default: `1e-8`) | +| **agg_client** | Federated aggregation client | +| **w** | Optional local non-negative weights | + +**Outputs (FederatedLMMResults)** + +* `params`: fixed-effect coefficients +* `bse`, `tvalues`, `pvalues`: inference for fixed effects +* `conf_int_low`, `conf_int_high`: confidence interval bounds +* `sigma2`: residual variance +* `sigma_u2`: random-intercept variance +* `cov_params`: covariance matrix of fixed effects +* `ll_reml`, `aic`, `bic`: fit quality metrics +* `nobs`, `n_groups`, `df_model`, `df_resid` +* `converged`, `n_iter` +* `predict(X)`: predicted mean outcome + +### Key Differences from statsmodels + +| Aspect | statsmodels MixedLM | Exaflow FederatedLMM | +| -------------------- | ------------------------ | ---------------------------- | +| Data access | Centralized | Stays local per worker | +| Optimization target | REML | REML | +| Random effects | Broad support | Random intercept | +| Computation | Single-process | Federated aggregated updates | +| API surface | Full statsmodels object | Simplified results container | + +### Approximation vs Exactness + +| Component | Status in FederatedLMM | +| --------------------------- | ------------------------------------ | +| Fixed effects | Exact for aggregated model equations | +| Variance components | Iterative REML optimization | +| Inference stats | Available for fixed effects | +| Random slopes | Not supported | diff --git a/exaflow/algorithms/flower/flowertune_llm_medical.json b/exaflow/algorithms/flower/flowertune_llm_medical.json new file mode 100644 index 000000000..e1588e81c --- /dev/null +++ b/exaflow/algorithms/flower/flowertune_llm_medical.json @@ -0,0 +1,27 @@ +{ + "name": "flowertune_llm_medical", + "desc": "Federated LLM fine-tuning contract wiring for simulation-first orchestration.", + "label": "FlowerTune LLM Medical", + "enabled": true, + "type": "flower", + "components": ["FLOWER"], + "inputdata": { + "y": { + "label": "Variable (dependent)", + "desc": "Target variable placeholder for compatibility.", + "types": ["text", "int"], + "stattypes": ["nominal"], + "required": false, + "multiple": false + }, + "x": { + "label": "Covariates (independent)", + "desc": "Feature variables placeholder for compatibility.", + "types": ["real", "int", "text"], + "stattypes": ["numerical", "nominal"], + "required": false, + "multiple": true + }, + "validation": false + } +} diff --git a/exaflow/algorithms/flower/flowertune_llm_medical/README.md b/exaflow/algorithms/flower/flowertune_llm_medical/README.md new file mode 100644 index 000000000..68edf5ce1 --- /dev/null +++ b/exaflow/algorithms/flower/flowertune_llm_medical/README.md @@ -0,0 +1,32 @@ +# flowertune_llm_medical (Phase 1) + +This package provides simulation-first Flower runtime wiring for Exaflow. + +Current implementation scope: +- strict RunConfig validation + env mapping contract +- lightweight federated training runtime (tiny adapter model for smoke runs) +- backend switch in contract: `model.backend = tiny | hf_peft` (`tiny` default) +- dataset adapter from worker CSVs when `inputdata.x`/`inputdata.y` are provided +- synthetic fallback partition only when no CSV dataset input is provided +- FedAvg strategy with checkpoint cadence and final summary reporting +- final summary payload validation against MetricsOut policy before posting +- checkpoint semantics: `final_checkpoint_ref` is always a filesystem path in completed runs + +Backend behavior: +- `tiny`: full Phase 1 smoke runtime. +- `hf_peft`: HF model/tokenizer load + LoRA attach (PEFT), adapter-only local training/evaluation, adapter tensor exchange. + - Requires: `torch`, `transformers`, `peft` (preflight-validated). + +Play mode (local simulation with prints): +- Run from project root: + `python -m exaflow.algorithms.flower.flowertune_llm_medical.main --backend tiny --rounds 2 --clients 2 --local-steps 1` +- You will see per-client and per-round metrics in stdout, then final metrics and `final_checkpoint_ref`. + +Execution scope policy: +- Phase 1 supports a single active Flower job at a time (`algorithm_execution_lock`). +- Flower controller endpoints are request-scoped via `request_id` for this algorithm. + +Planned next scope: +- replace tiny adapter runtime with full LLM/PEFT (LoRA) runtime +- round-level MetricsOut event emission +- deployment runtime support (SuperLink/SuperNodes) diff --git a/exaflow/algorithms/flower/flowertune_llm_medical/__init__.py b/exaflow/algorithms/flower/flowertune_llm_medical/__init__.py new file mode 100644 index 000000000..9553b7eab --- /dev/null +++ b/exaflow/algorithms/flower/flowertune_llm_medical/__init__.py @@ -0,0 +1 @@ +"""Flowertune LLM medical algorithm package (Phase 1 skeleton).""" diff --git a/exaflow/algorithms/flower/flowertune_llm_medical/client.py b/exaflow/algorithms/flower/flowertune_llm_medical/client.py new file mode 100644 index 000000000..cc0693abb --- /dev/null +++ b/exaflow/algorithms/flower/flowertune_llm_medical/client.py @@ -0,0 +1,6 @@ +"""Process entrypoint for flowertune_llm_medical Flower client.""" + +from exaflow.algorithms.flower.flowertune_llm_medical.client_app import start_client_app + +if __name__ == "__main__": + start_client_app() diff --git a/exaflow/algorithms/flower/flowertune_llm_medical/client_app.py b/exaflow/algorithms/flower/flowertune_llm_medical/client_app.py new file mode 100644 index 000000000..6e849ea45 --- /dev/null +++ b/exaflow/algorithms/flower/flowertune_llm_medical/client_app.py @@ -0,0 +1,143 @@ +"""Flower ClientApp runtime for flowertune_llm_medical Phase 1.""" + +from __future__ import annotations + +import json +import os +import time +from math import log2 +from typing import Dict + +import flwr as fl +import numpy as np +from flwr.common import NDArrays + +from exaflow.algorithms.flower.flowertune_llm_medical.controller_io import get_inputdata +from exaflow.algorithms.flower.flowertune_llm_medical.controller_io import ( + get_parameters, +) +from exaflow.algorithms.flower.flowertune_llm_medical.controller_io import get_run_env +from exaflow.algorithms.flower.flowertune_llm_medical.dataset import load_partition +from exaflow.algorithms.flower.flowertune_llm_medical.dataset import load_text_partition +from exaflow.algorithms.flower.flowertune_llm_medical.models import create_backend_model +from exaflow.algorithms.flower.flowertune_llm_medical.run_config import parse_run_config + + +def _parse_env_json_list(name: str): + raw = os.getenv(name) + if not raw: + return [] + return json.loads(raw) + + +class FederatedClient(fl.client.NumPyClient): + def __init__( + self, model, train_data, eval_data, num_train, num_val, requested_metrics + ): + self.model = model + self.train_data = train_data + self.eval_data = eval_data + self.num_train = num_train + self.num_val = num_val + self.requested_metrics = requested_metrics + + def get_parameters(self, config): + _ = config + return self.model.get_parameters() + + def fit(self, parameters: NDArrays, config: Dict): + self.model.set_parameters(parameters) + local_steps = int(os.getenv("LOCAL_STEPS", "1")) + if isinstance(self.train_data, tuple): + train_loss = self.model.fit_round( + self.train_data[0], self.train_data[1], local_steps + ) + else: + train_loss = self.model.fit_round(self.train_data, local_steps) + updated = self.model.get_parameters() + return updated, self.num_train, {"train_loss": float(train_loss)} + + def evaluate(self, parameters: NDArrays, config: Dict): + _ = config + self.model.set_parameters(parameters) + if isinstance(self.eval_data, tuple): + metrics = self.model.evaluate_round(self.eval_data[0], self.eval_data[1]) + else: + metrics = self.model.evaluate_round(self.eval_data) + reported = {k: v for k, v in metrics.items() if k in self.requested_metrics} + loss = float(reported.get("loss", metrics.get("loss", 0.0))) + return loss, self.num_val, reported + + +def _connect_with_retries(client): + timeout = max(2, int(os.getenv("TIMEOUT", "30"))) + max_attempts = max(2, int(log2(timeout)) + 1) + attempts = 0 + while True: + try: + fl.client.start_client( + server_address=os.environ["SERVER_ADDRESS"], + client=client.to_client(), + ) + return + except Exception: + attempts += 1 + if attempts >= max_attempts: + raise + time.sleep(min(2**attempts, 10)) + + +def start_client_app() -> None: + parameters = get_parameters() + config = parse_run_config(parameters) + os.environ.update(get_run_env()) + + inputdata = get_inputdata() + if config.model.backend.value == "hf_peft": + train_texts, eval_texts = load_text_partition( + inputdata, + seed=config.seed, + val_split_ratio=config.dataset.val_split_ratio, + ) + n_features = 16 + train_data = train_texts + eval_data = eval_texts + num_train = len(train_texts) + num_val = len(eval_texts) + else: + x_train, y_train, x_val, y_val = load_partition( + inputdata, + seed=config.seed, + val_split_ratio=config.dataset.val_split_ratio, + ) + n_features = int(x_train.shape[1]) + train_data = (x_train, y_train) + eval_data = (x_val, y_val) + num_train = len(y_train) + num_val = len(y_val) + + model = create_backend_model( + backend=config.model.backend.value, + n_features=n_features, + learning_rate=config.optimizer.learning_rate, + model_name=config.model.model_name, + local_steps=config.local_training.local_steps or 1, + max_seq_length=config.model.max_seq_length, + lora_r=config.peft.r or 4, + lora_alpha=config.peft.alpha or 8, + lora_dropout=config.peft.dropout or 0.0, + target_modules=config.peft.target_modules or ["q_proj", "v_proj"], + ) + requested_metrics = [m.value for m in config.evaluation.metrics] + # Parse env field to guarantee serialization consistency path is exercised. + _ = _parse_env_json_list("EVAL_METRICS") + + client = FederatedClient( + model=model, + train_data=train_data, + eval_data=eval_data, + num_train=num_train, + num_val=num_val, + requested_metrics=requested_metrics, + ) + _connect_with_retries(client) diff --git a/exaflow/algorithms/flower/flowertune_llm_medical/controller_io.py b/exaflow/algorithms/flower/flowertune_llm_medical/controller_io.py new file mode 100644 index 000000000..e467776f5 --- /dev/null +++ b/exaflow/algorithms/flower/flowertune_llm_medical/controller_io.py @@ -0,0 +1,60 @@ +"""Controller I/O helpers for flowertune_llm_medical.""" + +from __future__ import annotations + +import os +from typing import Any +from typing import Dict + +import requests + +CONTROLLER_IP = os.getenv("CONTROLLER_IP", "127.0.0.1") +CONTROLLER_PORT = os.getenv("CONTROLLER_PORT", "5000") +BASE_URL = f"http://{CONTROLLER_IP}:{CONTROLLER_PORT}" +INPUT_URL = f"{BASE_URL}/flower/input" +PARAMETERS_URL = f"{BASE_URL}/flower/parameters" +RUN_ENV_URL = f"{BASE_URL}/flower/run_env" +EVENT_URL = f"{BASE_URL}/flower/event" +RESULT_URL = f"{BASE_URL}/flower/result" +HEADERS = {"Content-type": "application/json", "Accept": "application/json"} + + +def _get_json(url: str) -> Dict[str, Any]: + response = requests.get(url, timeout=30) + response.raise_for_status() + return response.json() + + +def get_inputdata() -> Dict[str, Any]: + request_id = os.getenv("REQUEST_ID") + url = INPUT_URL if not request_id else f"{INPUT_URL}?request_id={request_id}" + return _get_json(url) + + +def get_parameters() -> Dict[str, Any]: + request_id = os.getenv("REQUEST_ID") + url = ( + PARAMETERS_URL + if not request_id + else f"{PARAMETERS_URL}?request_id={request_id}" + ) + return _get_json(url) + + +def get_run_env() -> Dict[str, str]: + request_id = os.getenv("REQUEST_ID") + url = RUN_ENV_URL if not request_id else f"{RUN_ENV_URL}?request_id={request_id}" + data = _get_json(url) + return {str(k): str(v) for k, v in data.items()} + + +def post_result(result: Dict[str, Any]) -> None: + response = requests.post(RESULT_URL, json=result, headers=HEADERS, timeout=30) + response.raise_for_status() + + +def post_event(event: Dict[str, Any]) -> None: + request_id = os.getenv("REQUEST_ID") + url = EVENT_URL if not request_id else f"{EVENT_URL}?request_id={request_id}" + response = requests.post(url, json=event, headers=HEADERS, timeout=30) + response.raise_for_status() diff --git a/exaflow/algorithms/flower/flowertune_llm_medical/dataset.py b/exaflow/algorithms/flower/flowertune_llm_medical/dataset.py new file mode 100644 index 000000000..bc6ee9142 --- /dev/null +++ b/exaflow/algorithms/flower/flowertune_llm_medical/dataset.py @@ -0,0 +1,162 @@ +"""Dataset adapter for flowertune_llm_medical Phase 1 runtime.""" + +from __future__ import annotations + +import os +from typing import Dict +from typing import List +from typing import Tuple + +import numpy as np +import pandas as pd + + +class DatasetLoadError(RuntimeError): + """Raised when explicit dataset loading fails.""" + + +def _synthetic_partition(seed: int, size: int = 256, n_features: int = 16): + rng = np.random.default_rng(seed) + x = rng.normal(size=(size, n_features)).astype(np.float32) + true_w = rng.normal(size=(n_features, 1)).astype(np.float32) + logits = x @ true_w + 0.1 * rng.normal(size=(size, 1)).astype(np.float32) + y = (logits[:, 0] > 0).astype(np.float32) + return x, y + + +def _split_train_val(x: np.ndarray, y: np.ndarray, val_ratio: float): + n = x.shape[0] + if n < 2: + return x, y, x, y + val_size = max(1, int(n * val_ratio)) + if val_size >= n: + val_size = max(1, n - 1) + split = n - val_size + return x[:split], y[:split], x[split:], y[split:] + + +def _split_text_train_val( + texts: List[str], val_ratio: float +) -> Tuple[List[str], List[str]]: + n = len(texts) + if n < 2: + return texts, texts + val_size = max(1, int(n * val_ratio)) + if val_size >= n: + val_size = max(1, n - 1) + split = n - val_size + return texts[:split], texts[split:] + + +def _to_numeric_frame(df: pd.DataFrame) -> np.ndarray: + # Convert mixed types to numeric matrix quickly for tiny smoke training. + for col in df.columns: + if df[col].dtype == "object": + df[col] = df[col].astype("category").cat.codes + return df.astype(float).to_numpy(dtype=np.float32) + + +def _row_to_text(row: pd.Series, x_vars: List[str], y_var: str) -> str: + feature_pairs = [f"{name}={row[name]}" for name in x_vars] + return f"Patient data: {', '.join(feature_pairs)}; target={row[y_var]}" + + +def load_partition( + inputdata: Dict, + *, + seed: int, + val_split_ratio: float, +) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Load a local partition from worker CSVs or fallback to synthetic data.""" + + csv_paths = [p for p in os.getenv("CSV_PATHS", "").split(",") if p] + x_vars = inputdata.get("x") or [] + y_vars = inputdata.get("y") or [] + + if csv_paths: + if not x_vars or not y_vars: + raise DatasetLoadError( + "CSV_PATHS provided but inputdata.x/y are missing; cannot load dataset." + ) + try: + frames = [pd.read_csv(path) for path in csv_paths] + full_df = pd.concat(frames, ignore_index=True) + + y_col = y_vars[0] + missing_cols = [ + col for col in [*x_vars, y_col] if col not in full_df.columns + ] + if missing_cols: + raise DatasetLoadError( + f"Dataset columns not found in CSV input: {missing_cols}" + ) + + features = _to_numeric_frame(full_df[x_vars].copy()) + target_raw = full_df[y_col] + if target_raw.dtype == "object": + target = target_raw.astype("category").cat.codes.to_numpy( + dtype=np.float32 + ) + else: + target = target_raw.to_numpy(dtype=np.float32) + target = (target > np.median(target)).astype(np.float32) + + return _split_train_val(features, target, val_split_ratio) + except DatasetLoadError: + raise + except Exception as exc: # noqa: BLE001 + raise DatasetLoadError(f"Failed to load CSV dataset: {exc}") from exc + + # Contract fallback for Phase 1 smoke runs. + x, y = _synthetic_partition(seed=seed) + return _split_train_val(x, y, val_split_ratio) + + +def load_text_partition( + inputdata: Dict, + *, + seed: int, + val_split_ratio: float, +) -> Tuple[List[str], List[str]]: + """Load local text partition for hf_peft backend.""" + + csv_paths = [p for p in os.getenv("CSV_PATHS", "").split(",") if p] + x_vars = inputdata.get("x") or [] + y_vars = inputdata.get("y") or [] + + if csv_paths: + if not x_vars or not y_vars: + raise DatasetLoadError( + "CSV_PATHS provided but inputdata.x/y are missing; cannot build text dataset." + ) + try: + frames = [pd.read_csv(path) for path in csv_paths] + full_df = pd.concat(frames, ignore_index=True) + y_col = y_vars[0] + missing_cols = [ + col for col in [*x_vars, y_col] if col not in full_df.columns + ] + if missing_cols: + raise DatasetLoadError( + f"Dataset columns not found in CSV input: {missing_cols}" + ) + texts = [ + _row_to_text(row, x_vars=x_vars, y_var=y_col) + for _, row in full_df.iterrows() + ] + return _split_text_train_val(texts, val_split_ratio) + except DatasetLoadError: + raise + except Exception as exc: # noqa: BLE001 + raise DatasetLoadError(f"Failed to load CSV text dataset: {exc}") from exc + + rng = np.random.default_rng(seed) + texts = [] + for i in range(256): + age = int(rng.integers(18, 90)) + bmi = float(rng.normal(27.0, 4.5)) + label = "high_risk" if (age > 60 and bmi > 28) else "low_risk" + texts.append( + f"Patient profile #{i}: age={age}, bmi={bmi:.2f}. predicted_outcome={label}" + ) + return _split_text_train_val(texts, val_split_ratio) diff --git a/exaflow/algorithms/flower/flowertune_llm_medical/main.py b/exaflow/algorithms/flower/flowertune_llm_medical/main.py new file mode 100644 index 000000000..e63b7f394 --- /dev/null +++ b/exaflow/algorithms/flower/flowertune_llm_medical/main.py @@ -0,0 +1,227 @@ +"""Play mode entrypoint for local federated training simulation. + +Run from IDE "Play" to observe round-by-round progress and final metrics +without starting the full Exaflow controller/worker stack. +""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path +from typing import List +from typing import Tuple + +import numpy as np + +# Allow running both as module and direct script file. +if __package__ in (None, ""): + sys.path.insert(0, str(Path(__file__).resolve().parents[4])) + +from exaflow.algorithms.flower.flowertune_llm_medical.dataset import load_partition +from exaflow.algorithms.flower.flowertune_llm_medical.models import ModelLoadError +from exaflow.algorithms.flower.flowertune_llm_medical.models import create_backend_model +from exaflow.algorithms.flower.flowertune_llm_medical.models import preflight_backend +from exaflow.algorithms.flower.flowertune_llm_medical.run_config import parse_run_config + + +def _build_run_config_from_args(args) -> dict: + return { + "schema_version": "1.1", + "runtime": "simulation", + "seed": args.seed, + "federation": { + "num_rounds": args.rounds, + "num_clients": args.clients, + "fraction_fit": 1.0, + "fraction_evaluate": 1.0, + "min_fit_clients": args.clients, + "min_evaluate_clients": args.clients, + "min_available_clients": args.clients, + }, + "model": { + "backend": args.backend, + "model_name": args.model_name, + "task_type": "causal_lm", + "max_seq_length": 128, + "quantization": "none", + }, + "peft": { + "enabled": True, + "method": "lora", + "r": 4, + "alpha": 8, + "dropout": 0.0, + "target_modules": ["q_proj"], + }, + "optimizer": { + "learning_rate": args.learning_rate, + "weight_decay": 0.0, + "max_grad_norm": 1.0, + }, + "local_training": { + "batch_size": 8, + "gradient_accumulation_steps": 1, + "local_steps": args.local_steps, + "fp16": False, + "bf16": True, + }, + "dataset": { + "dataset_name": "synthetic", + "split": "train", + "partitioner": "iid", + "num_partitions": args.clients, + "partition_id_strategy": "flower_simulation", + "val_split_ratio": args.val_split_ratio, + }, + "evaluation": { + "enabled": True, + "evaluate_every_n_rounds": 1, + "metrics": ["loss", "perplexity"], + }, + "artifacts": { + "artifact_dir": args.artifact_dir, + "checkpoint_every_n_rounds": 1, + "save_final_model": True, + "save_optimizer_state": False, + }, + } + + +def _weighted_average_parameters( + parameter_sets: List[Tuple[List[np.ndarray], int]], +) -> List[np.ndarray]: + total_examples = sum(num_examples for _, num_examples in parameter_sets) + if total_examples <= 0: + raise ValueError("Cannot aggregate parameters with zero examples.") + avg_params = None + for params, num_examples in parameter_sets: + weight = num_examples / total_examples + if avg_params is None: + avg_params = [weight * p for p in params] + else: + for i, p in enumerate(params): + avg_params[i] += weight * p + return avg_params + + +def _save_play_checkpoint(path: Path, parameters: List[np.ndarray]) -> Path: + path.parent.mkdir(parents=True, exist_ok=True) + np.savez(path, weights=parameters[0], bias=parameters[1]) + return path + + +def run_play_mode(args) -> int: + raw_cfg = _build_run_config_from_args(args) + config = parse_run_config(raw_cfg) + + try: + preflight_backend(config.model.backend.value) + except ModelLoadError as exc: + print(f"[ERROR] backend preflight failed: {exc}") + return 1 + + # Build synthetic local partitions for each simulated client. + client_partitions = [] + for cid in range(config.federation.num_clients): + x_train, y_train, x_val, y_val = load_partition( + {"x": [], "y": []}, + seed=config.seed + (cid * 101), + val_split_ratio=config.dataset.val_split_ratio, + ) + client_partitions.append((x_train, y_train, x_val, y_val)) + + n_features = client_partitions[0][0].shape[1] + template_model = create_backend_model( + backend=config.model.backend.value, + n_features=n_features, + learning_rate=config.optimizer.learning_rate, + model_name=config.model.model_name, + local_steps=config.local_training.local_steps or 1, + max_seq_length=config.model.max_seq_length, + lora_r=config.peft.r or 4, + lora_alpha=config.peft.alpha or 8, + lora_dropout=config.peft.dropout or 0.0, + target_modules=config.peft.target_modules or ["q_proj", "v_proj"], + ) + global_params = template_model.get_parameters() + + best_round = None + best_loss = None + final_metrics = {"loss": 0.0, "perplexity": 0.0} + + print( + f"[START] backend={config.model.backend.value} rounds={config.federation.num_rounds} clients={config.federation.num_clients}" + ) + + for rnd in range(1, config.federation.num_rounds + 1): + local_updates = [] + round_eval = [] + for cid, (x_train, y_train, x_val, y_val) in enumerate(client_partitions): + client_model = create_backend_model( + backend=config.model.backend.value, + n_features=n_features, + learning_rate=config.optimizer.learning_rate, + model_name=config.model.model_name, + local_steps=config.local_training.local_steps or 1, + max_seq_length=config.model.max_seq_length, + lora_r=config.peft.r or 4, + lora_alpha=config.peft.alpha or 8, + lora_dropout=config.peft.dropout or 0.0, + target_modules=config.peft.target_modules or ["q_proj", "v_proj"], + ) + client_model.set_parameters(global_params) + train_loss = client_model.fit_round( + x_train, y_train, config.local_training.local_steps + ) + local_updates.append((client_model.get_parameters(), len(y_train))) + metrics = client_model.evaluate_round(x_val, y_val) + round_eval.append((metrics, len(y_val))) + print( + f"[ROUND {rnd}] client={cid} train_loss={train_loss:.6f} val_loss={metrics['loss']:.6f}" + ) + + global_params = _weighted_average_parameters(local_updates) + + total_val = sum(n for _, n in round_eval) + agg_loss = sum(m["loss"] * n for m, n in round_eval) / max(1, total_val) + agg_ppl = sum(m["perplexity"] * n for m, n in round_eval) / max(1, total_val) + final_metrics = {"loss": float(agg_loss), "perplexity": float(agg_ppl)} + + if best_loss is None or agg_loss < best_loss: + best_loss = agg_loss + best_round = rnd + + print(f"[ROUND {rnd}] aggregate loss={agg_loss:.6f} perplexity={agg_ppl:.6f}") + + ckpt_path = _save_play_checkpoint( + Path(config.artifacts.artifact_dir) / "play_final_model.npz", global_params + ) + print("[DONE] Training completed") + print( + f"[RESULT] rounds={config.federation.num_rounds} best_round={best_round} best_loss={best_loss:.6f}" + ) + print( + f"[RESULT] final_loss={final_metrics['loss']:.6f} final_perplexity={final_metrics['perplexity']:.6f}" + ) + print(f"[RESULT] final_checkpoint_ref={ckpt_path}") + return 0 + + +def main(): + parser = argparse.ArgumentParser(description="Play mode training runner.") + parser.add_argument("--backend", choices=["tiny", "hf_peft"], default="tiny") + parser.add_argument("--model-name", default="tiny") + parser.add_argument("--rounds", type=int, default=2) + parser.add_argument("--clients", type=int, default=2) + parser.add_argument("--local-steps", type=int, default=1) + parser.add_argument("--learning-rate", type=float, default=0.01) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--val-split-ratio", type=float, default=0.2) + parser.add_argument("--artifact-dir", default="runs/play-mode") + args = parser.parse_args() + raise SystemExit(run_play_mode(args)) + + +if __name__ == "__main__": + main() diff --git a/exaflow/algorithms/flower/flowertune_llm_medical/metrics_out.py b/exaflow/algorithms/flower/flowertune_llm_medical/metrics_out.py new file mode 100644 index 000000000..20e84e88c --- /dev/null +++ b/exaflow/algorithms/flower/flowertune_llm_medical/metrics_out.py @@ -0,0 +1,166 @@ +"""MetricsOut validation for flowertune_llm_medical final summary payloads.""" + +from __future__ import annotations + +from datetime import datetime +from typing import Any +from typing import Dict +from typing import List + +from pydantic import BaseModel +from pydantic import Field +from pydantic import ValidationError + + +class _ErrorInfo(BaseModel): + code: str + message: str + details: Dict[str, Any] | None = None + + +class _Artifacts(BaseModel): + final_checkpoint_ref: str = Field(..., min_length=1) + intermediate_checkpoints: List[str] | None = None + + +class _RoundArtifacts(BaseModel): + checkpoint_ref: str | None = None + + +class _RoundClients(BaseModel): + participated: int = Field(..., ge=0) + expected: int = Field(..., ge=1) + failed: int = Field(0, ge=0) + + +class _RoundEvent(BaseModel): + schema_version: str + payload_type: str + job_id: str + algorithm_name: str + runtime: str + status: str + round: int = Field(..., ge=1) + rounds_total: int = Field(..., ge=1) + timestamp: str + requested_metrics: List[str] + reported_metrics: List[str] + metrics: Dict[str, float] + clients: _RoundClients + artifacts: _RoundArtifacts + error: _ErrorInfo | None = None + + +class _FinalSummary(BaseModel): + schema_version: str + payload_type: str + job_id: str + algorithm_name: str + runtime: str + status: str + timestamp: str + rounds_total: int + rounds_completed: int + requested_metrics: List[str] + reported_metrics: List[str] + aggregate_metrics: Dict[str, float] + artifacts: _Artifacts | None = None + error: _ErrorInfo | None = None + + +def _is_utc_z(ts: str) -> bool: + if not ts.endswith("Z"): + return False + try: + datetime.fromisoformat(ts.replace("Z", "+00:00")) + return True + except ValueError: + return False + + +def validate_final_summary_payload(payload: Dict[str, Any]) -> None: + """Validate final summary payload against Phase 1.1 policy.""" + + try: + parsed = _FinalSummary.parse_obj(payload) + except ValidationError as exc: + raise ValueError(f"Invalid final_summary payload shape: {exc}") from exc + + if parsed.schema_version != "1.1": + raise ValueError("final_summary.schema_version must be '1.1'") + if parsed.payload_type != "final_summary": + raise ValueError("payload_type must be 'final_summary'") + if parsed.runtime != "simulation": + raise ValueError("runtime must be 'simulation'") + if parsed.status not in {"COMPLETED", "FAILED", "CANCELLED", "TIMEOUT"}: + raise ValueError("invalid final_summary status") + if not _is_utc_z(parsed.timestamp): + raise ValueError("timestamp must be UTC and end with 'Z'") + + if parsed.rounds_completed > parsed.rounds_total: + raise ValueError("rounds_completed must be <= rounds_total") + + metric_keys = sorted(parsed.aggregate_metrics.keys()) + if parsed.status == "COMPLETED": + if parsed.error is not None: + raise ValueError("error must be omitted for COMPLETED status") + if parsed.artifacts is None: + raise ValueError("artifacts are required for COMPLETED status") + if sorted(parsed.reported_metrics) != metric_keys: + raise ValueError( + "reported_metrics must equal aggregate_metrics keys for COMPLETED status" + ) + else: + if parsed.error is None: + raise ValueError("error is required for non-COMPLETED status") + if parsed.aggregate_metrics and not set(parsed.reported_metrics).issubset( + parsed.aggregate_metrics.keys() + ): + raise ValueError( + "reported_metrics must be subset of aggregate_metrics keys for non-COMPLETED status" + ) + + +def validate_round_event_payload(payload: Dict[str, Any]) -> None: + """Validate round event payload against Phase 1.1 policy.""" + + try: + parsed = _RoundEvent.parse_obj(payload) + except ValidationError as exc: + raise ValueError(f"Invalid round_event payload shape: {exc}") from exc + + if parsed.schema_version != "1.1": + raise ValueError("round_event.schema_version must be '1.1'") + if parsed.payload_type != "round_event": + raise ValueError("payload_type must be 'round_event'") + if parsed.runtime != "simulation": + raise ValueError("runtime must be 'simulation'") + if parsed.status not in {"RUNNING", "FAILED"}: + raise ValueError("round_event.status must be RUNNING or FAILED") + if not _is_utc_z(parsed.timestamp): + raise ValueError("timestamp must be UTC and end with 'Z'") + if parsed.round > parsed.rounds_total: + raise ValueError("round must be <= rounds_total") + if parsed.clients.participated > parsed.clients.expected: + raise ValueError("clients.participated must be <= clients.expected") + if parsed.clients.failed > parsed.clients.expected: + raise ValueError("clients.failed must be <= clients.expected") + if (parsed.clients.participated + parsed.clients.failed) > parsed.clients.expected: + raise ValueError( + "clients.participated + clients.failed must be <= clients.expected" + ) + + if parsed.status == "RUNNING": + if parsed.error is not None: + raise ValueError("error must be omitted for RUNNING round_event") + if sorted(parsed.reported_metrics) != sorted(parsed.metrics.keys()): + raise ValueError("reported_metrics must equal metrics keys for RUNNING") + else: + if parsed.error is None: + raise ValueError("error is required for FAILED round_event") + if parsed.metrics and not set(parsed.reported_metrics).issubset( + parsed.metrics.keys() + ): + raise ValueError( + "reported_metrics must be subset of metrics keys for FAILED round_event" + ) diff --git a/exaflow/algorithms/flower/flowertune_llm_medical/models.py b/exaflow/algorithms/flower/flowertune_llm_medical/models.py new file mode 100644 index 000000000..4be6f5d5a --- /dev/null +++ b/exaflow/algorithms/flower/flowertune_llm_medical/models.py @@ -0,0 +1,260 @@ +"""Model backend implementations for flowertune_llm_medical.""" + +from __future__ import annotations + +import importlib +import importlib.util +from dataclasses import dataclass +from typing import Dict +from typing import List + +import numpy as np + + +class ModelLoadError(RuntimeError): + """Raised when requested model backend cannot be initialized.""" + + +@dataclass +class TinyAdapterModel: + """A small logistic adapter model to exercise federated runtime wiring.""" + + n_features: int + learning_rate: float + + def __post_init__(self): + self.weights = np.zeros((self.n_features,), dtype=np.float32) + self.bias = np.zeros((1,), dtype=np.float32) + + def get_parameters(self): + return [self.weights.copy(), self.bias.copy()] + + def set_parameters(self, params): + self.weights = np.asarray(params[0], dtype=np.float32).reshape(self.n_features) + self.bias = np.asarray(params[1], dtype=np.float32).reshape(1) + + def _predict_proba(self, x: np.ndarray) -> np.ndarray: + logits = x @ self.weights + self.bias[0] + logits = np.clip(logits, -30.0, 30.0) + return 1.0 / (1.0 + np.exp(-logits)) + + @staticmethod + def _binary_loss(y_true: np.ndarray, y_prob: np.ndarray) -> float: + eps = 1e-7 + y_prob = np.clip(y_prob, eps, 1 - eps) + return float( + -np.mean(y_true * np.log(y_prob) + (1 - y_true) * np.log(1 - y_prob)) + ) + + def fit_round(self, x: np.ndarray, y: np.ndarray, local_steps: int) -> float: + n = max(1, x.shape[0]) + for _ in range(local_steps): + probs = self._predict_proba(x) + err = probs - y + grad_w = (x.T @ err) / n + grad_b = np.mean(err) + self.weights -= self.learning_rate * grad_w.astype(np.float32) + self.bias[0] -= self.learning_rate * float(grad_b) + final_prob = self._predict_proba(x) + return self._binary_loss(y, final_prob) + + def evaluate_round(self, x: np.ndarray, y: np.ndarray): + prob = self._predict_proba(x) + loss = self._binary_loss(y, prob) + perplexity = float(np.exp(min(20.0, loss))) + return {"loss": loss, "perplexity": perplexity} + + +class HFPeftAdapterModel: + """HF+PEFT backend exchanging only adapter tensors.""" + + def __init__( + self, + *, + model_name: str, + learning_rate: float, + local_steps: int, + max_seq_length: int, + lora_r: int, + lora_alpha: int, + lora_dropout: float, + target_modules: List[str], + ): + self.model_name = model_name + self.learning_rate = learning_rate + self.local_steps = local_steps + self.max_seq_length = max_seq_length + + self.torch = importlib.import_module("torch") + transformers = importlib.import_module("transformers") + peft = importlib.import_module("peft") + + self.AutoTokenizer = transformers.AutoTokenizer + self.AutoModelForCausalLM = transformers.AutoModelForCausalLM + self.LoraConfig = peft.LoraConfig + self.get_peft_model = peft.get_peft_model + self.get_peft_model_state_dict = peft.get_peft_model_state_dict + self.set_peft_model_state_dict = peft.set_peft_model_state_dict + + self.tokenizer = self.AutoTokenizer.from_pretrained(self.model_name) + if self.tokenizer.pad_token is None: + self.tokenizer.pad_token = self.tokenizer.eos_token + + base_model = self.AutoModelForCausalLM.from_pretrained(self.model_name) + peft_config = self.LoraConfig( + r=lora_r, + lora_alpha=lora_alpha, + lora_dropout=lora_dropout, + bias="none", + task_type="CAUSAL_LM", + target_modules=target_modules, + ) + self.model = self.get_peft_model(base_model, peft_config) + self.model.train() + self.optimizer = self.torch.optim.AdamW( + self.model.parameters(), lr=self.learning_rate + ) + + adapter_state = self.get_peft_model_state_dict(self.model) + self.adapter_keys = sorted(adapter_state.keys()) + + def _tokenize(self, texts: List[str]): + tokens = self.tokenizer( + texts, + padding=True, + truncation=True, + max_length=self.max_seq_length, + return_tensors="pt", + ) + tokens["labels"] = tokens["input_ids"].clone() + return tokens + + def get_parameters(self): + state = self.get_peft_model_state_dict(self.model) + params = [] + for key in self.adapter_keys: + params.append(state[key].detach().cpu().numpy().astype(np.float32)) + return params + + def set_parameters(self, params): + state = self.get_peft_model_state_dict(self.model) + for key, arr in zip(self.adapter_keys, params): + tensor = self.torch.tensor(arr) + tensor = tensor.to(dtype=state[key].dtype) + state[key] = tensor + self.set_peft_model_state_dict(self.model, state) + + def fit_round(self, train_texts: List[str], local_steps: int) -> float: + if not train_texts: + return 0.0 + self.model.train() + losses = [] + batch_size = min(4, max(1, len(train_texts))) + for step in range(local_steps): + start = (step * batch_size) % len(train_texts) + batch = train_texts[start : start + batch_size] + if not batch: + batch = train_texts[:batch_size] + inputs = self._tokenize(batch) + outputs = self.model(**inputs) + loss = outputs.loss + self.optimizer.zero_grad() + loss.backward() + self.optimizer.step() + losses.append(float(loss.detach().cpu().item())) + return float(np.mean(losses)) if losses else 0.0 + + def evaluate_round(self, eval_texts: List[str]): + if not eval_texts: + return {"loss": 0.0, "perplexity": 1.0} + self.model.eval() + with self.torch.no_grad(): + batch = eval_texts[: min(8, len(eval_texts))] + inputs = self._tokenize(batch) + outputs = self.model(**inputs) + loss = float(outputs.loss.detach().cpu().item()) + perplexity = float(np.exp(min(20.0, loss))) + return {"loss": loss, "perplexity": perplexity} + + +def create_model(n_features: int, learning_rate: float) -> TinyAdapterModel: + return TinyAdapterModel(n_features=n_features, learning_rate=learning_rate) + + +def _ensure_hf_peft_dependencies(): + missing = [ + package + for package in ("torch", "transformers", "peft") + if importlib.util.find_spec(package) is None + ] + if missing: + raise ModelLoadError( + "hf_peft backend requires missing dependencies: " + + ", ".join(sorted(missing)) + ) + + +def preflight_backend(backend: str) -> None: + """Validate backend dependencies before runtime starts.""" + if backend == "hf_peft": + _ensure_hf_peft_dependencies() + elif backend != "tiny": + raise ModelLoadError(f"Unsupported model backend: {backend}") + + +def create_backend_model( + *, + backend: str, + n_features: int, + learning_rate: float, + model_name: str, + local_steps: int, + max_seq_length: int, + lora_r: int, + lora_alpha: int, + lora_dropout: float, + target_modules: List[str], +): + if backend == "tiny": + return create_model(n_features=n_features, learning_rate=learning_rate) + if backend == "hf_peft": + _ensure_hf_peft_dependencies() + return HFPeftAdapterModel( + model_name=model_name, + learning_rate=learning_rate, + local_steps=local_steps, + max_seq_length=max_seq_length, + lora_r=lora_r, + lora_alpha=lora_alpha, + lora_dropout=lora_dropout, + target_modules=target_modules, + ) + raise ModelLoadError(f"Unsupported model backend: {backend}") + + +def initial_parameters_for_backend( + *, + backend: str, + n_features: int, + learning_rate: float, + model_name: str, + local_steps: int, + max_seq_length: int, + lora_r: int, + lora_alpha: int, + lora_dropout: float, + target_modules: List[str], +): + model = create_backend_model( + backend=backend, + n_features=n_features, + learning_rate=learning_rate, + model_name=model_name, + local_steps=local_steps, + max_seq_length=max_seq_length, + lora_r=lora_r, + lora_alpha=lora_alpha, + lora_dropout=lora_dropout, + target_modules=target_modules, + ) + return model.get_parameters() diff --git a/exaflow/algorithms/flower/flowertune_llm_medical/run_config.py b/exaflow/algorithms/flower/flowertune_llm_medical/run_config.py new file mode 100644 index 000000000..becd50795 --- /dev/null +++ b/exaflow/algorithms/flower/flowertune_llm_medical/run_config.py @@ -0,0 +1,353 @@ +"""RunConfig validation and env mapping for flowertune_llm_medical.""" + +from __future__ import annotations + +import json +from enum import Enum +from typing import Any +from typing import Dict +from typing import List +from typing import Optional + +from pydantic import BaseModel +from pydantic import Field +from pydantic import root_validator +from pydantic import validator + + +class StrictBaseModel(BaseModel): + """Strict model base used for run configuration validation.""" + + class Config: + extra = "forbid" + anystr_strip_whitespace = True + validate_assignment = True + + +class Runtime(str, Enum): + SIMULATION = "simulation" + + +class Quantization(str, Enum): + NONE = "none" + BIT_8 = "8bit" + BIT_4 = "4bit" + + +class PeftMethod(str, Enum): + LORA = "lora" + + +class Partitioner(str, Enum): + IID = "iid" + + +class PartitionIdStrategy(str, Enum): + FLOWER_SIMULATION = "flower_simulation" + + +class LogLevel(str, Enum): + DEBUG = "DEBUG" + INFO = "INFO" + WARNING = "WARNING" + ERROR = "ERROR" + + +class TaskType(str, Enum): + CAUSAL_LM = "causal_lm" + + +class ModelBackend(str, Enum): + TINY = "tiny" + HF_PEFT = "hf_peft" + + +class EvalMetric(str, Enum): + LOSS = "loss" + PERPLEXITY = "perplexity" + + +class FederationConfig(StrictBaseModel): + num_rounds: int = Field(..., ge=1) + num_clients: int = Field(..., ge=1) + fraction_fit: float = Field(1.0, gt=0.0, le=1.0) + fraction_evaluate: float = Field(1.0, ge=0.0, le=1.0) + min_fit_clients: int = Field(..., ge=1) + min_evaluate_clients: int = Field(..., ge=0) + min_available_clients: int = Field(..., ge=1) + + @root_validator(skip_on_failure=True) + def _validate_client_counts(cls, values): + total_clients = values["num_clients"] + for key in ("min_fit_clients", "min_evaluate_clients", "min_available_clients"): + if values[key] > total_clients: + raise ValueError(f"{key} must be <= num_clients") + return values + + +class ModelConfig(StrictBaseModel): + backend: ModelBackend = ModelBackend.TINY + model_name: str = Field(..., min_length=1) + task_type: TaskType = TaskType.CAUSAL_LM + max_seq_length: int = Field(512, ge=16) + quantization: Quantization = Quantization.BIT_4 + + +class PeftConfig(StrictBaseModel): + enabled: bool + method: Optional[PeftMethod] = None + r: Optional[int] = Field(None, ge=1) + alpha: Optional[int] = Field(None, ge=1) + dropout: Optional[float] = Field(None, ge=0.0, le=1.0) + target_modules: Optional[List[str]] = None + + @root_validator(skip_on_failure=True) + def _validate_enabled_fields(cls, values): + if not values["enabled"]: + return values + required = ("method", "r", "alpha") + missing = [name for name in required if values.get(name) is None] + if missing: + raise ValueError( + f"peft enabled=true requires these fields: {', '.join(missing)}" + ) + return values + + @validator("target_modules") + def _validate_target_modules(cls, value): + if value is None: + return value + if not value: + raise ValueError("peft.target_modules must not be empty when provided") + return value + + +class OptimizerConfig(StrictBaseModel): + learning_rate: float = Field(2e-4, gt=0.0) + weight_decay: float = Field(0.0, ge=0.0) + max_grad_norm: float = Field(1.0, gt=0.0) + + +class LocalTrainingConfig(StrictBaseModel): + batch_size: int = Field(..., ge=1) + gradient_accumulation_steps: int = Field(1, ge=1) + local_steps: Optional[int] = Field(None, ge=1) + epochs: Optional[float] = Field(None, gt=0.0) + fp16: bool = False + bf16: bool = True + + @root_validator(skip_on_failure=True) + def _validate_steps_xor_epochs(cls, values): + has_steps = values.get("local_steps") is not None + has_epochs = values.get("epochs") is not None + if has_steps == has_epochs: + raise ValueError("exactly one of local_steps or epochs must be set") + if values.get("fp16") and values.get("bf16"): + raise ValueError("fp16 and bf16 cannot both be true") + return values + + +class DatasetConfig(StrictBaseModel): + dataset_name: str = Field(..., min_length=1) + dataset_config: Optional[str] = None + split: str = "train" + partitioner: Partitioner = Partitioner.IID + num_partitions: int = Field(..., ge=1) + partition_id_strategy: PartitionIdStrategy = PartitionIdStrategy.FLOWER_SIMULATION + val_split_ratio: float = Field(0.1, ge=0.0, lt=1.0) + + +class EvaluationConfig(StrictBaseModel): + enabled: bool = True + evaluate_every_n_rounds: int = Field(1, ge=1) + metrics: List[EvalMetric] = Field( + default_factory=lambda: [EvalMetric.LOSS, EvalMetric.PERPLEXITY] + ) + + +class ArtifactsConfig(StrictBaseModel): + artifact_dir: str = Field(..., min_length=1) + checkpoint_every_n_rounds: int = Field(1, ge=1) + save_final_model: bool = True + save_optimizer_state: bool = False + + +class LoggingConfig(StrictBaseModel): + log_level: LogLevel = LogLevel.INFO + report_round_metrics: bool = True + report_client_metrics: bool = False + + +class TimeoutsConfig(StrictBaseModel): + round_timeout_sec: int = Field(1800, ge=1) + job_timeout_sec: int = Field(14400, ge=1) + + +class RunConfigContract(StrictBaseModel): + schema_version: str = "1.1" + runtime: Runtime = Runtime.SIMULATION + seed: int = Field(42, ge=0) + federation: FederationConfig + model: ModelConfig + peft: PeftConfig + optimizer: OptimizerConfig = Field(default_factory=OptimizerConfig) + local_training: LocalTrainingConfig + dataset: DatasetConfig + evaluation: EvaluationConfig = Field(default_factory=EvaluationConfig) + artifacts: ArtifactsConfig + logging: LoggingConfig = Field(default_factory=LoggingConfig) + timeouts: TimeoutsConfig = Field(default_factory=TimeoutsConfig) + + @validator("schema_version") + def _validate_schema_version(cls, value): + if value != "1.1": + raise ValueError("schema_version must be '1.1'") + return value + + @root_validator(skip_on_failure=True) + def _validate_cross_section_constraints(cls, values): + federation = values["federation"] + dataset = values["dataset"] + evaluation = values["evaluation"] + timeouts = values["timeouts"] + artifacts = values["artifacts"] + + if dataset.partitioner == Partitioner.IID and ( + dataset.num_partitions != federation.num_clients + ): + raise ValueError( + "for iid partitioner, dataset.num_partitions must equal federation.num_clients" + ) + if not evaluation.enabled: + if federation.fraction_evaluate != 0: + raise ValueError( + "federation.fraction_evaluate must be 0 when evaluation.enabled=false" + ) + if federation.min_evaluate_clients != 0: + raise ValueError( + "federation.min_evaluate_clients must be 0 when evaluation.enabled=false" + ) + if timeouts.job_timeout_sec < timeouts.round_timeout_sec: + raise ValueError("timeouts.job_timeout_sec must be >= round_timeout_sec") + if artifacts.checkpoint_every_n_rounds > federation.num_rounds: + raise ValueError( + "artifacts.checkpoint_every_n_rounds must be <= federation.num_rounds" + ) + return values + + +def parse_run_config(parameters: Dict[str, Any]) -> RunConfigContract: + """Parse and validate raw parameters into a strict run config model.""" + + return RunConfigContract.parse_obj(parameters or {}) + + +def serialize_run_config(config: RunConfigContract) -> Dict[str, Any]: + """Return canonical, serializable run config dictionary.""" + + return config.dict(exclude_none=True) + + +def build_env_mapping(config: RunConfigContract, request_id: str) -> Dict[str, str]: + """Build canonical env var mapping from a validated run configuration.""" + + def _bool(v: bool) -> str: + return "true" if v else "false" + + def _set( + env: Dict[str, str], + key: str, + value: Any, + *, + as_bool: bool = False, + as_json: bool = False, + ) -> None: + if value is None: + return + if as_bool: + env[key] = _bool(bool(value)) + return + if as_json: + env[key] = json.dumps(value) + return + env[key] = str(value) + + resolved_artifact_dir = config.artifacts.artifact_dir.replace( + "${request_id}", request_id + ) + env: Dict[str, str] = {} + _set(env, "RUN_CONFIG_SCHEMA_VERSION", config.schema_version) + _set(env, "RUN_RUNTIME", config.runtime.value) + _set(env, "SEED", config.seed) + + _set(env, "NUM_ROUNDS", config.federation.num_rounds) + _set(env, "NUM_CLIENTS", config.federation.num_clients) + _set(env, "FRACTION_FIT", config.federation.fraction_fit) + _set(env, "FRACTION_EVALUATE", config.federation.fraction_evaluate) + _set(env, "MIN_FIT_CLIENTS", config.federation.min_fit_clients) + _set(env, "MIN_EVALUATE_CLIENTS", config.federation.min_evaluate_clients) + _set(env, "MIN_AVAILABLE_CLIENTS", config.federation.min_available_clients) + + _set(env, "MODEL_NAME", config.model.model_name) + _set(env, "MODEL_BACKEND", config.model.backend.value) + _set(env, "TASK_TYPE", config.model.task_type.value) + _set(env, "MAX_SEQ_LENGTH", config.model.max_seq_length) + _set(env, "QUANTIZATION", config.model.quantization.value) + + _set(env, "PEFT_ENABLED", config.peft.enabled, as_bool=True) + _set(env, "PEFT_METHOD", config.peft.method.value if config.peft.method else None) + _set(env, "LORA_R", config.peft.r) + _set(env, "LORA_ALPHA", config.peft.alpha) + _set(env, "LORA_DROPOUT", config.peft.dropout) + _set(env, "LORA_TARGET_MODULES", config.peft.target_modules, as_json=True) + + _set(env, "LEARNING_RATE", config.optimizer.learning_rate) + _set(env, "WEIGHT_DECAY", config.optimizer.weight_decay) + _set(env, "MAX_GRAD_NORM", config.optimizer.max_grad_norm) + + _set(env, "BATCH_SIZE", config.local_training.batch_size) + _set(env, "GRAD_ACC_STEPS", config.local_training.gradient_accumulation_steps) + _set(env, "LOCAL_STEPS", config.local_training.local_steps) + _set(env, "LOCAL_EPOCHS", config.local_training.epochs) + _set(env, "FP16", config.local_training.fp16, as_bool=True) + _set(env, "BF16", config.local_training.bf16, as_bool=True) + + _set(env, "DATASET_NAME", config.dataset.dataset_name) + _set(env, "DATASET_CONFIG", config.dataset.dataset_config) + _set(env, "DATASET_SPLIT", config.dataset.split) + _set(env, "PARTITIONER", config.dataset.partitioner.value) + _set(env, "NUM_PARTITIONS", config.dataset.num_partitions) + _set(env, "PARTITION_ID_STRATEGY", config.dataset.partition_id_strategy.value) + _set(env, "VAL_SPLIT_RATIO", config.dataset.val_split_ratio) + + _set(env, "EVAL_ENABLED", config.evaluation.enabled, as_bool=True) + _set(env, "EVAL_EVERY_N_ROUNDS", config.evaluation.evaluate_every_n_rounds) + _set( + env, + "EVAL_METRICS", + [metric.value for metric in config.evaluation.metrics], + as_json=True, + ) + + _set(env, "ARTIFACT_DIR", resolved_artifact_dir) + _set(env, "CKPT_EVERY_N_ROUNDS", config.artifacts.checkpoint_every_n_rounds) + _set(env, "SAVE_FINAL_MODEL", config.artifacts.save_final_model, as_bool=True) + _set(env, "SAVE_OPT_STATE", config.artifacts.save_optimizer_state, as_bool=True) + + _set(env, "FLOWERTUNE_LOG_LEVEL", config.logging.log_level.value) + _set( + env, + "REPORT_ROUND_METRICS", + config.logging.report_round_metrics, + as_bool=True, + ) + _set( + env, + "REPORT_CLIENT_METRICS", + config.logging.report_client_metrics, + as_bool=True, + ) + + _set(env, "ROUND_TIMEOUT_SEC", config.timeouts.round_timeout_sec) + _set(env, "JOB_TIMEOUT_SEC", config.timeouts.job_timeout_sec) + return env diff --git a/exaflow/algorithms/flower/flowertune_llm_medical/server.py b/exaflow/algorithms/flower/flowertune_llm_medical/server.py new file mode 100644 index 000000000..d35d7183e --- /dev/null +++ b/exaflow/algorithms/flower/flowertune_llm_medical/server.py @@ -0,0 +1,6 @@ +"""Process entrypoint for flowertune_llm_medical Flower server.""" + +from exaflow.algorithms.flower.flowertune_llm_medical.server_app import start_server_app + +if __name__ == "__main__": + start_server_app() diff --git a/exaflow/algorithms/flower/flowertune_llm_medical/server_app.py b/exaflow/algorithms/flower/flowertune_llm_medical/server_app.py new file mode 100644 index 000000000..889fd522e --- /dev/null +++ b/exaflow/algorithms/flower/flowertune_llm_medical/server_app.py @@ -0,0 +1,173 @@ +"""Flower ServerApp runtime for flowertune_llm_medical Phase 1.""" + +from __future__ import annotations + +import os + +import flwr as fl + +from exaflow.algorithms.flower.flowertune_llm_medical.controller_io import get_inputdata +from exaflow.algorithms.flower.flowertune_llm_medical.controller_io import ( + get_parameters, +) +from exaflow.algorithms.flower.flowertune_llm_medical.controller_io import get_run_env +from exaflow.algorithms.flower.flowertune_llm_medical.controller_io import post_result +from exaflow.algorithms.flower.flowertune_llm_medical.dataset import DatasetLoadError +from exaflow.algorithms.flower.flowertune_llm_medical.metrics_out import ( + validate_final_summary_payload, +) +from exaflow.algorithms.flower.flowertune_llm_medical.models import ModelLoadError +from exaflow.algorithms.flower.flowertune_llm_medical.models import ( + initial_parameters_for_backend, +) +from exaflow.algorithms.flower.flowertune_llm_medical.models import preflight_backend +from exaflow.algorithms.flower.flowertune_llm_medical.run_config import parse_run_config +from exaflow.algorithms.flower.flowertune_llm_medical.strategy import build_strategy +from exaflow.algorithms.flower.flowertune_llm_medical.strategy import ( + save_final_checkpoint, +) +from exaflow.algorithms.flower.flowertune_llm_medical.strategy import utc_now + + +def _job_id() -> str: + return os.getenv("REQUEST_ID", "unknown") + + +def _failure_summary(config, code: str, message: str, details: dict): + payload = { + "schema_version": "1.1", + "payload_type": "final_summary", + "job_id": _job_id(), + "request_id": _job_id(), + "algorithm_name": "flowertune_llm_medical", + "runtime": "simulation", + "status": "FAILED", + "timestamp": utc_now(), + "rounds_total": config.federation.num_rounds, + "rounds_completed": 0, + "requested_metrics": [m.value for m in config.evaluation.metrics], + "reported_metrics": [], + "aggregate_metrics": {}, + "error": {"code": code, "message": message, "details": details}, + } + validate_final_summary_payload(payload) + return payload + + +def start_server_app() -> None: + parameters = get_parameters() + config = parse_run_config(parameters) + os.environ.update(get_run_env()) + + try: + preflight_backend(config.model.backend.value) + requested_metrics = [metric.value for metric in config.evaluation.metrics] + artifact_dir = os.getenv("ARTIFACT_DIR", config.artifacts.artifact_dir) + + strategy, state = build_strategy( + num_rounds=config.federation.num_rounds, + fraction_fit=config.federation.fraction_fit, + fraction_evaluate=config.federation.fraction_evaluate, + min_fit_clients=config.federation.min_fit_clients, + min_evaluate_clients=config.federation.min_evaluate_clients, + min_available_clients=config.federation.min_available_clients, + checkpoint_every_n_rounds=config.artifacts.checkpoint_every_n_rounds, + artifact_dir=artifact_dir, + requested_metrics=requested_metrics, + ) + + inputdata = get_inputdata() + n_features = len(inputdata.get("x") or []) or 16 + init_arrays = initial_parameters_for_backend( + backend=config.model.backend.value, + n_features=n_features, + learning_rate=config.optimizer.learning_rate, + model_name=config.model.model_name, + local_steps=config.local_training.local_steps or 1, + max_seq_length=config.model.max_seq_length, + lora_r=config.peft.r or 4, + lora_alpha=config.peft.alpha or 8, + lora_dropout=config.peft.dropout or 0.0, + target_modules=config.peft.target_modules or ["q_proj", "v_proj"], + ) + strategy.initial_parameters = fl.common.ndarrays_to_parameters(init_arrays) + + fl.server.start_server( + server_address=os.environ["SERVER_ADDRESS"], + strategy=strategy, + config=fl.server.ServerConfig(num_rounds=config.federation.num_rounds), + ) + + final_checkpoint_ref = "" + if state.latest_parameters is not None: + final_checkpoint_ref = save_final_checkpoint( + artifact_dir, state.latest_parameters + ) + elif state.checkpoint_refs: + final_checkpoint_ref = state.checkpoint_refs[-1] + elif strategy.initial_parameters is not None: + final_checkpoint_ref = save_final_checkpoint( + artifact_dir, strategy.initial_parameters + ) + + reported_metrics = sorted(state.last_aggregate_metrics.keys()) + + result = { + "schema_version": "1.1", + "payload_type": "final_summary", + "job_id": _job_id(), + "request_id": _job_id(), + "algorithm_name": "flowertune_llm_medical", + "runtime": "simulation", + "status": "COMPLETED", + "timestamp": utc_now(), + "rounds_total": state.rounds_total, + "rounds_completed": state.rounds_completed, + "requested_metrics": requested_metrics, + "reported_metrics": reported_metrics, + "aggregate_metrics": state.last_aggregate_metrics, + "clients": { + "expected": config.federation.num_clients, + "avg_participated_per_round": state.avg_participated, + }, + "artifacts": { + "final_checkpoint_ref": final_checkpoint_ref, + "intermediate_checkpoints": state.checkpoint_refs, + }, + } + if state.best_round is not None: + result["best_round"] = state.best_round + + validate_final_summary_payload(result) + post_result(result) + + except DatasetLoadError as exc: + post_result( + _failure_summary( + config, + code="DATASET_LOAD_ERROR", + message="Server runtime failed while preparing dataset metadata.", + details={"exception": str(exc)}, + ) + ) + raise + except ModelLoadError as exc: + post_result( + _failure_summary( + config, + code="MODEL_LOAD_ERROR", + message="Model backend preflight failed.", + details={"exception": str(exc), "backend": config.model.backend.value}, + ) + ) + raise + except Exception as exc: # noqa: BLE001 + post_result( + _failure_summary( + config, + code="RUNTIME_ERROR", + message="Server runtime failed during Flower execution.", + details={"exception": str(exc)}, + ) + ) + raise diff --git a/exaflow/algorithms/flower/flowertune_llm_medical/strategy.py b/exaflow/algorithms/flower/flowertune_llm_medical/strategy.py new file mode 100644 index 000000000..74aece20f --- /dev/null +++ b/exaflow/algorithms/flower/flowertune_llm_medical/strategy.py @@ -0,0 +1,137 @@ +"""Flower strategy helpers for flowertune_llm_medical Phase 1 runtime.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from datetime import datetime +from datetime import timezone +from typing import Dict +from typing import List +from typing import Optional + +import numpy as np +from flwr.common import parameters_to_ndarrays +from flwr.server.client_proxy import ClientProxy +from flwr.server.strategy import FedAvg + + +def utc_now() -> str: + return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z") + + +@dataclass +class StrategyState: + rounds_total: int + requested_metrics: List[str] + last_aggregate_metrics: Dict[str, float] + rounds_completed: int + avg_participated: float + best_round: Optional[int] + best_loss: Optional[float] + checkpoint_refs: List[str] + latest_parameters: Optional[object] + + +class ContractFedAvg(FedAvg): + """FedAvg strategy that records summary stats and checkpoints.""" + + def __init__( + self, + state: StrategyState, + checkpoint_every_n_rounds: int, + artifact_dir: str, + *args, + **kwargs, + ): + super().__init__(*args, **kwargs) + self.state = state + self.checkpoint_every_n_rounds = checkpoint_every_n_rounds + self.artifact_dir = artifact_dir + os.makedirs(self.artifact_dir, exist_ok=True) + + def _save_checkpoint(self, rnd: int, parameters) -> None: + arrays = parameters_to_ndarrays(parameters) + ckpt_path = os.path.join(self.artifact_dir, f"round_{rnd:03d}.npz") + np.savez(ckpt_path, **{f"tensor_{i:04d}": arr for i, arr in enumerate(arrays)}) + self.state.checkpoint_refs.append(ckpt_path) + + def aggregate_fit(self, rnd, results, failures): + aggregated = super().aggregate_fit(rnd, results, failures) + if aggregated is None: + return None + parameters, fit_metrics = aggregated + self.state.latest_parameters = parameters + self.state.rounds_completed = rnd + if results: + avg_participated_prev = self.state.avg_participated * max(0, rnd - 1) + self.state.avg_participated = (avg_participated_prev + len(results)) / rnd + if rnd % self.checkpoint_every_n_rounds == 0: + self._save_checkpoint(rnd, parameters) + return parameters, fit_metrics + + def aggregate_evaluate( + self, + rnd: int, + results: List[tuple[ClientProxy, object]], + failures, + ): + aggregated = super().aggregate_evaluate(rnd, results, failures) + if aggregated is None: + return None + loss, metrics = aggregated + merged = {"loss": float(loss)} + merged.update({k: float(v) for k, v in metrics.items()}) + self.state.last_aggregate_metrics = merged + cur_loss = merged.get("loss") + if cur_loss is not None and ( + self.state.best_loss is None or cur_loss < self.state.best_loss + ): + self.state.best_loss = cur_loss + self.state.best_round = rnd + return aggregated + + +def build_strategy( + *, + num_rounds: int, + fraction_fit: float, + fraction_evaluate: float, + min_fit_clients: int, + min_evaluate_clients: int, + min_available_clients: int, + checkpoint_every_n_rounds: int, + artifact_dir: str, + requested_metrics: List[str], +): + state = StrategyState( + rounds_total=num_rounds, + requested_metrics=requested_metrics, + last_aggregate_metrics={}, + rounds_completed=0, + avg_participated=0.0, + best_round=None, + best_loss=None, + checkpoint_refs=[], + latest_parameters=None, + ) + + strategy = ContractFedAvg( + state=state, + checkpoint_every_n_rounds=checkpoint_every_n_rounds, + artifact_dir=artifact_dir, + fraction_fit=fraction_fit, + fraction_evaluate=fraction_evaluate, + min_fit_clients=min_fit_clients, + min_evaluate_clients=min_evaluate_clients, + min_available_clients=min_available_clients, + ) + return strategy, state + + +def save_final_checkpoint(artifact_dir: str, parameters) -> str: + os.makedirs(artifact_dir, exist_ok=True) + arrays = parameters_to_ndarrays(parameters) + ckpt_path = os.path.join(artifact_dir, "final_model.npz") + np.savez(ckpt_path, **{f"tensor_{i:04d}": arr for i, arr in enumerate(arrays)}) + return ckpt_path diff --git a/exaflow/controller/__init__.py b/exaflow/controller/__init__.py index bdef50d40..0be8b0c7c 100644 --- a/exaflow/controller/__init__.py +++ b/exaflow/controller/__init__.py @@ -1,7 +1,7 @@ import os from enum import Enum from enum import unique -from importlib.resources import open_text +from importlib.resources import files import envtoml @@ -20,10 +20,11 @@ class DeploymentType(str, Enum): # Initializing the configurations from the config file config = None if config_file := os.getenv("EXAFLOW_CONTROLLER_CONFIG_FILE"): - with open(config_file) as fp: + with open(config_file, "rb") as fp: config = AttrDict(envtoml.load(fp)) else: - with open_text(controller, "config.toml") as fp: + config_path = files(controller).joinpath("config.toml") + with config_path.open("rb") as fp: config = AttrDict(envtoml.load(fp)) worker_landscape_aggregator = None diff --git a/exaflow/controller/quart/endpoints.py b/exaflow/controller/quart/endpoints.py index 787b2261f..8124ecc25 100644 --- a/exaflow/controller/quart/endpoints.py +++ b/exaflow/controller/quart/endpoints.py @@ -2,6 +2,7 @@ import pydantic from quart import Blueprint +from quart import abort from quart import jsonify from quart import request @@ -127,9 +128,56 @@ async def run_algorithm(algorithm_name: str) -> str: @algorithms.route("/flower/input", methods=["GET"]) async def get_flower_input() -> dict: + request_id = request.args.get("request_id") + if not get_flower_controller().flower_execution_info.validate_execution_scope( + request_id + ): + abort(400, description="Flower request scope mismatch.") return get_flower_controller().flower_execution_info.get_inputdata() +@algorithms.route("/flower/parameters", methods=["GET"]) +async def get_flower_parameters() -> dict: + request_id = request.args.get("request_id") + if not get_flower_controller().flower_execution_info.validate_execution_scope( + request_id + ): + abort(400, description="Flower request scope mismatch.") + return get_flower_controller().flower_execution_info.get_parameters() + + +@algorithms.route("/flower/run_env", methods=["GET"]) +async def get_flower_run_env() -> dict: + request_id = request.args.get("request_id") + if not get_flower_controller().flower_execution_info.validate_execution_scope( + request_id + ): + abort(400, description="Flower request scope mismatch.") + return get_flower_controller().flower_execution_info.get_run_env() + + +@algorithms.route("/flower/events", methods=["GET"]) +async def get_flower_events() -> dict: + request_id = request.args.get("request_id") + if not get_flower_controller().flower_execution_info.validate_execution_scope( + request_id + ): + abort(400, description="Flower request scope mismatch.") + return {"events": get_flower_controller().flower_execution_info.get_events()} + + +@algorithms.route("/flower/event", methods=["POST"]) +async def add_flower_event(): + request_id = request.args.get("request_id") + if not get_flower_controller().flower_execution_info.validate_execution_scope( + request_id + ): + abort(400, description="Flower request scope mismatch.") + request_body = await request.json + await get_flower_controller().flower_execution_info.add_event(event=request_body) + return jsonify({"message": "Event added successfully"}), 200 + + @algorithms.route("/flower/result", methods=["POST"]) async def set_flower_result(): request_body = await request.json diff --git a/exaflow/controller/services/api/algorithm_request_validator.py b/exaflow/controller/services/api/algorithm_request_validator.py index f02e33591..efccb6cd1 100644 --- a/exaflow/controller/services/api/algorithm_request_validator.py +++ b/exaflow/controller/services/api/algorithm_request_validator.py @@ -4,6 +4,9 @@ from typing import List from typing import Optional +from pydantic import ValidationError + +from exaflow.algorithms.flower.flowertune_llm_medical.run_config import parse_run_config from exaflow.algorithms.specifications import AlgorithmSpecification from exaflow.algorithms.specifications import InputDataSpecification from exaflow.algorithms.specifications import InputDataSpecifications @@ -45,6 +48,16 @@ def validate_algorithm_request( ): algorithm_specs = _get_algorithm_specs(algorithm_name, algorithms_specs) + if _skip_wla_validation_for_flowertune_simulation( + algorithm_name=algorithm_name, + parameters=algorithm_request_dto.parameters, + ): + _validate_custom_algorithm_contracts( + algorithm_name=algorithm_name, + parameters=algorithm_request_dto.parameters, + ) + return + ( training_datasets, validation_datasets, @@ -75,6 +88,16 @@ def _get_algorithm_specs( return algorithms_specs[algorithm_name] +def _skip_wla_validation_for_flowertune_simulation( + algorithm_name: str, parameters: Optional[Dict[str, Any]] +) -> bool: + if algorithm_name != "flowertune_llm_medical": + return False + + runtime = (parameters or {}).get("runtime", "") + return str(runtime).lower() == "simulation" + + def _validate_algorithm_request_body( algorithm_request_dto: AlgorithmRequestDTO, algorithm_specs: AlgorithmSpecification, @@ -113,6 +136,22 @@ def _validate_algorithm_request_body( transformers_specs=transformers_specs, data_model_cdes=data_model_cdes, ) + _validate_custom_algorithm_contracts( + algorithm_name=algorithm_specs.name, + parameters=algorithm_request_dto.parameters, + ) + + +def _validate_custom_algorithm_contracts( + algorithm_name: str, parameters: Optional[Dict[str, Any]] +): + if algorithm_name != "flowertune_llm_medical": + return + + try: + parse_run_config(parameters or {}) + except ValidationError as exc: + raise BadUserInput(f"Invalid RunConfig for '{algorithm_name}': {exc}") from exc def _validate_inputdata( diff --git a/exaflow/controller/services/flower/flower_io_registry.py b/exaflow/controller/services/flower/flower_io_registry.py index 56daf7c0d..dfd18c85a 100644 --- a/exaflow/controller/services/flower/flower_io_registry.py +++ b/exaflow/controller/services/flower/flower_io_registry.py @@ -25,6 +25,11 @@ def __repr__(self): class FlowerIORegistry: def __init__(self, timeout, logger): self._inputdata: Optional[dict] = None + self._parameters: Optional[dict] = None + self._run_env: Optional[dict] = None + self._events: Optional[list] = None + self._request_id: Optional[str] = None + self._algorithm_name: Optional[str] = None self._result: Optional[Result] = None self.result_ready: Optional[asyncio.Event] = None self._logger = logger @@ -34,6 +39,11 @@ def __init__(self, timeout, logger): def _reset_sync(self): """Synchronously resets the algorithm execution info to initial state.""" self._inputdata = {} + self._parameters = {} + self._run_env = {} + self._events = [] + self._request_id = None + self._algorithm_name = None self._result = Result(content={}, status=Status.RUNNING) self.result_ready = asyncio.Event() self._logger.debug("Algorithm reset: input data cleared, status set to RUNNING") @@ -83,3 +93,51 @@ def get_inputdata(self) -> dict: """Returns the current input data.""" self._logger.debug(f"Input data retrieved: {self._inputdata}") return self._inputdata + + def set_parameters(self, parameters: dict): + """Sets validated algorithm parameters for Flower processes.""" + self._parameters = parameters or {} + self._logger.debug("Flower parameters updated.") + + def get_parameters(self) -> dict: + """Returns current algorithm parameters.""" + self._logger.debug("Flower parameters retrieved.") + return self._parameters + + def set_run_env(self, run_env: dict): + """Sets computed runtime env overrides for Flower processes.""" + self._run_env = run_env or {} + self._logger.debug("Flower run env updated.") + + def get_run_env(self) -> dict: + """Returns current runtime env overrides.""" + self._logger.debug("Flower run env retrieved.") + return self._run_env + + async def add_event(self, event: Dict[str, Any]): + """Adds a runtime event payload (e.g. round_event).""" + self._events.append(event) + self._logger.debug("Flower event appended. total_events=%s", len(self._events)) + + def get_events(self) -> list: + """Returns runtime events for the active execution.""" + return list(self._events or []) + + def set_execution_context(self, request_id: str, algorithm_name: str): + """Sets active execution context for scoped Flower access.""" + self._request_id = request_id + self._algorithm_name = algorithm_name + self._logger.debug( + "Flower execution context set. request_id=%s algorithm=%s", + request_id, + algorithm_name, + ) + + def validate_execution_scope(self, request_id: Optional[str]) -> bool: + """Validates that the caller request id matches the active Flower execution.""" + if request_id is None: + # Backward-compatible access for existing Flower algorithms. + return True + if not self._request_id: + return False + return request_id == self._request_id diff --git a/exaflow/controller/services/flower/strategies.py b/exaflow/controller/services/flower/strategies.py index ce19c736f..1a5ff3213 100644 --- a/exaflow/controller/services/flower/strategies.py +++ b/exaflow/controller/services/flower/strategies.py @@ -1,7 +1,30 @@ import asyncio +import os from typing import List +from typing import Tuple + +import numpy as np from exaflow import flower_algorithm_folder_paths +from exaflow.algorithms.flower.flowertune_llm_medical.dataset import load_partition +from exaflow.algorithms.flower.flowertune_llm_medical.dataset import load_text_partition +from exaflow.algorithms.flower.flowertune_llm_medical.metrics_out import ( + validate_final_summary_payload, +) +from exaflow.algorithms.flower.flowertune_llm_medical.metrics_out import ( + validate_round_event_payload, +) +from exaflow.algorithms.flower.flowertune_llm_medical.models import ModelLoadError +from exaflow.algorithms.flower.flowertune_llm_medical.models import create_backend_model +from exaflow.algorithms.flower.flowertune_llm_medical.models import preflight_backend +from exaflow.algorithms.flower.flowertune_llm_medical.run_config import ( + build_env_mapping, +) +from exaflow.algorithms.flower.flowertune_llm_medical.run_config import parse_run_config +from exaflow.algorithms.flower.flowertune_llm_medical.run_config import ( + serialize_run_config, +) +from exaflow.algorithms.flower.flowertune_llm_medical.strategy import utc_now from exaflow.controller import config as ctrl_config from exaflow.controller.federation_info_logs import log_experiment_execution from exaflow.controller.services.errors import WorkerTaskTimeoutError @@ -19,6 +42,7 @@ class FlowerStrategy(AlgorithmExecutionStrategyI): async def execute(self) -> str: async with self._controller.algorithm_execution_lock: + await self._controller.flower_execution_info.reset() data_model = self._algorithm_request_dto.inputdata.data_model datasets = self._algorithm_request_dto.inputdata.datasets + ( self._algorithm_request_dto.inputdata.validation_datasets @@ -26,6 +50,29 @@ async def execute(self) -> str: else [] ) + self._controller.flower_execution_info.set_inputdata( + inputdata=self._algorithm_request_dto.inputdata.dict() + ) + self._controller.flower_execution_info.set_execution_context( + request_id=self._request_id, algorithm_name=self._algorithm_name + ) + raw_parameters = self._algorithm_request_dto.parameters or {} + self._controller.flower_execution_info.set_parameters(raw_parameters) + self._controller.flower_execution_info.set_run_env({}) + + run_config = None + if self._algorithm_name == "flowertune_llm_medical": + run_config = parse_run_config(raw_parameters) + normalized_parameters = serialize_run_config(run_config) + run_env = build_env_mapping(run_config, self._request_id) + self._controller.flower_execution_info.set_parameters( + normalized_parameters + ) + self._controller.flower_execution_info.set_run_env(run_env) + + if run_config.runtime.value == "simulation": + return await self._execute_flowertune_local_simulation(run_config) + self._safe_worker_call( "garbage collect on global worker", self._global_worker_tasks_handler.garbage_collect, @@ -36,9 +83,6 @@ async def execute(self) -> str: handler.garbage_collect, ) - self._controller.flower_execution_info.set_inputdata( - inputdata=self._algorithm_request_dto.inputdata.dict() - ) server_pid = None clients_pids = {} server_address = f"{self._controller.worker_landscape_aggregator.get_global_worker().ip}:{ctrl_config.flower.server_port}" @@ -89,6 +133,274 @@ async def execute(self) -> str: clients_pids, ) + async def _execute_flowertune_local_simulation(self, run_config): + try: + preflight_backend(run_config.model.backend.value) + inputdata = self._controller.flower_execution_info.get_inputdata() + num_clients = run_config.federation.num_clients + num_rounds = run_config.federation.num_rounds + local_steps = run_config.local_training.local_steps or 1 + requested_metrics = [m.value for m in run_config.evaluation.metrics] + + artifact_dir = run_config.artifacts.artifact_dir.replace( + "${request_id}", self._request_id + ) + os.makedirs(artifact_dir, exist_ok=True) + + if run_config.model.backend.value == "tiny": + client_partitions = [] + for cid in range(num_clients): + x_train, y_train, x_val, y_val = load_partition( + inputdata, + seed=run_config.seed + (cid * 101), + val_split_ratio=run_config.dataset.val_split_ratio, + ) + client_partitions.append((x_train, y_train, x_val, y_val)) + + n_features = client_partitions[0][0].shape[1] + template = create_backend_model( + backend=run_config.model.backend.value, + n_features=n_features, + learning_rate=run_config.optimizer.learning_rate, + model_name=run_config.model.model_name, + local_steps=local_steps, + max_seq_length=run_config.model.max_seq_length, + lora_r=run_config.peft.r or 4, + lora_alpha=run_config.peft.alpha or 8, + lora_dropout=run_config.peft.dropout or 0.0, + target_modules=run_config.peft.target_modules + or ["q_proj", "v_proj"], + ) + else: + client_partitions = [] + for cid in range(num_clients): + train_texts, val_texts = load_text_partition( + inputdata, + seed=run_config.seed + (cid * 101), + val_split_ratio=run_config.dataset.val_split_ratio, + ) + client_partitions.append((train_texts, val_texts)) + + n_features = 1 + template = create_backend_model( + backend=run_config.model.backend.value, + n_features=n_features, + learning_rate=run_config.optimizer.learning_rate, + model_name=run_config.model.model_name, + local_steps=local_steps, + max_seq_length=run_config.model.max_seq_length, + lora_r=run_config.peft.r or 4, + lora_alpha=run_config.peft.alpha or 8, + lora_dropout=run_config.peft.dropout or 0.0, + target_modules=run_config.peft.target_modules + or ["q_proj", "v_proj"], + ) + + global_params = template.get_parameters() + best_round = None + best_loss = None + final_metrics = {} + checkpoint_refs = [] + + for rnd in range(1, num_rounds + 1): + local_updates: List[Tuple[List[np.ndarray], int]] = [] + round_eval: List[Tuple[dict, int]] = [] + + for cid in range(num_clients): + model = create_backend_model( + backend=run_config.model.backend.value, + n_features=n_features, + learning_rate=run_config.optimizer.learning_rate, + model_name=run_config.model.model_name, + local_steps=local_steps, + max_seq_length=run_config.model.max_seq_length, + lora_r=run_config.peft.r or 4, + lora_alpha=run_config.peft.alpha or 8, + lora_dropout=run_config.peft.dropout or 0.0, + target_modules=run_config.peft.target_modules + or ["q_proj", "v_proj"], + ) + model.set_parameters(global_params) + + if run_config.model.backend.value == "tiny": + x_train, y_train, x_val, y_val = client_partitions[cid] + model.fit_round(x_train, y_train, local_steps) + metrics = model.evaluate_round(x_val, y_val) + local_updates.append( + (model.get_parameters(), max(1, len(y_train))) + ) + round_eval.append((metrics, max(1, len(y_val)))) + else: + train_texts, val_texts = client_partitions[cid] + model.fit_round(train_texts, local_steps) + metrics = model.evaluate_round(val_texts) + local_updates.append( + (model.get_parameters(), max(1, len(train_texts))) + ) + round_eval.append((metrics, max(1, len(val_texts)))) + + global_params = self._weighted_average_parameters(local_updates) + + total_eval = sum(n for _, n in round_eval) + agg_loss = sum(m.get("loss", 0.0) * n for m, n in round_eval) / max( + 1, total_eval + ) + agg_perplexity = sum( + m.get("perplexity", 1.0) * n for m, n in round_eval + ) / max(1, total_eval) + final_metrics = { + "loss": float(agg_loss), + "perplexity": float(agg_perplexity), + } + + if best_loss is None or agg_loss < best_loss: + best_loss = agg_loss + best_round = rnd + + checkpoint_ref = None + if rnd % run_config.artifacts.checkpoint_every_n_rounds == 0: + checkpoint_ref = self._save_checkpoint( + artifact_dir, rnd, global_params + ) + checkpoint_refs.append(checkpoint_ref) + + round_event = { + "schema_version": "1.1", + "payload_type": "round_event", + "job_id": self._request_id, + "request_id": self._request_id, + "algorithm_name": "flowertune_llm_medical", + "runtime": "simulation", + "status": "RUNNING", + "round": rnd, + "rounds_total": num_rounds, + "timestamp": utc_now(), + "requested_metrics": requested_metrics, + "reported_metrics": sorted(final_metrics.keys()), + "metrics": final_metrics, + "clients": { + "participated": num_clients, + "expected": num_clients, + "failed": 0, + }, + "artifacts": {"checkpoint_ref": checkpoint_ref}, + } + validate_round_event_payload(round_event) + await self._controller.flower_execution_info.add_event(round_event) + + final_checkpoint_ref = self._save_final_checkpoint( + artifact_dir, global_params + ) + + result = { + "schema_version": "1.1", + "payload_type": "final_summary", + "job_id": self._request_id, + "request_id": self._request_id, + "algorithm_name": "flowertune_llm_medical", + "runtime": "simulation", + "status": "COMPLETED", + "timestamp": utc_now(), + "rounds_total": num_rounds, + "rounds_completed": num_rounds, + "requested_metrics": requested_metrics, + "reported_metrics": sorted(final_metrics.keys()), + "aggregate_metrics": final_metrics, + "clients": { + "expected": num_clients, + "avg_participated_per_round": float(num_clients), + }, + "artifacts": { + "final_checkpoint_ref": final_checkpoint_ref, + "intermediate_checkpoints": checkpoint_refs, + }, + } + if best_round is not None: + result["best_round"] = best_round + + validate_final_summary_payload(result) + await self._controller.flower_execution_info.set_result(result) + return result + + except ModelLoadError as exc: + return await self._set_failure_summary( + run_config, + code="MODEL_LOAD_ERROR", + message="Model backend preflight failed.", + details={ + "exception": str(exc), + "backend": run_config.model.backend.value, + }, + ) + except Exception as exc: # noqa: BLE001 + return await self._set_failure_summary( + run_config, + code="RUNTIME_ERROR", + message="Local simulation runtime failed.", + details={"exception": str(exc)}, + ) + + async def _set_failure_summary( + self, run_config, *, code: str, message: str, details: dict + ): + requested_metrics = [m.value for m in run_config.evaluation.metrics] + result = { + "schema_version": "1.1", + "payload_type": "final_summary", + "job_id": self._request_id, + "request_id": self._request_id, + "algorithm_name": "flowertune_llm_medical", + "runtime": "simulation", + "status": "FAILED", + "timestamp": utc_now(), + "rounds_total": run_config.federation.num_rounds, + "rounds_completed": 0, + "requested_metrics": requested_metrics, + "reported_metrics": [], + "aggregate_metrics": {}, + "error": {"code": code, "message": message, "details": details}, + } + validate_final_summary_payload(result) + await self._controller.flower_execution_info.set_result(result) + return result + + @staticmethod + def _weighted_average_parameters( + parameter_sets: List[Tuple[List[np.ndarray], int]], + ) -> List[np.ndarray]: + total_examples = sum(num_examples for _, num_examples in parameter_sets) + if total_examples <= 0: + raise ValueError("Cannot aggregate parameters with zero examples") + avg_params = None + for params, num_examples in parameter_sets: + weight = num_examples / total_examples + if avg_params is None: + avg_params = [weight * p for p in params] + else: + for i, p in enumerate(params): + avg_params[i] += weight * p + return avg_params + + @staticmethod + def _save_checkpoint( + artifact_dir: str, rnd: int, parameters: List[np.ndarray] + ) -> str: + os.makedirs(artifact_dir, exist_ok=True) + ckpt_path = os.path.join(artifact_dir, f"round_{rnd:03d}.npz") + np.savez( + ckpt_path, **{f"tensor_{i:04d}": arr for i, arr in enumerate(parameters)} + ) + return ckpt_path + + @staticmethod + def _save_final_checkpoint(artifact_dir: str, parameters: List[np.ndarray]) -> str: + os.makedirs(artifact_dir, exist_ok=True) + ckpt_path = os.path.join(artifact_dir, "final_model.npz") + np.savez( + ckpt_path, **{f"tensor_{i:04d}": arr for i, arr in enumerate(parameters)} + ) + return ckpt_path + def _safe_worker_call(self, action_desc, func, *args, **kwargs): try: func(*args, **kwargs) @@ -100,7 +412,6 @@ def _safe_worker_call(self, action_desc, func, *args, **kwargs): async def _cleanup( self, algorithm_name, server_task_handler, server_pid, clients_pids ): - await self._controller.flower_execution_info.reset() if server_pid is not None: self._safe_worker_call( f"stop flower server pid={server_pid}",