Convert JobAPI to Recipe for Kaplan-Meier example #3894
Conversation
There was a problem hiding this comment.
Pull request overview
This PR converts the Kaplan-Meier homomorphic encryption example from the JobAPI to the Recipe API, enabling both simulation and production deployment modes. The key changes include:
- Replaced
km_job.pywith a new Recipe-basedjob.pythat supports both simulation and production environments - Added production deployment infrastructure including provisioning configuration (
project.yml) and convenience scripts (start_all.sh) - Updated controller initialization to explicitly pass empty
persistor_idparameter - Enhanced documentation with comprehensive instructions for both deployment modes
Reviewed changes
Copilot reviewed 7 out of 9 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| examples/advanced/kaplan-meier-he/job.py | New Recipe-based job configuration replacing JobAPI implementation |
| examples/advanced/kaplan-meier-he/km_job.py | Removed old JobAPI implementation |
| examples/advanced/kaplan-meier-he/server.py | Updated controller initialization with explicit persistor_id |
| examples/advanced/kaplan-meier-he/server_he.py | Updated controller initialization with explicit persistor_id |
| examples/advanced/kaplan-meier-he/project.yml | New provisioning configuration for production deployment with CKKS HE scheme |
| examples/advanced/kaplan-meier-he/start_all.sh | Convenience script for starting all production components locally |
| examples/advanced/kaplan-meier-he/README.md | Comprehensive documentation update covering both simulation and production modes |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Greptile SummaryThis PR successfully migrates the Kaplan-Meier homomorphic encryption example from the deprecated JobAPI pattern to the modern Recipe API. The changes enable both simulation and production deployment modes. Key Changes:
Technical Improvements:
Confidence Score: 5/5
Important Files Changed
Sequence DiagramsequenceDiagram
participant Admin as Admin Console
participant Server as FL Server<br/>(KM_HE Controller)
participant C1 as Client 1
participant C2 as Client 2
participant CN as Client N
Note over Admin,CN: Production Mode: Provisioned HE context via startup kits
Note over Admin,CN: Simulation Mode: Manual HE context preparation
Admin->>Server: Submit KM_HE Job
Note over Server,CN: Round 1: Collect Maximum Histogram Index
Server->>C1: Empty start message
Server->>C2: Empty start message
Server->>CN: Empty start message
C1->>C1: Generate local histogram<br/>from event data
C2->>C2: Generate local histogram<br/>from event data
CN->>CN: Generate local histogram<br/>from event data
C1->>Server: max_idx (cleartext)
C2->>Server: max_idx (cleartext)
CN->>Server: max_idx (cleartext)
Server->>Server: Aggregate: max_idx_global = max(all indices) + 1
Note over Server,CN: Round 2: Collect Encrypted Histograms
Server->>C1: max_idx_global (cleartext)
Server->>C2: max_idx_global (cleartext)
Server->>CN: max_idx_global (cleartext)
C1->>C1: Normalize histogram to<br/>global length
C1->>C1: Encrypt with CKKS:<br/>ts.ckks_vector(context, hist)
C2->>C2: Normalize histogram to<br/>global length
C2->>C2: Encrypt with CKKS:<br/>ts.ckks_vector(context, hist)
CN->>CN: Normalize histogram to<br/>global length
CN->>CN: Encrypt with CKKS:<br/>ts.ckks_vector(context, hist)
C1->>Server: hist_obs, hist_cen (ciphertext)
C2->>Server: hist_obs, hist_cen (ciphertext)
CN->>Server: hist_obs, hist_cen (ciphertext)
Server->>Server: Aggregate encrypted vectors:<br/>global = c1 + c2 + ... + cn
Note over Server,CN: Round 3: Distribute Global Encrypted Histograms
Server->>C1: hist_obs_global, hist_cen_global (ciphertext)
Server->>C2: hist_obs_global, hist_cen_global (ciphertext)
Server->>CN: hist_obs_global, hist_cen_global (ciphertext)
C1->>C1: Decrypt: ts.ckks_vector_from(context, data).decrypt()
C1->>C1: Round floats to integers
C1->>C1: Unfold to event list
C1->>C1: Perform KM analysis
C1->>C1: Save km_curve_fl_he.png & km_global.json
C2->>C2: Decrypt & perform KM analysis
CN->>CN: Decrypt & perform KM analysis
C1->>Server: Empty response
C2->>Server: Empty response
CN->>Server: Empty response
Server->>Admin: Job Complete
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
There was a problem hiding this comment.
Additional Comments (1)
-
examples/advanced/kaplan-meier-he/start_all.sh, line 71 (link)style:
stop_all.shscript doesn't existUpdate to mention using the admin console instead:
9 files reviewed, 1 comment
There was a problem hiding this comment.
Additional Comments (4)
-
examples/advanced/kaplan-meier-he/client_he.py, line 140-141 (link)logic: HE scheme mismatch - code uses BFV (
ts.bfv_vector) butproject.ymlprovisions CKKS scheme for production. Production mode will fail when it tries to use BFV functions with CKKS context. Either changeproject.ymlscheme to BFV or update client/server code to use scheme-agnostic operations. -
examples/advanced/kaplan-meier-he/server_he.py, line 95-97 (link)logic: HE scheme mismatch - code uses BFV (
ts.bfv_vector_from) butproject.ymlprovisions CKKS scheme for production. Production mode will fail at deserialization. Either changeproject.ymlscheme to BFV or update code to be scheme-agnostic. -
examples/advanced/kaplan-meier-he/project.yml, line 50 (link)logic: CKKS scheme incompatible with code -
client_he.pyandserver_he.pyuse BFV-specific functions (ts.bfv_vector,ts.bfv_vector_from) for integer histogram operations. Change to BFV scheme or update code. -
examples/advanced/kaplan-meier-he/job.py, line 143 (link)logic: string replacement will fail if
--workspace_dirdoesn't contain "/km/". usepathlibor explicit path construction instead.
9 files reviewed, 4 comments
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
|
/build |
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
|
/build |
chesterxgchen
left a comment
chesterxgchen
left a comment
There was a problem hiding this comment.
LGTM, we need to discuss how to avoid to use provision (heavy handed) in order to use HE
|
/build |
There was a problem hiding this comment.
Greptile Overview
Greptile Summary
This PR successfully converts the Kaplan-Meier HE example from JobAPI to the Recipe pattern, enabling both simulation and production deployment modes with homomorphic encryption support.
Major Changes:
- Converted
km_job.pyto newjob.pywith Recipe pattern (KMRecipeclass) - Renamed client/server files (removed
src/directory structure) - Added
project.ymlfor production provisioning withHEBuilder(CKKS scheme) - Added
start_all.shconvenience script for local production testing - Enhanced HE context handling to support both simulation (base64-encoded
.txt) and production (raw binary.tenseal) formats - Updated
prepare_he_context.pyto generate relinearization keys for CKKS compatibility - Comprehensive README update with detailed instructions for both modes
Issues Found:
job.pyhas type hint mismatch: parameters acceptNonebut are typed asstr- Fragile path manipulation using string replacement instead of
os.path - Missing
osimport for proposed path fix
Confidence Score: 3/5
- This PR is mostly safe to merge but requires fixing critical type hint and path handling issues
- Score reflects solid architecture changes and comprehensive documentation, but marked down due to: (1) type hint mismatch that could cause runtime issues when
Noneis passed for HE context paths without encryption, (2) fragile string-based path manipulation that will fail with custom paths, (3) missing import. These are fixable issues that don't break core functionality but should be addressed before merge. - Primary attention needed for
job.py- fix type hints, path handling, and add missing import. All other files are well-implemented.
Important Files Changed
File Analysis
| Filename | Score | Overview |
|---|---|---|
| examples/advanced/kaplan-meier-he/job.py | 3/5 | Converted JobAPI to Recipe pattern, added production support. Issues: missing import, type hint mismatch, fragile path manipulation |
| examples/advanced/kaplan-meier-he/client_he.py | 5/5 | Renamed from kaplan_meier_train_he.py, added support for both simulation and production HE context formats |
| examples/advanced/kaplan-meier-he/server_he.py | 5/5 | Renamed from kaplan_meier_wf_he.py, added support for both simulation and production HE context formats |
| examples/advanced/kaplan-meier-he/project.yml | 5/5 | New project configuration for provisioning with HEBuilder using CKKS scheme for production deployment |
| examples/advanced/kaplan-meier-he/README.md | 5/5 | Comprehensive documentation update covering Recipe API, simulation/production modes, and HE context provisioning |
Sequence Diagram
sequenceDiagram
participant User
participant job.py
participant Recipe
participant FedJob
participant SimEnv/ProdEnv
participant Server
participant Client
User->>job.py: Run with --encryption flag
job.py->>job.py: Parse args & determine mode
alt Simulation Mode
job.py->>job.py: Load HE context from /tmp/nvflare/he_context/*.txt
else Production Mode
job.py->>job.py: Use provisioned context filenames
end
job.py->>Recipe: Create KMRecipe(encryption, num_clients, paths)
Recipe->>FedJob: Create FedJob(KM or KM_HE)
alt With HE
Recipe->>Server: Add KM_HE controller
Recipe->>Client: Add ScriptRunner(client_he.py)
else Without HE
Recipe->>Server: Add KM controller
Recipe->>Client: Add ScriptRunner(client.py)
end
job.py->>FedJob: Export job
alt Simulation Mode
job.py->>SimEnv: Create SimEnv(num_clients, workspace)
job.py->>Recipe: Execute with SimEnv
else Production Mode
job.py->>ProdEnv: Create ProdEnv(startup_kit)
job.py->>Recipe: Execute with ProdEnv
end
Recipe->>Server: Start workflow
alt With HE (3 rounds)
Server->>Client: Round 1: Request max index
Client->>Server: Send max_idx (plaintext)
Server->>Client: Round 2: Distribute global max_idx
Client->>Server: Send encrypted histograms
Server->>Client: Round 3: Send encrypted global histograms
Client->>Client: Decrypt & perform KM analysis
else Without HE (2 rounds)
Server->>Client: Round 1: Request histograms
Client->>Server: Send local histograms (plaintext)
Server->>Client: Round 2: Send aggregated histograms
Client->>Client: Perform KM analysis
end
Client->>Client: Save km_curve and km_global.json
Recipe-->>User: Return job status
Additional Comments (3)
|
2 participants
Fixes # .
Description
Convert KM example from JobAPI to Recipe, also add production instructions with provisioned HE context
Types of changes
./runtest.sh.