Compare commits
78 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 0aa3b54827 | |||
|
|
693efbb72b | ||
|
|
7e371fd758 | ||
|
|
56baae78b6 | ||
|
|
1174bb3f54 | ||
|
|
96ab75e736 | ||
|
|
1e20d73f7f | ||
|
|
9f218af6b0 | ||
|
|
99392034e5 | ||
|
|
f988a3ae20 | ||
|
|
57b7979999 | ||
|
|
75b56251ef | ||
|
|
2649476571 | ||
|
|
bba4daaddd | ||
|
|
0742078fad | ||
|
|
79605db260 | ||
| 4ae76cfd83 | |||
|
|
14b27e4c30 | ||
|
|
09349b6274 | ||
|
|
2ace55c241 | ||
|
|
6981d09ac6 | ||
| 860bce1ba3 | |||
|
|
c190267292 | ||
|
|
206b47dfae | ||
|
|
9c628e70a6 | ||
|
|
db32450c68 | ||
|
|
a01763b69c | ||
|
|
e5ebf243bd | ||
|
|
cff81e3594 | ||
|
|
b5a54147f5 | ||
|
|
150e2f7259 | ||
|
|
e5cea1ef96 | ||
|
|
08b2986f81 | ||
|
|
8b69b2cb29 | ||
| 9920cd61c2 | |||
| d27f6573fe | |||
| f323da7b56 | |||
| 5c3c9292fc | |||
| 9fb81ef556 | |||
| 26e37e0a0a | |||
| 116330fc3f | |||
| 00b6382a88 | |||
| 4c51e7c2b3 | |||
| d3691989bb | |||
| 272e02dadc | |||
| 19c8a028d2 | |||
| 11691398da | |||
| 054dcbfa50 | |||
| e393263ddf | |||
| 11224537d3 | |||
| 26502b5274 | |||
| ad4da333a0 | |||
| a4594997e8 | |||
| e53e792e79 | |||
| 6c4095825d | |||
| 7a776066ec | |||
| c919713075 | |||
| ee9c71cbbb | |||
| c9ac25ea80 | |||
| 0c9bea54c0 | |||
| c2e2a8ff7a | |||
| 7bf7848a93 | |||
| 4edf70c4a0 | |||
| 83d3a4482c | |||
| b37ec60b5d | |||
| 00cf8a2ad6 | |||
| 84e6501063 | |||
| 4989cfcb3c | |||
| 16ea436e45 | |||
| 02913e341d | |||
| d26e89aa98 | |||
| e22b0bc7c3 | |||
| 8d34228d7b | |||
| 10c7e292b9 | |||
| e6018af23f | |||
| aaf647efdf | |||
| 424be007ef | |||
| 1ce8b9d3a7 |
152
.claude/commands/opsx/apply.md
Normal file
152
.claude/commands/opsx/apply.md
Normal file
@@ -0,0 +1,152 @@
|
||||
---
|
||||
name: "OPSX: Apply"
|
||||
description: Implement tasks from an OpenSpec change (Experimental)
|
||||
category: Workflow
|
||||
tags: [workflow, artifacts, experimental]
|
||||
---
|
||||
|
||||
Implement tasks from an OpenSpec change.
|
||||
|
||||
**Input**: Optionally specify a change name (e.g., `/opsx:apply add-auth`). If omitted, check if it can be inferred from conversation context. If vague or ambiguous you MUST prompt for available changes.
|
||||
|
||||
**Steps**
|
||||
|
||||
1. **Select the change**
|
||||
|
||||
If a name is provided, use it. Otherwise:
|
||||
- Infer from conversation context if the user mentioned a change
|
||||
- Auto-select if only one active change exists
|
||||
- If ambiguous, run `openspec list --json` to get available changes and use the **AskUserQuestion tool** to let the user select
|
||||
|
||||
Always announce: "Using change: <name>" and how to override (e.g., `/opsx:apply <other>`).
|
||||
|
||||
2. **Check status to understand the schema**
|
||||
```bash
|
||||
openspec status --change "<name>" --json
|
||||
```
|
||||
Parse the JSON to understand:
|
||||
- `schemaName`: The workflow being used (e.g., "spec-driven")
|
||||
- Which artifact contains the tasks (typically "tasks" for spec-driven, check status for others)
|
||||
|
||||
3. **Get apply instructions**
|
||||
|
||||
```bash
|
||||
openspec instructions apply --change "<name>" --json
|
||||
```
|
||||
|
||||
This returns:
|
||||
- `contextFiles`: artifact ID -> array of concrete file paths (varies by schema)
|
||||
- Progress (total, complete, remaining)
|
||||
- Task list with status
|
||||
- Dynamic instruction based on current state
|
||||
|
||||
**Handle states:**
|
||||
- If `state: "blocked"` (missing artifacts): show message, suggest using `/opsx:continue`
|
||||
- If `state: "all_done"`: congratulate, suggest archive
|
||||
- Otherwise: proceed to implementation
|
||||
|
||||
4. **Read context files**
|
||||
|
||||
Read every file path listed under `contextFiles` from the apply instructions output.
|
||||
The files depend on the schema being used:
|
||||
- **spec-driven**: proposal, specs, design, tasks
|
||||
- Other schemas: follow the contextFiles from CLI output
|
||||
|
||||
5. **Show current progress**
|
||||
|
||||
Display:
|
||||
- Schema being used
|
||||
- Progress: "N/M tasks complete"
|
||||
- Remaining tasks overview
|
||||
- Dynamic instruction from CLI
|
||||
|
||||
6. **Implement tasks (loop until done or blocked)**
|
||||
|
||||
For each pending task:
|
||||
- Show which task is being worked on
|
||||
- Make the code changes required
|
||||
- Keep changes minimal and focused
|
||||
- Mark task complete in the tasks file: `- [ ]` → `- [x]`
|
||||
- Continue to next task
|
||||
|
||||
**Pause if:**
|
||||
- Task is unclear → ask for clarification
|
||||
- Implementation reveals a design issue → suggest updating artifacts
|
||||
- Error or blocker encountered → report and wait for guidance
|
||||
- User interrupts
|
||||
|
||||
7. **On completion or pause, show status**
|
||||
|
||||
Display:
|
||||
- Tasks completed this session
|
||||
- Overall progress: "N/M tasks complete"
|
||||
- If all done: suggest archive
|
||||
- If paused: explain why and wait for guidance
|
||||
|
||||
**Output During Implementation**
|
||||
|
||||
```
|
||||
## Implementing: <change-name> (schema: <schema-name>)
|
||||
|
||||
Working on task 3/7: <task description>
|
||||
[...implementation happening...]
|
||||
✓ Task complete
|
||||
|
||||
Working on task 4/7: <task description>
|
||||
[...implementation happening...]
|
||||
✓ Task complete
|
||||
```
|
||||
|
||||
**Output On Completion**
|
||||
|
||||
```
|
||||
## Implementation Complete
|
||||
|
||||
**Change:** <change-name>
|
||||
**Schema:** <schema-name>
|
||||
**Progress:** 7/7 tasks complete ✓
|
||||
|
||||
### Completed This Session
|
||||
- [x] Task 1
|
||||
- [x] Task 2
|
||||
...
|
||||
|
||||
All tasks complete! You can archive this change with `/opsx:archive`.
|
||||
```
|
||||
|
||||
**Output On Pause (Issue Encountered)**
|
||||
|
||||
```
|
||||
## Implementation Paused
|
||||
|
||||
**Change:** <change-name>
|
||||
**Schema:** <schema-name>
|
||||
**Progress:** 4/7 tasks complete
|
||||
|
||||
### Issue Encountered
|
||||
<description of the issue>
|
||||
|
||||
**Options:**
|
||||
1. <option 1>
|
||||
2. <option 2>
|
||||
3. Other approach
|
||||
|
||||
What would you like to do?
|
||||
```
|
||||
|
||||
**Guardrails**
|
||||
- Keep going through tasks until done or blocked
|
||||
- Always read context files before starting (from the apply instructions output)
|
||||
- If task is ambiguous, pause and ask before implementing
|
||||
- If implementation reveals issues, pause and suggest artifact updates
|
||||
- Keep code changes minimal and scoped to each task
|
||||
- Update task checkbox immediately after completing each task
|
||||
- Pause on errors, blockers, or unclear requirements - don't guess
|
||||
- Use contextFiles from CLI output, don't assume specific file names
|
||||
|
||||
**Fluid Workflow Integration**
|
||||
|
||||
This skill supports the "actions on a change" model:
|
||||
|
||||
- **Can be invoked anytime**: Before all artifacts are done (if tasks exist), after partial implementation, interleaved with other actions
|
||||
- **Allows artifact updates**: If implementation reveals design issues, suggest updating artifacts - not phase-locked, work fluidly
|
||||
157
.claude/commands/opsx/archive.md
Normal file
157
.claude/commands/opsx/archive.md
Normal file
@@ -0,0 +1,157 @@
|
||||
---
|
||||
name: "OPSX: Archive"
|
||||
description: Archive a completed change in the experimental workflow
|
||||
category: Workflow
|
||||
tags: [workflow, archive, experimental]
|
||||
---
|
||||
|
||||
Archive a completed change in the experimental workflow.
|
||||
|
||||
**Input**: Optionally specify a change name after `/opsx:archive` (e.g., `/opsx:archive add-auth`). If omitted, check if it can be inferred from conversation context. If vague or ambiguous you MUST prompt for available changes.
|
||||
|
||||
**Steps**
|
||||
|
||||
1. **If no change name provided, prompt for selection**
|
||||
|
||||
Run `openspec list --json` to get available changes. Use the **AskUserQuestion tool** to let the user select.
|
||||
|
||||
Show only active changes (not already archived).
|
||||
Include the schema used for each change if available.
|
||||
|
||||
**IMPORTANT**: Do NOT guess or auto-select a change. Always let the user choose.
|
||||
|
||||
2. **Check artifact completion status**
|
||||
|
||||
Run `openspec status --change "<name>" --json` to check artifact completion.
|
||||
|
||||
Parse the JSON to understand:
|
||||
- `schemaName`: The workflow being used
|
||||
- `artifacts`: List of artifacts with their status (`done` or other)
|
||||
|
||||
**If any artifacts are not `done`:**
|
||||
- Display warning listing incomplete artifacts
|
||||
- Prompt user for confirmation to continue
|
||||
- Proceed if user confirms
|
||||
|
||||
3. **Check task completion status**
|
||||
|
||||
Read the tasks file (typically `tasks.md`) to check for incomplete tasks.
|
||||
|
||||
Count tasks marked with `- [ ]` (incomplete) vs `- [x]` (complete).
|
||||
|
||||
**If incomplete tasks found:**
|
||||
- Display warning showing count of incomplete tasks
|
||||
- Prompt user for confirmation to continue
|
||||
- Proceed if user confirms
|
||||
|
||||
**If no tasks file exists:** Proceed without task-related warning.
|
||||
|
||||
4. **Assess delta spec sync state**
|
||||
|
||||
Check for delta specs at `openspec/changes/<name>/specs/`. If none exist, proceed without sync prompt.
|
||||
|
||||
**If delta specs exist:**
|
||||
- Compare each delta spec with its corresponding main spec at `openspec/specs/<capability>/spec.md`
|
||||
- Determine what changes would be applied (adds, modifications, removals, renames)
|
||||
- Show a combined summary before prompting
|
||||
|
||||
**Prompt options:**
|
||||
- If changes needed: "Sync now (recommended)", "Archive without syncing"
|
||||
- If already synced: "Archive now", "Sync anyway", "Cancel"
|
||||
|
||||
If user chooses sync, use Task tool (subagent_type: "general-purpose", prompt: "Use Skill tool to invoke openspec-sync-specs for change '<name>'. Delta spec analysis: <include the analyzed delta spec summary>"). Proceed to archive regardless of choice.
|
||||
|
||||
5. **Perform the archive**
|
||||
|
||||
Create the archive directory if it doesn't exist:
|
||||
```bash
|
||||
mkdir -p openspec/changes/archive
|
||||
```
|
||||
|
||||
Generate target name using current date: `YYYY-MM-DD-<change-name>`
|
||||
|
||||
**Check if target already exists:**
|
||||
- If yes: Fail with error, suggest renaming existing archive or using different date
|
||||
- If no: Move the change directory to archive
|
||||
|
||||
```bash
|
||||
mv openspec/changes/<name> openspec/changes/archive/YYYY-MM-DD-<name>
|
||||
```
|
||||
|
||||
6. **Display summary**
|
||||
|
||||
Show archive completion summary including:
|
||||
- Change name
|
||||
- Schema that was used
|
||||
- Archive location
|
||||
- Spec sync status (synced / sync skipped / no delta specs)
|
||||
- Note about any warnings (incomplete artifacts/tasks)
|
||||
|
||||
**Output On Success**
|
||||
|
||||
```
|
||||
## Archive Complete
|
||||
|
||||
**Change:** <change-name>
|
||||
**Schema:** <schema-name>
|
||||
**Archived to:** openspec/changes/archive/YYYY-MM-DD-<name>/
|
||||
**Specs:** ✓ Synced to main specs
|
||||
|
||||
All artifacts complete. All tasks complete.
|
||||
```
|
||||
|
||||
**Output On Success (No Delta Specs)**
|
||||
|
||||
```
|
||||
## Archive Complete
|
||||
|
||||
**Change:** <change-name>
|
||||
**Schema:** <schema-name>
|
||||
**Archived to:** openspec/changes/archive/YYYY-MM-DD-<name>/
|
||||
**Specs:** No delta specs
|
||||
|
||||
All artifacts complete. All tasks complete.
|
||||
```
|
||||
|
||||
**Output On Success With Warnings**
|
||||
|
||||
```
|
||||
## Archive Complete (with warnings)
|
||||
|
||||
**Change:** <change-name>
|
||||
**Schema:** <schema-name>
|
||||
**Archived to:** openspec/changes/archive/YYYY-MM-DD-<name>/
|
||||
**Specs:** Sync skipped (user chose to skip)
|
||||
|
||||
**Warnings:**
|
||||
- Archived with 2 incomplete artifacts
|
||||
- Archived with 3 incomplete tasks
|
||||
- Delta spec sync was skipped (user chose to skip)
|
||||
|
||||
Review the archive if this was not intentional.
|
||||
```
|
||||
|
||||
**Output On Error (Archive Exists)**
|
||||
|
||||
```
|
||||
## Archive Failed
|
||||
|
||||
**Change:** <change-name>
|
||||
**Target:** openspec/changes/archive/YYYY-MM-DD-<name>/
|
||||
|
||||
Target archive directory already exists.
|
||||
|
||||
**Options:**
|
||||
1. Rename the existing archive
|
||||
2. Delete the existing archive if it's a duplicate
|
||||
3. Wait until a different date to archive
|
||||
```
|
||||
|
||||
**Guardrails**
|
||||
- Always prompt for change selection if not provided
|
||||
- Use artifact graph (openspec status --json) for completion checking
|
||||
- Don't block archive on warnings - just inform and confirm
|
||||
- Preserve .openspec.yaml when moving to archive (it moves with the directory)
|
||||
- Show clear summary of what happened
|
||||
- If sync is requested, use the Skill tool to invoke `openspec-sync-specs` (agent-driven)
|
||||
- If delta specs exist, always run the sync assessment and show the combined summary before prompting
|
||||
173
.claude/commands/opsx/explore.md
Normal file
173
.claude/commands/opsx/explore.md
Normal file
@@ -0,0 +1,173 @@
|
||||
---
|
||||
name: "OPSX: Explore"
|
||||
description: "Enter explore mode - think through ideas, investigate problems, clarify requirements"
|
||||
category: Workflow
|
||||
tags: [workflow, explore, experimental, thinking]
|
||||
---
|
||||
|
||||
Enter explore mode. Think deeply. Visualize freely. Follow the conversation wherever it goes.
|
||||
|
||||
**IMPORTANT: Explore mode is for thinking, not implementing.** You may read files, search code, and investigate the codebase, but you must NEVER write code or implement features. If the user asks you to implement something, remind them to exit explore mode first and create a change proposal. You MAY create OpenSpec artifacts (proposals, designs, specs) if the user asks—that's capturing thinking, not implementing.
|
||||
|
||||
**This is a stance, not a workflow.** There are no fixed steps, no required sequence, no mandatory outputs. You're a thinking partner helping the user explore.
|
||||
|
||||
**Input**: The argument after `/opsx:explore` is whatever the user wants to think about. Could be:
|
||||
- A vague idea: "real-time collaboration"
|
||||
- A specific problem: "the auth system is getting unwieldy"
|
||||
- A change name: "add-dark-mode" (to explore in context of that change)
|
||||
- A comparison: "postgres vs sqlite for this"
|
||||
- Nothing (just enter explore mode)
|
||||
|
||||
---
|
||||
|
||||
## The Stance
|
||||
|
||||
- **Curious, not prescriptive** - Ask questions that emerge naturally, don't follow a script
|
||||
- **Open threads, not interrogations** - Surface multiple interesting directions and let the user follow what resonates. Don't funnel them through a single path of questions.
|
||||
- **Visual** - Use ASCII diagrams liberally when they'd help clarify thinking
|
||||
- **Adaptive** - Follow interesting threads, pivot when new information emerges
|
||||
- **Patient** - Don't rush to conclusions, let the shape of the problem emerge
|
||||
- **Grounded** - Explore the actual codebase when relevant, don't just theorize
|
||||
|
||||
---
|
||||
|
||||
## What You Might Do
|
||||
|
||||
Depending on what the user brings, you might:
|
||||
|
||||
**Explore the problem space**
|
||||
- Ask clarifying questions that emerge from what they said
|
||||
- Challenge assumptions
|
||||
- Reframe the problem
|
||||
- Find analogies
|
||||
|
||||
**Investigate the codebase**
|
||||
- Map existing architecture relevant to the discussion
|
||||
- Find integration points
|
||||
- Identify patterns already in use
|
||||
- Surface hidden complexity
|
||||
|
||||
**Compare options**
|
||||
- Brainstorm multiple approaches
|
||||
- Build comparison tables
|
||||
- Sketch tradeoffs
|
||||
- Recommend a path (if asked)
|
||||
|
||||
**Visualize**
|
||||
```
|
||||
┌─────────────────────────────────────────┐
|
||||
│ Use ASCII diagrams liberally │
|
||||
├─────────────────────────────────────────┤
|
||||
│ │
|
||||
│ ┌────────┐ ┌────────┐ │
|
||||
│ │ State │────────▶│ State │ │
|
||||
│ │ A │ │ B │ │
|
||||
│ └────────┘ └────────┘ │
|
||||
│ │
|
||||
│ System diagrams, state machines, │
|
||||
│ data flows, architecture sketches, │
|
||||
│ dependency graphs, comparison tables │
|
||||
│ │
|
||||
└─────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
**Surface risks and unknowns**
|
||||
- Identify what could go wrong
|
||||
- Find gaps in understanding
|
||||
- Suggest spikes or investigations
|
||||
|
||||
---
|
||||
|
||||
## OpenSpec Awareness
|
||||
|
||||
You have full context of the OpenSpec system. Use it naturally, don't force it.
|
||||
|
||||
### Check for context
|
||||
|
||||
At the start, quickly check what exists:
|
||||
```bash
|
||||
openspec list --json
|
||||
```
|
||||
|
||||
This tells you:
|
||||
- If there are active changes
|
||||
- Their names, schemas, and status
|
||||
- What the user might be working on
|
||||
|
||||
If the user mentioned a specific change name, read its artifacts for context.
|
||||
|
||||
### When no change exists
|
||||
|
||||
Think freely. When insights crystallize, you might offer:
|
||||
|
||||
- "This feels solid enough to start a change. Want me to create a proposal?"
|
||||
- Or keep exploring - no pressure to formalize
|
||||
|
||||
### When a change exists
|
||||
|
||||
If the user mentions a change or you detect one is relevant:
|
||||
|
||||
1. **Read existing artifacts for context**
|
||||
- `openspec/changes/<name>/proposal.md`
|
||||
- `openspec/changes/<name>/design.md`
|
||||
- `openspec/changes/<name>/tasks.md`
|
||||
- etc.
|
||||
|
||||
2. **Reference them naturally in conversation**
|
||||
- "Your design mentions using Redis, but we just realized SQLite fits better..."
|
||||
- "The proposal scopes this to premium users, but we're now thinking everyone..."
|
||||
|
||||
3. **Offer to capture when decisions are made**
|
||||
|
||||
| Insight Type | Where to Capture |
|
||||
|----------------------------|--------------------------------|
|
||||
| New requirement discovered | `specs/<capability>/spec.md` |
|
||||
| Requirement changed | `specs/<capability>/spec.md` |
|
||||
| Design decision made | `design.md` |
|
||||
| Scope changed | `proposal.md` |
|
||||
| New work identified | `tasks.md` |
|
||||
| Assumption invalidated | Relevant artifact |
|
||||
|
||||
Example offers:
|
||||
- "That's a design decision. Capture it in design.md?"
|
||||
- "This is a new requirement. Add it to specs?"
|
||||
- "This changes scope. Update the proposal?"
|
||||
|
||||
4. **The user decides** - Offer and move on. Don't pressure. Don't auto-capture.
|
||||
|
||||
---
|
||||
|
||||
## What You Don't Have To Do
|
||||
|
||||
- Follow a script
|
||||
- Ask the same questions every time
|
||||
- Produce a specific artifact
|
||||
- Reach a conclusion
|
||||
- Stay on topic if a tangent is valuable
|
||||
- Be brief (this is thinking time)
|
||||
|
||||
---
|
||||
|
||||
## Ending Discovery
|
||||
|
||||
There's no required ending. Discovery might:
|
||||
|
||||
- **Flow into a proposal**: "Ready to start? I can create a change proposal."
|
||||
- **Result in artifact updates**: "Updated design.md with these decisions"
|
||||
- **Just provide clarity**: User has what they need, moves on
|
||||
- **Continue later**: "We can pick this up anytime"
|
||||
|
||||
When things crystallize, you might offer a summary - but it's optional. Sometimes the thinking IS the value.
|
||||
|
||||
---
|
||||
|
||||
## Guardrails
|
||||
|
||||
- **Don't implement** - Never write code or implement features. Creating OpenSpec artifacts is fine, writing application code is not.
|
||||
- **Don't fake understanding** - If something is unclear, dig deeper
|
||||
- **Don't rush** - Discovery is thinking time, not task time
|
||||
- **Don't force structure** - Let patterns emerge naturally
|
||||
- **Don't auto-capture** - Offer to save insights, don't just do it
|
||||
- **Do visualize** - A good diagram is worth many paragraphs
|
||||
- **Do explore the codebase** - Ground discussions in reality
|
||||
- **Do question assumptions** - Including the user's and your own
|
||||
106
.claude/commands/opsx/propose.md
Normal file
106
.claude/commands/opsx/propose.md
Normal file
@@ -0,0 +1,106 @@
|
||||
---
|
||||
name: "OPSX: Propose"
|
||||
description: Propose a new change - create it and generate all artifacts in one step
|
||||
category: Workflow
|
||||
tags: [workflow, artifacts, experimental]
|
||||
---
|
||||
|
||||
Propose a new change - create the change and generate all artifacts in one step.
|
||||
|
||||
I'll create a change with artifacts:
|
||||
- proposal.md (what & why)
|
||||
- design.md (how)
|
||||
- tasks.md (implementation steps)
|
||||
|
||||
When ready to implement, run /opsx:apply
|
||||
|
||||
---
|
||||
|
||||
**Input**: The argument after `/opsx:propose` is the change name (kebab-case), OR a description of what the user wants to build.
|
||||
|
||||
**Steps**
|
||||
|
||||
1. **If no input provided, ask what they want to build**
|
||||
|
||||
Use the **AskUserQuestion tool** (open-ended, no preset options) to ask:
|
||||
> "What change do you want to work on? Describe what you want to build or fix."
|
||||
|
||||
From their description, derive a kebab-case name (e.g., "add user authentication" → `add-user-auth`).
|
||||
|
||||
**IMPORTANT**: Do NOT proceed without understanding what the user wants to build.
|
||||
|
||||
2. **Create the change directory**
|
||||
```bash
|
||||
openspec new change "<name>"
|
||||
```
|
||||
This creates a scaffolded change at `openspec/changes/<name>/` with `.openspec.yaml`.
|
||||
|
||||
3. **Get the artifact build order**
|
||||
```bash
|
||||
openspec status --change "<name>" --json
|
||||
```
|
||||
Parse the JSON to get:
|
||||
- `applyRequires`: array of artifact IDs needed before implementation (e.g., `["tasks"]`)
|
||||
- `artifacts`: list of all artifacts with their status and dependencies
|
||||
|
||||
4. **Create artifacts in sequence until apply-ready**
|
||||
|
||||
Use the **TodoWrite tool** to track progress through the artifacts.
|
||||
|
||||
Loop through artifacts in dependency order (artifacts with no pending dependencies first):
|
||||
|
||||
a. **For each artifact that is `ready` (dependencies satisfied)**:
|
||||
- Get instructions:
|
||||
```bash
|
||||
openspec instructions <artifact-id> --change "<name>" --json
|
||||
```
|
||||
- The instructions JSON includes:
|
||||
- `context`: Project background (constraints for you - do NOT include in output)
|
||||
- `rules`: Artifact-specific rules (constraints for you - do NOT include in output)
|
||||
- `template`: The structure to use for your output file
|
||||
- `instruction`: Schema-specific guidance for this artifact type
|
||||
- `outputPath`: Where to write the artifact
|
||||
- `dependencies`: Completed artifacts to read for context
|
||||
- Read any completed dependency files for context
|
||||
- Create the artifact file using `template` as the structure
|
||||
- Apply `context` and `rules` as constraints - but do NOT copy them into the file
|
||||
- Show brief progress: "Created <artifact-id>"
|
||||
|
||||
b. **Continue until all `applyRequires` artifacts are complete**
|
||||
- After creating each artifact, re-run `openspec status --change "<name>" --json`
|
||||
- Check if every artifact ID in `applyRequires` has `status: "done"` in the artifacts array
|
||||
- Stop when all `applyRequires` artifacts are done
|
||||
|
||||
c. **If an artifact requires user input** (unclear context):
|
||||
- Use **AskUserQuestion tool** to clarify
|
||||
- Then continue with creation
|
||||
|
||||
5. **Show final status**
|
||||
```bash
|
||||
openspec status --change "<name>"
|
||||
```
|
||||
|
||||
**Output**
|
||||
|
||||
After completing all artifacts, summarize:
|
||||
- Change name and location
|
||||
- List of artifacts created with brief descriptions
|
||||
- What's ready: "All artifacts created! Ready for implementation."
|
||||
- Prompt: "Run `/opsx:apply` to start implementing."
|
||||
|
||||
**Artifact Creation Guidelines**
|
||||
|
||||
- Follow the `instruction` field from `openspec instructions` for each artifact type
|
||||
- The schema defines what each artifact should contain - follow it
|
||||
- Read dependency artifacts for context before creating new ones
|
||||
- Use `template` as the structure for your output file - fill in its sections
|
||||
- **IMPORTANT**: `context` and `rules` are constraints for YOU, not content for the file
|
||||
- Do NOT copy `<context>`, `<rules>`, `<project_context>` blocks into the artifact
|
||||
- These guide what you write, but should never appear in the output
|
||||
|
||||
**Guardrails**
|
||||
- Create ALL artifacts needed for implementation (as defined by schema's `apply.requires`)
|
||||
- Always read dependency artifacts before creating a new one
|
||||
- If context is critically unclear, ask the user - but prefer making reasonable decisions to keep momentum
|
||||
- If a change with that name already exists, ask if user wants to continue it or create a new one
|
||||
- Verify each artifact file exists after writing before proceeding to next
|
||||
156
.claude/skills/openspec-apply-change/SKILL.md
Normal file
156
.claude/skills/openspec-apply-change/SKILL.md
Normal file
@@ -0,0 +1,156 @@
|
||||
---
|
||||
name: openspec-apply-change
|
||||
description: Implement tasks from an OpenSpec change. Use when the user wants to start implementing, continue implementation, or work through tasks.
|
||||
license: MIT
|
||||
compatibility: Requires openspec CLI.
|
||||
metadata:
|
||||
author: openspec
|
||||
version: "1.0"
|
||||
generatedBy: "1.3.1"
|
||||
---
|
||||
|
||||
Implement tasks from an OpenSpec change.
|
||||
|
||||
**Input**: Optionally specify a change name. If omitted, check if it can be inferred from conversation context. If vague or ambiguous you MUST prompt for available changes.
|
||||
|
||||
**Steps**
|
||||
|
||||
1. **Select the change**
|
||||
|
||||
If a name is provided, use it. Otherwise:
|
||||
- Infer from conversation context if the user mentioned a change
|
||||
- Auto-select if only one active change exists
|
||||
- If ambiguous, run `openspec list --json` to get available changes and use the **AskUserQuestion tool** to let the user select
|
||||
|
||||
Always announce: "Using change: <name>" and how to override (e.g., `/opsx:apply <other>`).
|
||||
|
||||
2. **Check status to understand the schema**
|
||||
```bash
|
||||
openspec status --change "<name>" --json
|
||||
```
|
||||
Parse the JSON to understand:
|
||||
- `schemaName`: The workflow being used (e.g., "spec-driven")
|
||||
- Which artifact contains the tasks (typically "tasks" for spec-driven, check status for others)
|
||||
|
||||
3. **Get apply instructions**
|
||||
|
||||
```bash
|
||||
openspec instructions apply --change "<name>" --json
|
||||
```
|
||||
|
||||
This returns:
|
||||
- `contextFiles`: artifact ID -> array of concrete file paths (varies by schema - could be proposal/specs/design/tasks or spec/tests/implementation/docs)
|
||||
- Progress (total, complete, remaining)
|
||||
- Task list with status
|
||||
- Dynamic instruction based on current state
|
||||
|
||||
**Handle states:**
|
||||
- If `state: "blocked"` (missing artifacts): show message, suggest using openspec-continue-change
|
||||
- If `state: "all_done"`: congratulate, suggest archive
|
||||
- Otherwise: proceed to implementation
|
||||
|
||||
4. **Read context files**
|
||||
|
||||
Read every file path listed under `contextFiles` from the apply instructions output.
|
||||
The files depend on the schema being used:
|
||||
- **spec-driven**: proposal, specs, design, tasks
|
||||
- Other schemas: follow the contextFiles from CLI output
|
||||
|
||||
5. **Show current progress**
|
||||
|
||||
Display:
|
||||
- Schema being used
|
||||
- Progress: "N/M tasks complete"
|
||||
- Remaining tasks overview
|
||||
- Dynamic instruction from CLI
|
||||
|
||||
6. **Implement tasks (loop until done or blocked)**
|
||||
|
||||
For each pending task:
|
||||
- Show which task is being worked on
|
||||
- Make the code changes required
|
||||
- Keep changes minimal and focused
|
||||
- Mark task complete in the tasks file: `- [ ]` → `- [x]`
|
||||
- Continue to next task
|
||||
|
||||
**Pause if:**
|
||||
- Task is unclear → ask for clarification
|
||||
- Implementation reveals a design issue → suggest updating artifacts
|
||||
- Error or blocker encountered → report and wait for guidance
|
||||
- User interrupts
|
||||
|
||||
7. **On completion or pause, show status**
|
||||
|
||||
Display:
|
||||
- Tasks completed this session
|
||||
- Overall progress: "N/M tasks complete"
|
||||
- If all done: suggest archive
|
||||
- If paused: explain why and wait for guidance
|
||||
|
||||
**Output During Implementation**
|
||||
|
||||
```
|
||||
## Implementing: <change-name> (schema: <schema-name>)
|
||||
|
||||
Working on task 3/7: <task description>
|
||||
[...implementation happening...]
|
||||
✓ Task complete
|
||||
|
||||
Working on task 4/7: <task description>
|
||||
[...implementation happening...]
|
||||
✓ Task complete
|
||||
```
|
||||
|
||||
**Output On Completion**
|
||||
|
||||
```
|
||||
## Implementation Complete
|
||||
|
||||
**Change:** <change-name>
|
||||
**Schema:** <schema-name>
|
||||
**Progress:** 7/7 tasks complete ✓
|
||||
|
||||
### Completed This Session
|
||||
- [x] Task 1
|
||||
- [x] Task 2
|
||||
...
|
||||
|
||||
All tasks complete! Ready to archive this change.
|
||||
```
|
||||
|
||||
**Output On Pause (Issue Encountered)**
|
||||
|
||||
```
|
||||
## Implementation Paused
|
||||
|
||||
**Change:** <change-name>
|
||||
**Schema:** <schema-name>
|
||||
**Progress:** 4/7 tasks complete
|
||||
|
||||
### Issue Encountered
|
||||
<description of the issue>
|
||||
|
||||
**Options:**
|
||||
1. <option 1>
|
||||
2. <option 2>
|
||||
3. Other approach
|
||||
|
||||
What would you like to do?
|
||||
```
|
||||
|
||||
**Guardrails**
|
||||
- Keep going through tasks until done or blocked
|
||||
- Always read context files before starting (from the apply instructions output)
|
||||
- If task is ambiguous, pause and ask before implementing
|
||||
- If implementation reveals issues, pause and suggest artifact updates
|
||||
- Keep code changes minimal and scoped to each task
|
||||
- Update task checkbox immediately after completing each task
|
||||
- Pause on errors, blockers, or unclear requirements - don't guess
|
||||
- Use contextFiles from CLI output, don't assume specific file names
|
||||
|
||||
**Fluid Workflow Integration**
|
||||
|
||||
This skill supports the "actions on a change" model:
|
||||
|
||||
- **Can be invoked anytime**: Before all artifacts are done (if tasks exist), after partial implementation, interleaved with other actions
|
||||
- **Allows artifact updates**: If implementation reveals design issues, suggest updating artifacts - not phase-locked, work fluidly
|
||||
114
.claude/skills/openspec-archive-change/SKILL.md
Normal file
114
.claude/skills/openspec-archive-change/SKILL.md
Normal file
@@ -0,0 +1,114 @@
|
||||
---
|
||||
name: openspec-archive-change
|
||||
description: Archive a completed change in the experimental workflow. Use when the user wants to finalize and archive a change after implementation is complete.
|
||||
license: MIT
|
||||
compatibility: Requires openspec CLI.
|
||||
metadata:
|
||||
author: openspec
|
||||
version: "1.0"
|
||||
generatedBy: "1.3.1"
|
||||
---
|
||||
|
||||
Archive a completed change in the experimental workflow.
|
||||
|
||||
**Input**: Optionally specify a change name. If omitted, check if it can be inferred from conversation context. If vague or ambiguous you MUST prompt for available changes.
|
||||
|
||||
**Steps**
|
||||
|
||||
1. **If no change name provided, prompt for selection**
|
||||
|
||||
Run `openspec list --json` to get available changes. Use the **AskUserQuestion tool** to let the user select.
|
||||
|
||||
Show only active changes (not already archived).
|
||||
Include the schema used for each change if available.
|
||||
|
||||
**IMPORTANT**: Do NOT guess or auto-select a change. Always let the user choose.
|
||||
|
||||
2. **Check artifact completion status**
|
||||
|
||||
Run `openspec status --change "<name>" --json` to check artifact completion.
|
||||
|
||||
Parse the JSON to understand:
|
||||
- `schemaName`: The workflow being used
|
||||
- `artifacts`: List of artifacts with their status (`done` or other)
|
||||
|
||||
**If any artifacts are not `done`:**
|
||||
- Display warning listing incomplete artifacts
|
||||
- Use **AskUserQuestion tool** to confirm user wants to proceed
|
||||
- Proceed if user confirms
|
||||
|
||||
3. **Check task completion status**
|
||||
|
||||
Read the tasks file (typically `tasks.md`) to check for incomplete tasks.
|
||||
|
||||
Count tasks marked with `- [ ]` (incomplete) vs `- [x]` (complete).
|
||||
|
||||
**If incomplete tasks found:**
|
||||
- Display warning showing count of incomplete tasks
|
||||
- Use **AskUserQuestion tool** to confirm user wants to proceed
|
||||
- Proceed if user confirms
|
||||
|
||||
**If no tasks file exists:** Proceed without task-related warning.
|
||||
|
||||
4. **Assess delta spec sync state**
|
||||
|
||||
Check for delta specs at `openspec/changes/<name>/specs/`. If none exist, proceed without sync prompt.
|
||||
|
||||
**If delta specs exist:**
|
||||
- Compare each delta spec with its corresponding main spec at `openspec/specs/<capability>/spec.md`
|
||||
- Determine what changes would be applied (adds, modifications, removals, renames)
|
||||
- Show a combined summary before prompting
|
||||
|
||||
**Prompt options:**
|
||||
- If changes needed: "Sync now (recommended)", "Archive without syncing"
|
||||
- If already synced: "Archive now", "Sync anyway", "Cancel"
|
||||
|
||||
If user chooses sync, use Task tool (subagent_type: "general-purpose", prompt: "Use Skill tool to invoke openspec-sync-specs for change '<name>'. Delta spec analysis: <include the analyzed delta spec summary>"). Proceed to archive regardless of choice.
|
||||
|
||||
5. **Perform the archive**
|
||||
|
||||
Create the archive directory if it doesn't exist:
|
||||
```bash
|
||||
mkdir -p openspec/changes/archive
|
||||
```
|
||||
|
||||
Generate target name using current date: `YYYY-MM-DD-<change-name>`
|
||||
|
||||
**Check if target already exists:**
|
||||
- If yes: Fail with error, suggest renaming existing archive or using different date
|
||||
- If no: Move the change directory to archive
|
||||
|
||||
```bash
|
||||
mv openspec/changes/<name> openspec/changes/archive/YYYY-MM-DD-<name>
|
||||
```
|
||||
|
||||
6. **Display summary**
|
||||
|
||||
Show archive completion summary including:
|
||||
- Change name
|
||||
- Schema that was used
|
||||
- Archive location
|
||||
- Whether specs were synced (if applicable)
|
||||
- Note about any warnings (incomplete artifacts/tasks)
|
||||
|
||||
**Output On Success**
|
||||
|
||||
```
|
||||
## Archive Complete
|
||||
|
||||
**Change:** <change-name>
|
||||
**Schema:** <schema-name>
|
||||
**Archived to:** openspec/changes/archive/YYYY-MM-DD-<name>/
|
||||
**Specs:** ✓ Synced to main specs (or "No delta specs" or "Sync skipped")
|
||||
|
||||
All artifacts complete. All tasks complete.
|
||||
```
|
||||
|
||||
**Guardrails**
|
||||
- Always prompt for change selection if not provided
|
||||
- Use artifact graph (openspec status --json) for completion checking
|
||||
- Don't block archive on warnings - just inform and confirm
|
||||
- Preserve .openspec.yaml when moving to archive (it moves with the directory)
|
||||
- Show clear summary of what happened
|
||||
- If sync is requested, use openspec-sync-specs approach (agent-driven)
|
||||
- If delta specs exist, always run the sync assessment and show the combined summary before prompting
|
||||
288
.claude/skills/openspec-explore/SKILL.md
Normal file
288
.claude/skills/openspec-explore/SKILL.md
Normal file
@@ -0,0 +1,288 @@
|
||||
---
|
||||
name: openspec-explore
|
||||
description: Enter explore mode - a thinking partner for exploring ideas, investigating problems, and clarifying requirements. Use when the user wants to think through something before or during a change.
|
||||
license: MIT
|
||||
compatibility: Requires openspec CLI.
|
||||
metadata:
|
||||
author: openspec
|
||||
version: "1.0"
|
||||
generatedBy: "1.3.1"
|
||||
---
|
||||
|
||||
Enter explore mode. Think deeply. Visualize freely. Follow the conversation wherever it goes.
|
||||
|
||||
**IMPORTANT: Explore mode is for thinking, not implementing.** You may read files, search code, and investigate the codebase, but you must NEVER write code or implement features. If the user asks you to implement something, remind them to exit explore mode first and create a change proposal. You MAY create OpenSpec artifacts (proposals, designs, specs) if the user asks—that's capturing thinking, not implementing.
|
||||
|
||||
**This is a stance, not a workflow.** There are no fixed steps, no required sequence, no mandatory outputs. You're a thinking partner helping the user explore.
|
||||
|
||||
---
|
||||
|
||||
## The Stance
|
||||
|
||||
- **Curious, not prescriptive** - Ask questions that emerge naturally, don't follow a script
|
||||
- **Open threads, not interrogations** - Surface multiple interesting directions and let the user follow what resonates. Don't funnel them through a single path of questions.
|
||||
- **Visual** - Use ASCII diagrams liberally when they'd help clarify thinking
|
||||
- **Adaptive** - Follow interesting threads, pivot when new information emerges
|
||||
- **Patient** - Don't rush to conclusions, let the shape of the problem emerge
|
||||
- **Grounded** - Explore the actual codebase when relevant, don't just theorize
|
||||
|
||||
---
|
||||
|
||||
## What You Might Do
|
||||
|
||||
Depending on what the user brings, you might:
|
||||
|
||||
**Explore the problem space**
|
||||
- Ask clarifying questions that emerge from what they said
|
||||
- Challenge assumptions
|
||||
- Reframe the problem
|
||||
- Find analogies
|
||||
|
||||
**Investigate the codebase**
|
||||
- Map existing architecture relevant to the discussion
|
||||
- Find integration points
|
||||
- Identify patterns already in use
|
||||
- Surface hidden complexity
|
||||
|
||||
**Compare options**
|
||||
- Brainstorm multiple approaches
|
||||
- Build comparison tables
|
||||
- Sketch tradeoffs
|
||||
- Recommend a path (if asked)
|
||||
|
||||
**Visualize**
|
||||
```
|
||||
┌─────────────────────────────────────────┐
|
||||
│ Use ASCII diagrams liberally │
|
||||
├─────────────────────────────────────────┤
|
||||
│ │
|
||||
│ ┌────────┐ ┌────────┐ │
|
||||
│ │ State │────────▶│ State │ │
|
||||
│ │ A │ │ B │ │
|
||||
│ └────────┘ └────────┘ │
|
||||
│ │
|
||||
│ System diagrams, state machines, │
|
||||
│ data flows, architecture sketches, │
|
||||
│ dependency graphs, comparison tables │
|
||||
│ │
|
||||
└─────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
**Surface risks and unknowns**
|
||||
- Identify what could go wrong
|
||||
- Find gaps in understanding
|
||||
- Suggest spikes or investigations
|
||||
|
||||
---
|
||||
|
||||
## OpenSpec Awareness
|
||||
|
||||
You have full context of the OpenSpec system. Use it naturally, don't force it.
|
||||
|
||||
### Check for context
|
||||
|
||||
At the start, quickly check what exists:
|
||||
```bash
|
||||
openspec list --json
|
||||
```
|
||||
|
||||
This tells you:
|
||||
- If there are active changes
|
||||
- Their names, schemas, and status
|
||||
- What the user might be working on
|
||||
|
||||
### When no change exists
|
||||
|
||||
Think freely. When insights crystallize, you might offer:
|
||||
|
||||
- "This feels solid enough to start a change. Want me to create a proposal?"
|
||||
- Or keep exploring - no pressure to formalize
|
||||
|
||||
### When a change exists
|
||||
|
||||
If the user mentions a change or you detect one is relevant:
|
||||
|
||||
1. **Read existing artifacts for context**
|
||||
- `openspec/changes/<name>/proposal.md`
|
||||
- `openspec/changes/<name>/design.md`
|
||||
- `openspec/changes/<name>/tasks.md`
|
||||
- etc.
|
||||
|
||||
2. **Reference them naturally in conversation**
|
||||
- "Your design mentions using Redis, but we just realized SQLite fits better..."
|
||||
- "The proposal scopes this to premium users, but we're now thinking everyone..."
|
||||
|
||||
3. **Offer to capture when decisions are made**
|
||||
|
||||
| Insight Type | Where to Capture |
|
||||
|----------------------------|--------------------------------|
|
||||
| New requirement discovered | `specs/<capability>/spec.md` |
|
||||
| Requirement changed | `specs/<capability>/spec.md` |
|
||||
| Design decision made | `design.md` |
|
||||
| Scope changed | `proposal.md` |
|
||||
| New work identified | `tasks.md` |
|
||||
| Assumption invalidated | Relevant artifact |
|
||||
|
||||
Example offers:
|
||||
- "That's a design decision. Capture it in design.md?"
|
||||
- "This is a new requirement. Add it to specs?"
|
||||
- "This changes scope. Update the proposal?"
|
||||
|
||||
4. **The user decides** - Offer and move on. Don't pressure. Don't auto-capture.
|
||||
|
||||
---
|
||||
|
||||
## What You Don't Have To Do
|
||||
|
||||
- Follow a script
|
||||
- Ask the same questions every time
|
||||
- Produce a specific artifact
|
||||
- Reach a conclusion
|
||||
- Stay on topic if a tangent is valuable
|
||||
- Be brief (this is thinking time)
|
||||
|
||||
---
|
||||
|
||||
## Handling Different Entry Points
|
||||
|
||||
**User brings a vague idea:**
|
||||
```
|
||||
User: I'm thinking about adding real-time collaboration
|
||||
|
||||
You: Real-time collab is a big space. Let me think about this...
|
||||
|
||||
COLLABORATION SPECTRUM
|
||||
════════════════════════════════════════════
|
||||
|
||||
Awareness Coordination Sync
|
||||
│ │ │
|
||||
▼ ▼ ▼
|
||||
┌────────┐ ┌────────┐ ┌────────┐
|
||||
│Presence│ │Cursors │ │ CRDT │
|
||||
│ "3 │ │ Multi │ │Conflict│
|
||||
│online" │ │ select │ │ free │
|
||||
└────────┘ └────────┘ └────────┘
|
||||
│ │ │
|
||||
trivial moderate complex
|
||||
|
||||
Where's your head at?
|
||||
```
|
||||
|
||||
**User brings a specific problem:**
|
||||
```
|
||||
User: The auth system is a mess
|
||||
|
||||
You: [reads codebase]
|
||||
|
||||
┌─────────────────────────────────────────────┐
|
||||
│ CURRENT AUTH FLOW │
|
||||
└─────────────────────────────────────────────┘
|
||||
│
|
||||
┌─────────────┼─────────────┐
|
||||
▼ ▼ ▼
|
||||
┌─────────┐ ┌─────────┐ ┌─────────┐
|
||||
│ Google │ │ GitHub │ │ Email │
|
||||
│ OAuth │ │ OAuth │ │ Magic │
|
||||
└────┬────┘ └────┬────┘ └────┬────┘
|
||||
│ │ │
|
||||
└─────────────┼─────────────┘
|
||||
▼
|
||||
┌───────────┐
|
||||
│ Session │
|
||||
└─────┬─────┘
|
||||
│
|
||||
▼
|
||||
┌───────────┐
|
||||
│ Perms │
|
||||
└───────────┘
|
||||
|
||||
I see three tangles. Which one's burning?
|
||||
```
|
||||
|
||||
**User is stuck mid-implementation:**
|
||||
```
|
||||
User: /opsx:explore add-auth-system
|
||||
The OAuth integration is more complex than expected
|
||||
|
||||
You: [reads change artifacts]
|
||||
|
||||
You're on task 4: "Implement OAuth flow"
|
||||
|
||||
Let me trace what's involved...
|
||||
|
||||
[draws diagram, explores options, suggests paths]
|
||||
|
||||
Want to update the design to reflect this?
|
||||
Or add a spike task to investigate?
|
||||
```
|
||||
|
||||
**User wants to compare options:**
|
||||
```
|
||||
User: Should we use Postgres or SQLite?
|
||||
|
||||
You: Generic answer is boring. What's the context?
|
||||
|
||||
User: A CLI tool that tracks local dev environments
|
||||
|
||||
You: That changes everything.
|
||||
|
||||
┌─────────────────────────────────────────────────┐
|
||||
│ CLI TOOL DATA STORAGE │
|
||||
└─────────────────────────────────────────────────┘
|
||||
|
||||
Key constraints:
|
||||
• No daemon running
|
||||
• Must work offline
|
||||
• Single user
|
||||
|
||||
SQLite Postgres
|
||||
Deployment embedded ✓ needs server ✗
|
||||
Offline yes ✓ no ✗
|
||||
Single file yes ✓ no ✗
|
||||
|
||||
SQLite. Not even close.
|
||||
|
||||
Unless... is there a sync component?
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Ending Discovery
|
||||
|
||||
There's no required ending. Discovery might:
|
||||
|
||||
- **Flow into a proposal**: "Ready to start? I can create a change proposal."
|
||||
- **Result in artifact updates**: "Updated design.md with these decisions"
|
||||
- **Just provide clarity**: User has what they need, moves on
|
||||
- **Continue later**: "We can pick this up anytime"
|
||||
|
||||
When it feels like things are crystallizing, you might summarize:
|
||||
|
||||
```
|
||||
## What We Figured Out
|
||||
|
||||
**The problem**: [crystallized understanding]
|
||||
|
||||
**The approach**: [if one emerged]
|
||||
|
||||
**Open questions**: [if any remain]
|
||||
|
||||
**Next steps** (if ready):
|
||||
- Create a change proposal
|
||||
- Keep exploring: just keep talking
|
||||
```
|
||||
|
||||
But this summary is optional. Sometimes the thinking IS the value.
|
||||
|
||||
---
|
||||
|
||||
## Guardrails
|
||||
|
||||
- **Don't implement** - Never write code or implement features. Creating OpenSpec artifacts is fine, writing application code is not.
|
||||
- **Don't fake understanding** - If something is unclear, dig deeper
|
||||
- **Don't rush** - Discovery is thinking time, not task time
|
||||
- **Don't force structure** - Let patterns emerge naturally
|
||||
- **Don't auto-capture** - Offer to save insights, don't just do it
|
||||
- **Do visualize** - A good diagram is worth many paragraphs
|
||||
- **Do explore the codebase** - Ground discussions in reality
|
||||
- **Do question assumptions** - Including the user's and your own
|
||||
110
.claude/skills/openspec-propose/SKILL.md
Normal file
110
.claude/skills/openspec-propose/SKILL.md
Normal file
@@ -0,0 +1,110 @@
|
||||
---
|
||||
name: openspec-propose
|
||||
description: Propose a new change with all artifacts generated in one step. Use when the user wants to quickly describe what they want to build and get a complete proposal with design, specs, and tasks ready for implementation.
|
||||
license: MIT
|
||||
compatibility: Requires openspec CLI.
|
||||
metadata:
|
||||
author: openspec
|
||||
version: "1.0"
|
||||
generatedBy: "1.3.1"
|
||||
---
|
||||
|
||||
Propose a new change - create the change and generate all artifacts in one step.
|
||||
|
||||
I'll create a change with artifacts:
|
||||
- proposal.md (what & why)
|
||||
- design.md (how)
|
||||
- tasks.md (implementation steps)
|
||||
|
||||
When ready to implement, run /opsx:apply
|
||||
|
||||
---
|
||||
|
||||
**Input**: The user's request should include a change name (kebab-case) OR a description of what they want to build.
|
||||
|
||||
**Steps**
|
||||
|
||||
1. **If no clear input provided, ask what they want to build**
|
||||
|
||||
Use the **AskUserQuestion tool** (open-ended, no preset options) to ask:
|
||||
> "What change do you want to work on? Describe what you want to build or fix."
|
||||
|
||||
From their description, derive a kebab-case name (e.g., "add user authentication" → `add-user-auth`).
|
||||
|
||||
**IMPORTANT**: Do NOT proceed without understanding what the user wants to build.
|
||||
|
||||
2. **Create the change directory**
|
||||
```bash
|
||||
openspec new change "<name>"
|
||||
```
|
||||
This creates a scaffolded change at `openspec/changes/<name>/` with `.openspec.yaml`.
|
||||
|
||||
3. **Get the artifact build order**
|
||||
```bash
|
||||
openspec status --change "<name>" --json
|
||||
```
|
||||
Parse the JSON to get:
|
||||
- `applyRequires`: array of artifact IDs needed before implementation (e.g., `["tasks"]`)
|
||||
- `artifacts`: list of all artifacts with their status and dependencies
|
||||
|
||||
4. **Create artifacts in sequence until apply-ready**
|
||||
|
||||
Use the **TodoWrite tool** to track progress through the artifacts.
|
||||
|
||||
Loop through artifacts in dependency order (artifacts with no pending dependencies first):
|
||||
|
||||
a. **For each artifact that is `ready` (dependencies satisfied)**:
|
||||
- Get instructions:
|
||||
```bash
|
||||
openspec instructions <artifact-id> --change "<name>" --json
|
||||
```
|
||||
- The instructions JSON includes:
|
||||
- `context`: Project background (constraints for you - do NOT include in output)
|
||||
- `rules`: Artifact-specific rules (constraints for you - do NOT include in output)
|
||||
- `template`: The structure to use for your output file
|
||||
- `instruction`: Schema-specific guidance for this artifact type
|
||||
- `outputPath`: Where to write the artifact
|
||||
- `dependencies`: Completed artifacts to read for context
|
||||
- Read any completed dependency files for context
|
||||
- Create the artifact file using `template` as the structure
|
||||
- Apply `context` and `rules` as constraints - but do NOT copy them into the file
|
||||
- Show brief progress: "Created <artifact-id>"
|
||||
|
||||
b. **Continue until all `applyRequires` artifacts are complete**
|
||||
- After creating each artifact, re-run `openspec status --change "<name>" --json`
|
||||
- Check if every artifact ID in `applyRequires` has `status: "done"` in the artifacts array
|
||||
- Stop when all `applyRequires` artifacts are done
|
||||
|
||||
c. **If an artifact requires user input** (unclear context):
|
||||
- Use **AskUserQuestion tool** to clarify
|
||||
- Then continue with creation
|
||||
|
||||
5. **Show final status**
|
||||
```bash
|
||||
openspec status --change "<name>"
|
||||
```
|
||||
|
||||
**Output**
|
||||
|
||||
After completing all artifacts, summarize:
|
||||
- Change name and location
|
||||
- List of artifacts created with brief descriptions
|
||||
- What's ready: "All artifacts created! Ready for implementation."
|
||||
- Prompt: "Run `/opsx:apply` or ask me to implement to start working on the tasks."
|
||||
|
||||
**Artifact Creation Guidelines**
|
||||
|
||||
- Follow the `instruction` field from `openspec instructions` for each artifact type
|
||||
- The schema defines what each artifact should contain - follow it
|
||||
- Read dependency artifacts for context before creating new ones
|
||||
- Use `template` as the structure for your output file - fill in its sections
|
||||
- **IMPORTANT**: `context` and `rules` are constraints for YOU, not content for the file
|
||||
- Do NOT copy `<context>`, `<rules>`, `<project_context>` blocks into the artifact
|
||||
- These guide what you write, but should never appear in the output
|
||||
|
||||
**Guardrails**
|
||||
- Create ALL artifacts needed for implementation (as defined by schema's `apply.requires`)
|
||||
- Always read dependency artifacts before creating a new one
|
||||
- If context is critically unclear, ask the user - but prefer making reasonable decisions to keep momentum
|
||||
- If a change with that name already exists, ask if user wants to continue it or create a new one
|
||||
- Verify each artifact file exists after writing before proceeding to next
|
||||
54
.env.example
54
.env.example
@@ -11,6 +11,22 @@ MLFLOW_PORT="80"
|
||||
MLFLOW_USERNAME="aignosi"
|
||||
MLFLOW_PASSWORD="mlflow_password"
|
||||
|
||||
# Worker runtime name for PluginStore.install_runtime (PredictionsBatch / MinimalRetrain workers).
|
||||
RUNTIME="single"
|
||||
|
||||
# Plugin store (Git-backed catalog + runtime install).
|
||||
STORE_BASE_URL="http://gitea.sientia.svc.cluster.local:3000"
|
||||
STORE_OWNER="sientia"
|
||||
STORE_REPO="model-library-store"
|
||||
STORE_BRANCH="main"
|
||||
STORE_USERNAME=""
|
||||
STORE_PASSWORD=""
|
||||
STORE_CACHE_TTL_SECONDS=""
|
||||
|
||||
PYPI_SERVER="http://library-distribution-server.library.svc.cluster.local:5000"
|
||||
PYPI_USERNAME=""
|
||||
PYPI_PASSWORD=""
|
||||
|
||||
OPC_ID="1"
|
||||
OPC_URL="opc.tcp://sientia-opc-simulator-opc.sientia.svc.cluster.local:4840"
|
||||
|
||||
@@ -33,41 +49,3 @@ MINIO_ACCESS_KEY="sientia"
|
||||
MINIO_SECRET_KEY="sientia"
|
||||
MINIO_REGION_NAME="sa-east-1"
|
||||
MINIO_DEFAULT_BUCKET="sientia"
|
||||
|
||||
# Suffix for every Temporal task queue name; required, non-empty.
|
||||
# The import worker listens on "import_model-${RUNTIME}-queue".
|
||||
RUNTIME="local"
|
||||
|
||||
# --- Model import (.sientia bundles) -----------------------------------------
|
||||
# The import worker is registered only when IMPORT_MODEL_ENABLED is truthy
|
||||
# (1/true/yes/on). Do not turn it on before the V12 migration is applied to
|
||||
# public.experiment_run in the BFF database.
|
||||
IMPORT_MODEL_ENABLED="false"
|
||||
|
||||
# base64 of the 32-byte AES-256-GCM key that wraps bundle passwords.
|
||||
# Generate a throwaway one with: openssl rand -base64 32
|
||||
IMPORT_PASSWORD_KEY="replace-with-base64-of-32-random-bytes"
|
||||
|
||||
# Where the frontend uploads the bundle. The object key must sit under the prefix.
|
||||
IMPORT_BUNDLE_BUCKET="sientia"
|
||||
IMPORT_BUNDLE_PREFIX="imported_models/"
|
||||
|
||||
# MongoDB collection holding the model listing the import writes its document into.
|
||||
IMPORT_MODELS_COLLECTION="models"
|
||||
|
||||
# The import log lives in the Spring Boot BFF database on the POSTGRES_* server
|
||||
# above: same host, port, user and password, another database name. Only that
|
||||
# name is configuration here; the schema (public) and the table (experiment_run)
|
||||
# are constants of the code. Set IMPORT_STATUS_DB_{HOST,PORT,USER,PASSWORD} only
|
||||
# if the log really is on a different server than POSTGRES_* — they override it.
|
||||
IMPORT_STATUS_DB_NAME="sientia-core-mlops-bff"
|
||||
|
||||
# --- scripts/import_model_manual_run.py only ---------------------------------
|
||||
# Inputs for the manual run that stands in for the frontend. The script writes
|
||||
# one object and one row to whatever the variables above name.
|
||||
IMPORT_SCRIPT_BUNDLE_PATH="tests/fixtures/bundle/golden_model_v3.sientia"
|
||||
IMPORT_SCRIPT_BUNDLE_PASSWORD="the-password-the-bundle-was-written-with"
|
||||
IMPORT_SCRIPT_USERNAME="manual.run@example.com"
|
||||
# Cell 2's guard: restate POSTGRES_HOST here (or in the cell) after reading what cell 1
|
||||
# printed. Left empty on purpose — the script refuses to write until it matches.
|
||||
IMPORT_SCRIPT_CONFIRM_TARGET=""
|
||||
4
.github/workflows/quality-gate.yml
vendored
4
.github/workflows/quality-gate.yml
vendored
@@ -26,6 +26,6 @@ jobs:
|
||||
issues: write
|
||||
with:
|
||||
project_name: 'laborious'
|
||||
repositories: 'sientia-dataops-library, sientia-mlops-library'
|
||||
requirements_file: 'requirements-light.txt'
|
||||
repositories: 'sientia-dataops-library,sientia-model-library'
|
||||
requirements_file: 'requirements.txt'
|
||||
secrets: inherit
|
||||
2
.gitignore
vendored
2
.gitignore
vendored
@@ -56,4 +56,4 @@ mlruns/
|
||||
relatorio*
|
||||
openspec/
|
||||
.cursor/
|
||||
.claude/
|
||||
.claude/
|
||||
|
||||
310
README.md
310
README.md
@@ -21,7 +21,6 @@ A comprehensive, Temporal-based ML orchestration system for industrial data proc
|
||||
- [Minimal Retrain Workflow](#4-minimal-retrain-workflow-minimal_retrainpy)
|
||||
- [Drift Workflow](#5-drift-workflow-driftpy)
|
||||
- [Simple Metrics Workflow](#6-simple-metrics-workflow-simple_metricspy)
|
||||
- [Import Model Workflow](#7-import-model-workflow-import_modelpy)
|
||||
- [Installation & Setup](#installation--setup)
|
||||
- [Prerequisites](#prerequisites)
|
||||
- [Environment Setup](#environment-setup)
|
||||
@@ -32,7 +31,6 @@ A comprehensive, Temporal-based ML orchestration system for industrial data proc
|
||||
- [Running Tests and Coverage](#running-tests-and-coverage)
|
||||
- [Manual Test Execution](#manual-test-execution)
|
||||
- [Manual Application Execution](#manual-application-execution)
|
||||
- [Running One Model Import by Hand](#running-one-model-import-by-hand)
|
||||
- [Code Quality & Validation](#code-quality--validation)
|
||||
- [Overview](#overview)
|
||||
- [Validation Tools](#validation-tools)
|
||||
@@ -129,25 +127,19 @@ Laborious uses a Temporal-based architecture with strong separation of concerns
|
||||
### Key Components
|
||||
|
||||
#### **Worker (`laborious/worker/worker.py`)**
|
||||
- Temporal client setup, four workers via `sientia_do.temporal.worker.prepare_worker`, plus a
|
||||
fifth registered only when `IMPORT_MODEL_ENABLED` is truthy (`1`, `true`, `yes`, `on`)
|
||||
- Runtime-scoped task queues: `{workflow}-{RUNTIME}-queue` for all workflows
|
||||
- Temporal client setup, four workers via `sientia_do.temporal.worker.prepare_worker`
|
||||
- `RUNTIME` is required: installs the PluginStore runtime and suffixes the model-loading queues
|
||||
- Metrics server initialization, notification handler setup
|
||||
- Graceful shutdown and autoscaling-friendly behavior
|
||||
|
||||
**Breaking (schedulers):** drift and simple_metrics queues are no longer `drift-queue` /
|
||||
`simple_metrics-queue`. Use `drift-{RUNTIME}-queue` and `simple_metrics-{RUNTIME}-queue`
|
||||
matching the worker pod `RUNTIME` env (same as `predictions_batch` / `minimal_retrain`).
|
||||
|
||||
Task queues, one per worker:
|
||||
|
||||
| Queue | Worker | Registered |
|
||||
|---|---|---|
|
||||
| `predictions_batch-{RUNTIME}-queue` | PredictionsBatch (+ its two sub-workflows) | always |
|
||||
| `minimal_retrain-{RUNTIME}-queue` | MinimalRetrain | always |
|
||||
| `drift-{RUNTIME}-queue` | Drift | always |
|
||||
| `simple_metrics-{RUNTIME}-queue` | SimpleMetrics | always |
|
||||
| `import_model-{RUNTIME}-queue` | ImportModel | only when `IMPORT_MODEL_ENABLED` is truthy |
|
||||
| Queue | Worker |
|
||||
|---|---|
|
||||
| `predictions_batch-{RUNTIME}-queue` | PredictionsBatch (+ its two sub-workflows) |
|
||||
| `minimal_retrain-{RUNTIME}-queue` | MinimalRetrain |
|
||||
| `drift-core-queue` | Drift (pinned to `core`: no model code loaded) |
|
||||
| `simple_metrics-core-queue` | SimpleMetrics (pinned to `core`: no model code loaded) |
|
||||
|
||||
#### **Workflows (`laborious/workflows/`)**
|
||||
- `predictions_batch.py`: Batch prediction entry point
|
||||
@@ -156,8 +148,6 @@ Task queues, one per worker:
|
||||
- `minimal_retrain.py`: Automated model retraining and production update
|
||||
- `drift.py`: Data drift detection and monitoring
|
||||
- `simple_metrics.py`: Regression metrics calculation (RMSE, MSE, MAE, R²)
|
||||
- `import_model.py`: `.sientia` bundle import — registered model + model document, behind
|
||||
`IMPORT_MODEL_ENABLED`
|
||||
|
||||
#### **Activities (`laborious/activities/`)**
|
||||
- `gates.py`: Data quality validation, filtering, and data formatting operations
|
||||
@@ -185,7 +175,7 @@ Task queues, one per worker:
|
||||
#### **Data Services (`laborious/utils/`)**
|
||||
- `connectors_config.py`: Env-driven configuration builders
|
||||
- `models/minio_dataframe_payload.py`: MinIO-offloaded DataFrame payload model
|
||||
- `repository/model_repository.py`: MLFlow operations and retraining
|
||||
- ML models are loaded via `SientiaMLflowRepository` (wrapper-based, `@production` alias) constructed in `Activities` from `build_mlflow_config()`.
|
||||
- `repository/opc_repository.py`: OPC UA client, writes, session recovery (see [OPC UA Communication](#opc-ua-communication))
|
||||
- `repository/minio_manager.py`: MinIO object storage operations
|
||||
- `filters/conditional_filters.py` and `filters/mlflow_filters.py`
|
||||
@@ -485,7 +475,7 @@ flowchart LR
|
||||
"source_table_name": "laborious_data",
|
||||
"target_table_name": "drift_metrics",
|
||||
"interval": 60,
|
||||
"model_config": { "target": "temperature" },
|
||||
"model_config": { "target": "temperature", "alias": "production" },
|
||||
"drift_metrics": ["kolmogorov_smirnov", "jensen_shannon", "wasserstein"],
|
||||
"chunk_period": "min"
|
||||
}
|
||||
@@ -520,60 +510,11 @@ flowchart LR
|
||||
"data_table_name": "laborious_data",
|
||||
"target_table_name": "simple_metrics",
|
||||
"interval_minutes": 60,
|
||||
"model_config": { "target": "temperature" },
|
||||
"model_config": { "target": "temperature", "alias": "production" },
|
||||
"metrics": ["rmse", "mse", "mae", "r2"]
|
||||
}
|
||||
```
|
||||
|
||||
### 7. Import Model Workflow (`import_model.py`)
|
||||
|
||||
The **ImportModel** workflow turns an encrypted `.sientia` bundle already uploaded to MinIO into a
|
||||
registered MLflow model plus one document in the MongoDB `models` collection — or into a durable
|
||||
`ERROR` row carrying a sentence a person can read. It is registered **only** when
|
||||
`IMPORT_MODEL_ENABLED` is set, so the importer can live in its own runtime with its own credentials.
|
||||
|
||||
Full contract for the frontend and the BFF: [`docs/model-import.md`](docs/model-import.md).
|
||||
The file format: [`docs/sientia-bundle-format.md`](docs/sientia-bundle-format.md).
|
||||
|
||||
#### Execution Flow
|
||||
1. **Claim**: takes ownership of the `public.experiment_run` row the BFF inserted (it never inserts one)
|
||||
2. **Download**: object limits, then the mandatory `expected_digest` comparison
|
||||
3. **Open**: seven ordered gates — header, decryption, archive inspection, extraction, structure, content policy
|
||||
4. **Provision**: experiment, a new run with the recorded name, the artifact tree, a registered version (stage untouched)
|
||||
5. **List**: the MongoDB model document, `active: false`
|
||||
6. **Finish**: cleanup, then one terminal write classifying the outcome once
|
||||
|
||||
#### Architecture Diagram
|
||||
```mermaid
|
||||
flowchart LR
|
||||
A[1. claim_import_status] --> B[2. download_import_bundle] --> C[3. open_import_bundle]
|
||||
C --> D[4. create_import_experiment] --> E[5. upload_import_artifacts]
|
||||
E --> F[6. register_import_model_version] --> G[7. write_import_model_document]
|
||||
G --> H[8. cleanup_import_files] --> I[9. record_import_terminal_status]
|
||||
|
||||
A -.-> BFF[(sientia-core-mlops-bff)]
|
||||
B -.-> MinIO[(MinIO)]
|
||||
F -.-> MLflow[(MLflow)]
|
||||
G -.-> Mongo[(MongoDB)]
|
||||
I -.-> BFF
|
||||
```
|
||||
|
||||
#### Input Parameters
|
||||
```json
|
||||
{
|
||||
"import_run_id": 41,
|
||||
"object_key": "imported_models/sales_forecast_v3.sientia",
|
||||
"expected_digest": "9f86d081884c7d659a2feaa0c55ad015a3bf4f1b2b0b822cd15d6c15b0f00a08",
|
||||
"password_envelope": "Base64OfNonceCiphertextTag==",
|
||||
"bucket": "sientia"
|
||||
}
|
||||
```
|
||||
|
||||
A successful import leaves the registered version **unstaged** and the model document **inactive**:
|
||||
promoting and activating are manual operator steps (QTZPOC-21). The `V12` migration of
|
||||
`sientia-core-mlops-bff` must be applied before enabling the worker — see
|
||||
[`docs/model-import.md`](docs/model-import.md) § 7.
|
||||
|
||||
## 📋 Prerequisites
|
||||
|
||||
- Python 3.11+
|
||||
@@ -719,62 +660,6 @@ fi
|
||||
python -m laborious.worker.worker
|
||||
```
|
||||
|
||||
### Running One Model Import by Hand
|
||||
|
||||
[`scripts/import_model_manual_run.py`](scripts/import_model_manual_run.py) stands in for the frontend
|
||||
that does not exist yet. The Angular screen plus the Spring Boot BFF will, between them, hash a
|
||||
`.sientia` file, wrap the bundle password, upload the object and insert one `PENDING` row before
|
||||
starting `import_model`; this script does exactly those four things against whatever `.env` points
|
||||
at, then reads the outcome back and can undo itself.
|
||||
|
||||
It is a **development tool**. It writes to whatever `.env` names — a real object in a real bucket and
|
||||
a real row in a real `public.experiment_run` — and a `.env` copied from a cluster deployment can
|
||||
perfectly well name production. It is not for production use.
|
||||
|
||||
**Prerequisites**
|
||||
|
||||
- A reachable environment: the BFF Postgres, MinIO, MongoDB, MLflow and Temporal that `.env` names.
|
||||
- The `V12` migration applied to `public.experiment_run` in the BFF database. Without it the import
|
||||
provisions the model and then cannot record the outcome — see
|
||||
[`docs/model-import.md`](docs/model-import.md) § 7.
|
||||
- A worker running with `IMPORT_MODEL_ENABLED` truthy and the same `RUNTIME`, if you want cell 8 (or
|
||||
the Temporal UI) to actually execute the run.
|
||||
- The import variables from [`.env.example`](.env.example) filled in, including `IMPORT_PASSWORD_KEY`
|
||||
and the script's own `IMPORT_SCRIPT_BUNDLE_PATH`, `IMPORT_SCRIPT_BUNDLE_PASSWORD` and
|
||||
`IMPORT_SCRIPT_USERNAME`.
|
||||
|
||||
**Running it.** The file is split into `# %%` cells: run them one at a time in the VS Code
|
||||
interactive window, or run the module top to bottom.
|
||||
|
||||
```bash
|
||||
source ./venv/bin/activate
|
||||
|
||||
# top to bottom (cells 5 and 6 write; cell 2 must pass first)
|
||||
python scripts/import_model_manual_run.py
|
||||
```
|
||||
|
||||
| Cell | Does | Writes? |
|
||||
|---|---|---|
|
||||
| 1 | Resolves and prints every target — Postgres, MinIO, Mongo, MLflow, Temporal, the task queue | no |
|
||||
| 2 | **Confirmation guard** | no |
|
||||
| 3 | Picks the bundle, computes `expected_digest`, prints the parsed header | no |
|
||||
| 4 | Builds the AES-256-GCM password envelope | no |
|
||||
| 5 | Uploads the object, re-reads it and checks the digest still matches | **yes** |
|
||||
| 6 | Inserts the `PENDING` / `run_type='IMPORT'` row, prints `import_run_id` | **yes** |
|
||||
| 7 | Prints the workflow input as JSON, ready for the Temporal UI | no |
|
||||
| 8 | Optionally starts `import_model` and waits (`START_THE_WORKFLOW`) | starts a run |
|
||||
| 9 | Reads the outcome back: the row, the MLflow experiment/run/version, the model document | no |
|
||||
| 10 | Undo: deletes the object and the row always; the provisioning behind `DELETE_PROVISIONING` | deletes |
|
||||
|
||||
**The guard.** Cells 5 and 6 are the only cells that write, and neither runs until cell 2 passes.
|
||||
Cell 2 compares a confirmation value against the Postgres host cell 1 resolved: set
|
||||
`IMPORT_SCRIPT_CONFIRM_TARGET` in `.env`, or paste the host into the cell. It is a read-then-restate
|
||||
rather than a boolean on purpose — a flag can be flipped without ever looking at what it points at.
|
||||
|
||||
No credential is written in the script and none is printed: it prints hosts, ports, database, bucket
|
||||
and object names, the row id and the workflow input, and never the bundle password, the envelope key
|
||||
or the envelope's plaintext.
|
||||
|
||||
## Code Quality & Validation
|
||||
|
||||
### Overview
|
||||
@@ -857,7 +742,6 @@ tests/
|
||||
│ │ ├── test_conditional_filters.py
|
||||
│ │ └── test_mlflow_filters.py
|
||||
│ └── repository/
|
||||
│ ├── test_model_repository.py
|
||||
│ └── test_opc_repository.py
|
||||
```
|
||||
|
||||
@@ -885,9 +769,8 @@ pytest tests/laborious/workflows/test_predictions_batch.py
|
||||
### End-to-End Tests (`e2e/`)
|
||||
|
||||
The `e2e/` suite runs against real dependencies in Docker containers via
|
||||
[testcontainers](https://testcontainers.com/): PostgreSQL, MinIO, MongoDB, and — for the model-import
|
||||
suite — an `mlflow server` subprocess started from this repository's venv. It lives outside
|
||||
`testpaths`, so it is never collected by the unit gate.
|
||||
[testcontainers](https://testcontainers.com/): PostgreSQL and MinIO. It lives outside `testpaths`,
|
||||
so it is never collected by the unit gate.
|
||||
|
||||
**One-time setup, required before running anything under `e2e/`:**
|
||||
|
||||
@@ -924,13 +807,10 @@ docker ps -a --filter "label=org.testcontainers=true" -q | xargs -r docker rm -f
|
||||
source ./venv/bin/activate
|
||||
|
||||
# the predictions and OPC suites
|
||||
pytest e2e -m "not import_e2e"
|
||||
|
||||
# the .sientia model-import suite (opt-in; see e2e/scenarios.md § 4)
|
||||
pytest e2e -m import_e2e
|
||||
pytest e2e/ --override-ini testpaths=e2e -m integration
|
||||
```
|
||||
|
||||
Full scenario catalogue and the import suite's fidelity contract:
|
||||
Full scenario catalogue:
|
||||
[`e2e/scenarios.md`](e2e/scenarios.md).
|
||||
|
||||
## 📊 Monitoring and Metrics
|
||||
@@ -991,20 +871,20 @@ See [OPC UA Communication](#opc-ua-communication) for semantics, concurrency, an
|
||||
| `MLFLOW_PORT` | MLFlow server port | `5080` | Yes |
|
||||
| `MLFLOW_USERNAME` | MLFlow username | `aignosi` | Yes |
|
||||
| `MLFLOW_PASSWORD` | MLFlow password | `aignosi` | Yes |
|
||||
| `IMPORT_MODEL_ENABLED` | Register the import worker on `import_model-{RUNTIME}-queue` | _(unset)_ | No |
|
||||
| `IMPORT_PASSWORD_KEY` | base64 of the 32-byte AES-256-GCM key wrapping bundle passwords | _(none)_ | If importing |
|
||||
| `IMPORT_BUNDLE_BUCKET` | Bucket holding uploaded `.sientia` objects | `sientia` | No |
|
||||
| `IMPORT_BUNDLE_PREFIX` | Key prefix the uploads must sit under | `imported_models/` | No |
|
||||
| `IMPORT_BUNDLE_RETENTION_DAYS` | Lifecycle expiry applied to that prefix | `7` | No |
|
||||
| `IMPORT_WORK_DIR` | Worker scratch directory for bundles | `<tempdir>/sientia-import` | No |
|
||||
| `IMPORT_MAX_OBJECT_BYTES` | Gate 1 object size ceiling | `1073741824` | No |
|
||||
| `IMPORT_MAX_ARCHIVE_ENTRIES` | Gate 4 entry count ceiling | `5000` | No |
|
||||
| `IMPORT_MAX_UNCOMPRESSED_BYTES` | Gate 4 uncompressed size ceiling | `4294967296` | No |
|
||||
| `IMPORT_MAX_COMPRESSION_RATIO` | Gate 4 compression ratio ceiling | `200` | No |
|
||||
| `IMPORT_MODELS_COLLECTION` | MongoDB collection holding the model listing | `models` | No |
|
||||
| `IMPORT_STATUS_DB_NAME` | Import log database name (the BFF's, on the `POSTGRES_*` server) | `sientia-core-mlops-bff` | If importing |
|
||||
| `IMPORT_STATUS_DB_{HOST,PORT,USER,PASSWORD}` | Override, only when the import log is not on the `POSTGRES_*` server | the `POSTGRES_*` value | No |
|
||||
| `RUNTIME` | Plugin store runtime name installed at worker boot (required for model-loading workers) | `single` | Yes* |
|
||||
| `STORE_BASE_URL` | Plugin store Git server base URL | `http://localhost:3000` | Yes* |
|
||||
| `STORE_OWNER` | Plugin store repository owner | `sientia` | Yes* |
|
||||
| `STORE_REPO` | Plugin store repository name | `model-library-store` | Yes* |
|
||||
| `STORE_BRANCH` | Optional branch for the store repository | `main` | No |
|
||||
| `STORE_USERNAME` | HTTP username for the Git store | `None` | No |
|
||||
| `STORE_PASSWORD` | HTTP password/token for the Git store | `None` | No |
|
||||
| `STORE_CACHE_TTL_SECONDS` | Optional cache TTL for store metadata | `None` | No |
|
||||
| `PYPI_SERVER` | Private PyPI index URL for runtime wheels | `http://localhost:5000` | Yes* |
|
||||
| `PYPI_USERNAME` | Optional PyPI basic-auth username | `None` | No |
|
||||
| `PYPI_PASSWORD` | Optional PyPI basic-auth password | `None` | No |
|
||||
| `OPC_CONFIG` | OPC server configuration (JSON) | `{}` | No |
|
||||
|
||||
\* `RUNTIME`, PluginStore (`STORE_*`), and `PYPI_SERVER` are required for workers that install a runtime and load `SientiaModel` wrappers (`PredictionsBatch`, `MinimalRetrain`). Workers that only run drift/simple-metrics style jobs may omit them when those workflows are deployed separately.
|
||||
| `OPC_ID` | OPC server identifier | `1` | No |
|
||||
| `OPC_URL` | OPC server URL | `opc.tcp://localhost:4840` | No |
|
||||
| `OPC_SERVER_URI` | OPC server URI | `opc.tcp://localhost:4840` | No |
|
||||
@@ -1024,7 +904,7 @@ See [OPC UA Communication](#opc-ua-communication) for semantics, concurrency, an
|
||||
| `PROJECT_NAME` | Project name for metrics | `laborious` | No |
|
||||
| `HTTP_METRICS_PORT` | Prometheus metrics port | `9090` | No |
|
||||
| `HTTP_SDK_METRICS_PORT` | Temporal SDK metrics port | `9091` | No |
|
||||
| `POD_ID` | Kubernetes pod identifier | `None` | No |
|
||||
| `HOSTNAME` | Pod identifier (set by Kubernetes) used in metrics labels | `None` | No |
|
||||
| `MINIO_ENDPOINT_URL` | MinIO endpoint URL | `http://localhost:9000` | Yes |
|
||||
| `MINIO_ACCESS_KEY` | MinIO access key | `minioadmin` | Yes |
|
||||
| `MINIO_SECRET_KEY` | MinIO secret key | `minioadmin` | Yes |
|
||||
@@ -1104,33 +984,6 @@ Legacy MinIO object layout (relative key):
|
||||
`training_datasets/{model_name}/{object_prefix}_{timestamp}.parquet` where `object_prefix` is sanitized
|
||||
(slashes replaced by underscores) to keep a stable model-level directory.
|
||||
|
||||
## Model Flavors
|
||||
|
||||
`model_config.predict_flavor` / `model_config.transform_flavor` accept `sklearn`, `pyfunc` and
|
||||
`joblib`. The flavor is configuration, never discovered from the artifact, and there is
|
||||
no fallback between flavors: a failing loader propagates.
|
||||
|
||||
`model_config` also carries ad-hoc behavioral flags, unrelated to flavor, e.g.
|
||||
`skip_predict_timestamp` (default `False`): when `True`, `request_predict` does not reinsert the
|
||||
`timestamp` column derived from the index into the transformed data before calling predict. Full
|
||||
list: [docs/model-flavors.md § Other `model_config` flags](docs/model-flavors.md#other-model_config-flags).
|
||||
|
||||
`joblib` reads the artifact directly — download, `<artifact>/model.pkl`, `joblib.load` — and writes on
|
||||
retrain (`model.pkl` + `code/`, no `MLmodel`, so only this runtime reads it back).
|
||||
|
||||
Every flavor puts the artifact's own `code/` on `sys.path` before loading, **when that directory
|
||||
exists**: MLflow's loaders only do it when the `MLmodel` manifest declares `code`. `download_model`
|
||||
owns that step for all of them — it asks the registry whether the artifact carries `code/` and, only
|
||||
then, downloads that one directory. The models themselves are never pre-downloaded: `sklearn` and
|
||||
`pyfunc` load straight from the model URI and fetch what they need.
|
||||
|
||||
`download_model` does not branch on the flavor: past that shared first step it delegates to
|
||||
`load_predict_model` / `load_transform_model`, and its only branch left is `load_wrapper`, which is
|
||||
orthogonal to the flavor. It returns `(model, artifact_dir)`, the directory the retrain path re-logs
|
||||
`code/` from.
|
||||
|
||||
Full reference, including how to switch an existing model over: **[docs/model-flavors.md](docs/model-flavors.md)**.
|
||||
|
||||
## OPC UA Communication
|
||||
|
||||
Full reference: **[docs/opc-communication.md](docs/opc-communication.md)** (connection lifecycle, Tier-1 `Bad*` reconnect, connection lock / session readiness, metrics, PostgreSQL confidence **12** vs **14**, tests).
|
||||
@@ -1175,7 +1028,7 @@ For single OPC server, use individual environment variables:
|
||||
|
||||
### PI Web API Configuration
|
||||
|
||||
PI Web API configuration is built from environment variables using the `build_api_config` function from `sientia_do.connectors_config`. The configuration includes:
|
||||
PI Web API configuration is built from environment variables using the `build_api_config` function from `sientia_do.utils.connectors_config`. The configuration includes:
|
||||
|
||||
- `PI_WEB_API_BASE_URL`: Base URL of the PI Web API server
|
||||
- `PI_WEB_API_AUTH_TYPE`: Authentication type ('basic' or 'bearer')
|
||||
@@ -1208,9 +1061,30 @@ Where:
|
||||
|
||||
MongoDB pipeline configuration:
|
||||
|
||||
#### Predictions Batch Workflow configuration sample
|
||||
#### MongoDB input samples (updated)
|
||||
|
||||
This is the configuration for the Predictions Batch Workflow, to be inserted into the MongoDB pipeline collection.
|
||||
Updated examples are available in `input_sample.json` at the repository root.
|
||||
The sample already reflects the runtime-aware and alias-based flow:
|
||||
|
||||
- `model_config` uses `target`, `retention_minutes`, and `alias`.
|
||||
- `transform_flavor` / `predict_flavor` are not used anymore.
|
||||
|
||||
Example model document:
|
||||
|
||||
```json
|
||||
{
|
||||
"id": "4",
|
||||
"name": "vcm-nox",
|
||||
"active": false,
|
||||
"model_config": {
|
||||
"alias": "production",
|
||||
"retention_minutes": 60,
|
||||
"target": "CI-W3W01A3"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Example predictions_batch schedule document:
|
||||
|
||||
```json
|
||||
{
|
||||
@@ -1220,41 +1094,67 @@ This is the configuration for the Predictions Batch Workflow, to be inserted int
|
||||
"frequency": "30s",
|
||||
"max_retry_policy": 1,
|
||||
"query": "select * from sientia_data.laborious_data where model_id = 1 and \"timestamp\" > NOW() - INTERVAL '5 minutes' order by \"timestamp\" desc limit 30;",
|
||||
"retention_time": 60,
|
||||
"write_tags": [
|
||||
{
|
||||
"server_id": "server1",
|
||||
"server_id": "1",
|
||||
"type": "prediction",
|
||||
"addr": "ns=2;i=5",
|
||||
"data_type": "double"
|
||||
},
|
||||
{
|
||||
"server_id": "server1",
|
||||
"server_id": "1",
|
||||
"type": "confidence",
|
||||
"addr": "ns=2;i=6",
|
||||
"addr": "ns=2;i=5",
|
||||
"data_type": "double"
|
||||
}
|
||||
],
|
||||
"input_filters": {
|
||||
"EMPTY_DATA": {"POLICY": "STOP"},
|
||||
"SPECIFIC_VARIABLES_NULL_VALUES": {
|
||||
"POLICY": "CONTINUE",
|
||||
"config": {"variables": ["Counter"]}
|
||||
"input_filters": [
|
||||
{
|
||||
"filter_name": "EMPTY_DATA",
|
||||
"policy": "STOP"
|
||||
},
|
||||
{
|
||||
"filter_name": "SPECIFIC_VARIABLES_NULL_VALUES",
|
||||
"policy": "CONTINUE",
|
||||
"config": {
|
||||
"variables": ["Counter"]
|
||||
}
|
||||
}
|
||||
},
|
||||
"mlflow_transform_filters": {
|
||||
"API_ERROR": {"POLICY": "REPEAT"},
|
||||
"NAN_VALUES": {"POLICY": "STOP"}
|
||||
},
|
||||
"mlflow_predict_filters": {
|
||||
"API_ERROR": {"POLICY": "CONTINUE"}
|
||||
},
|
||||
],
|
||||
"mlflow_transform_filters": [
|
||||
{
|
||||
"filter_name": "API_ERROR",
|
||||
"policy": "REPEAT"
|
||||
},
|
||||
{
|
||||
"filter_name": "NAN_VALUES",
|
||||
"policy": "STOP"
|
||||
}
|
||||
],
|
||||
"mlflow_predict_filters": [
|
||||
{
|
||||
"filter_name": "API_ERROR",
|
||||
"policy": "CONTINUE"
|
||||
}
|
||||
],
|
||||
"path_priority": ["STOP", "CONTINUE", "REPEAT"],
|
||||
"active": true,
|
||||
"updated_at": {
|
||||
"$date": "2025-09-16T10:00:00.000Z"
|
||||
"$date": "2026-01-27T17:35:01.600Z"
|
||||
},
|
||||
"datetime_columns": ["timestamp", "created_at"],
|
||||
"predictions_storage_policy": "lts:1"
|
||||
"save_transform": false,
|
||||
"pi_web_api_output_config": {
|
||||
"endpoint": "/streamsets/value",
|
||||
"prediction_tags": {},
|
||||
"confidence_tags": {}
|
||||
},
|
||||
"model_config": {
|
||||
"alias": "production",
|
||||
"retention_minutes": 60,
|
||||
"target": "CI-W3W01A3"
|
||||
},
|
||||
"datetime_columns": ["timestamp", "created_at"]
|
||||
}
|
||||
```
|
||||
|
||||
@@ -1272,11 +1172,8 @@ This is the configuration created by the Orchestrator in Temporal.
|
||||
"EMPTY_DATA":{"config":{},"policy":"STOP"}
|
||||
},
|
||||
"model_config":{
|
||||
"is_compressed":true,
|
||||
"predict_flavor":"pyfunc",
|
||||
"retention_minutes":60,
|
||||
"retention_target":"artifact",
|
||||
"transform_function_keyword":"transform"
|
||||
"target":"sensor_or_label_column",
|
||||
"retention_minutes":60
|
||||
},
|
||||
"model_id":"352",
|
||||
"model_name":"courier",
|
||||
@@ -1315,7 +1212,7 @@ laborious/
|
||||
│ ├── prediction_process.py # Core prediction workflow
|
||||
│ └── format_and_export_prediction.py # Export workflow
|
||||
├── worker/ # Worker implementation
|
||||
│ └── worker.py # Main worker orchestrator (uses sientia_do prepare_worker)
|
||||
│ └── worker.py # Entrypoint; workers built via `sientia_do.temporal.worker.prepare_worker`
|
||||
├── utils/ # Utility functions
|
||||
│ ├── connectors_config.py # Environment-driven config builders
|
||||
│ ├── models/ # Data models
|
||||
@@ -1324,7 +1221,6 @@ laborious/
|
||||
│ │ ├── conditional_filters.py # Conditional data filters
|
||||
│ │ └── mlflow_filters.py # MLFlow response filters
|
||||
│ └── repository/ # Data access layer
|
||||
│ ├── model_repository.py # MLFlow model operations
|
||||
│ ├── opc_repository.py # OPC server operations
|
||||
│ └── minio_manager.py # MinIO object storage operations
|
||||
├── metrics.py # Prometheus metrics definitions
|
||||
@@ -1351,7 +1247,7 @@ laborious/
|
||||
2. **MLFlow Connection Issues**
|
||||
- Verify MLFlow server is running and accessible
|
||||
- Check authentication credentials and permissions
|
||||
- Ensure model names and versions exist
|
||||
- Ensure model names exist and the expected alias (for example `production`) is registered
|
||||
|
||||
3. **Database Connection Issues**
|
||||
- Verify PostgreSQL service is running
|
||||
@@ -1398,7 +1294,7 @@ export LOG_LEVEL=DEBUG
|
||||
- **Horizontal Scaling**: Deploy multiple worker instances
|
||||
- **Task Queue Distribution**: One worker pod per `RUNTIME`; queues are
|
||||
`predictions_batch-{RUNTIME}-queue`, `minimal_retrain-{RUNTIME}-queue`,
|
||||
`drift-{RUNTIME}-queue`, `simple_metrics-{RUNTIME}-queue`
|
||||
`drift-core-queue`, `simple_metrics-core-queue`
|
||||
- **Database Performance**: Optimize indexes and connection pooling
|
||||
- **MLFlow Performance**: Configure appropriate model serving resources
|
||||
|
||||
|
||||
149
docs/E2E_TEST_REPORT.md
Normal file
149
docs/E2E_TEST_REPORT.md
Normal file
@@ -0,0 +1,149 @@
|
||||
# E2E test run report
|
||||
|
||||
**Date:** 2026-05-08
|
||||
**Command:** `source venv/bin/activate && rtk pytest e2e/ -v --tb=short`
|
||||
**Environment:** Linux, Python 3.11.15, pytest 9.0.3
|
||||
|
||||
## Summary
|
||||
|
||||
| Metric | Count |
|
||||
|--------|------:|
|
||||
| Collected | 43 |
|
||||
| **Passed** | **37** |
|
||||
| **Failed** | **6** |
|
||||
|
||||
Full pytest output (compressed by `rtk`) was written to:
|
||||
|
||||
`~/.local/share/rtk/tee/1778268193_pytest.log`
|
||||
|
||||
---
|
||||
|
||||
## Failed tests (6)
|
||||
|
||||
1. `e2e/test_drift.py::test_drift_happy_path_persists_all_columns_with_reference_data`
|
||||
2. `e2e/test_drift.py::test_drift_uses_30pct_fallback_when_reference_unavailable`
|
||||
3. `e2e/test_drift.py::test_drift_chunk_period_seconds_preserves_seconds_in_chunk_start_date`
|
||||
4. `e2e/test_predictions_batch_prediction_process.py::test_scenario_2_1_3_input_gate_triggers_repeat`
|
||||
5. `e2e/test_predictions_batch_prediction_process.py::test_scenario_2_2_3_transform_gate_triggers_repeat`
|
||||
6. `e2e/test_predictions_batch_prediction_process.py::test_scenario_2_3_3_predict_gate_triggers_repeat`
|
||||
|
||||
**Follow-up:** Jensen–Shannon NULLs for the drift happy-path scenario are fully traced (histogram out-of-range + `density=True`, vs NannyML’s leftover bin) in [DRIFT_JS_NULL_INVESTIGATION.md](./DRIFT_JS_NULL_INVESTIGATION.md).
|
||||
|
||||
---
|
||||
|
||||
## Failure group A — Drift: `drift_metrics.value` NOT NULL (2 tests)
|
||||
|
||||
### Error
|
||||
|
||||
`psycopg2.errors.NotNullViolation`: null value in column `value` of relation `sientia_data.drift_metrics` violates not-null constraint.
|
||||
|
||||
Example failing row (from logs): `feature=sensor_2`, `method=jensen_shannon`, `value=null`, with `kolmogorov_smirnov` / `wasserstein` populated for the same chunk.
|
||||
|
||||
The bulk INSERT built by `Activities.export_data_to_postgres` includes parameters such as `'value__4': None` for `jensen_shannon` on a given chunk.
|
||||
|
||||
### Root cause
|
||||
|
||||
`calculate_drift` (real `DriftAnalysis` + `ModelMetrics.get_drift_metrics`) can emit **NaN / missing** values for some metric methods (here **Jensen–Shannon**) on some features/chunks. Pandas/SQLAlchemy turns that into SQL `NULL`, while the E2E schema (mirroring production) defines:
|
||||
|
||||
```sql
|
||||
value numeric NOT NULL
|
||||
```
|
||||
|
||||
in `e2e/db_schema.sql` for `sientia_data.drift_metrics`.
|
||||
|
||||
### Recommended fixes (pick one consistent with product rules)
|
||||
|
||||
1. **Application layer (preferred if NULLs are never valid in production):** Before export, sanitize the drift dataframe — e.g. drop rows where `value` is null/NaN, or replace with a defined sentinel (only if product agrees), or skip emitting that method row when the statistic is undefined.
|
||||
2. **Analytics layer:** Harden the Jensen–Shannon (and similar) paths so they always return a finite float for the supported inputs, or explicitly map “undefined” to an agreed numeric convention.
|
||||
3. **Schema (only if product allows missing metrics):** Align DDL with reality by making `value` nullable — **only** if production and downstream consumers already expect missing metrics; the E2E comment in `test_drift.py` suggests `feature`/`timestamp` nullable cases exist, but `value` is still listed in `NON_NULL_DRIFT_COLUMNS`.
|
||||
|
||||
### Tests affected
|
||||
|
||||
- `test_drift_happy_path_persists_all_columns_with_reference_data`
|
||||
- `test_drift_uses_30pct_fallback_when_reference_unavailable`
|
||||
|
||||
---
|
||||
|
||||
## Failure group B — Drift: chunk period seconds (1 test)
|
||||
|
||||
### Error
|
||||
|
||||
`AssertionError: expected at least one drift row to be persisted` (`e2e/test_drift.py:493`).
|
||||
|
||||
The workflow run completed without failing the test via `WorkflowFailureError`, but **no rows** were found in `sientia_data.drift_metrics` for the model.
|
||||
|
||||
### Likely cause
|
||||
|
||||
In `Drift.run`, persistence runs only when `if drift_data:` is truthy (`laborious/workflows/drift.py`). An **empty** drift result skips `export_data_to_postgres`, so the table stays empty.
|
||||
|
||||
Probable reasons:
|
||||
|
||||
- With **`chunk_period='s'`** and only **three** target rows (30 s spacing), `calculate_drift` / `DriftAnalysis` may produce **no output rows** (insufficient data per chunk or internal filters).
|
||||
- Less likely here: time-window mismatch — timestamps are built from `datetime.now(UTC)` with `interval: 60` minutes from `drift_base.json`, so data should still fall in the window.
|
||||
|
||||
### Recommended fixes
|
||||
|
||||
1. **Test data:** Increase the number of second-spaced points (and/or span multiple chunk boundaries) so the analyzer reliably emits at least one chunk row.
|
||||
2. **Product code:** If sub-minute chunking is required to always produce metrics when any data exists, adjust `ModelMetrics` / `DriftAnalysis` integration for small-N second buckets.
|
||||
3. **Diagnostics:** Run the same scenario with `pytest -s` and confirm logs for “empty `drift_data`” vs export errors.
|
||||
|
||||
Solução a ser aplicada:
|
||||
Dividir em dois testes:
|
||||
|
||||
1. Teste com dados suficientes para produzir pelo menos uma linha de drift
|
||||
2. Teste com dados insuficientes para produzir pelo menos uma linha de drift, mas ja esperando os erros e validando que nao foi persistido nada
|
||||
|
||||
### Test affected
|
||||
|
||||
- `test_drift_chunk_period_seconds_preserves_seconds_in_chunk_start_date`
|
||||
|
||||
---
|
||||
|
||||
## Failure group C — Prediction REPEAT path: unique constraint on predictions (3 tests)
|
||||
|
||||
### Error
|
||||
|
||||
`psycopg2.errors.UniqueViolation`: duplicate key value violates unique constraint **`unique_model_id_timestamp`** on `sientia_data.predictions` (`model_id`, `timestamp`).
|
||||
|
||||
### Root cause
|
||||
|
||||
REPEAT is handled by `Activities.repeat_last_prediction` (registered from `sientia_do.temporal.activities.postgres_sync.Postgres`, via `Storage` inheritance in `laborious/activities/storage.py`). The E2E tests seed a prior row with a **fixed** timestamp:
|
||||
|
||||
```python
|
||||
# e2e/test_predictions_batch_prediction_process.py — insert_sample_prediction
|
||||
VALUES ({model_id}, '2024-01-01 12:00:00+00:00', ...)
|
||||
```
|
||||
|
||||
The helper `assert_repeat` expects **two** rows with the same `(model_id, prediction, prediction_confidence, prediction_status)` but compares **only** those columns — not `timestamp` (`e2e/helpers.py`). So the intended behavior is: **duplicate business payload**, not necessarily **duplicate primary unique key** `(model_id, timestamp)`.
|
||||
|
||||
If `repeat_last_prediction` **INSERT**s a copy using the **same** `timestamp` as the last prediction, Postgres correctly rejects the second insert.
|
||||
|
||||
### Recommended fixes
|
||||
|
||||
1. **Implement REPEAT as UPSERT:** Use `ON CONFLICT (model_id, timestamp) DO UPDATE` (or the project’s existing `export_data_to_postgres` `on_conflict` pattern used in `FormatAndExportPrediction`) when writing the repeated prediction — if the product definition of REPEAT is “refresh same logical slot.”
|
||||
2. **Insert with the current batch timestamp:** Copy numeric/status fields from the last row but set `timestamp` to the **new** batch instant (e.g. the workflow’s `last_timestamp` / slice timestamp). This matches `assert_repeat`, which does not assert on `timestamp`.
|
||||
3. **Test-only change (weakest):** Relax assertions or change seed data — only if production behavior is “duplicate key is expected” (unlikely).
|
||||
|
||||
Solução a ser aplicada:
|
||||
Alterar o teste para usar o timestamp da ultima execucao do batch, ao inves do timestamp do primeiro batch.
|
||||
2. **Insert with the current batch timestamp:** Copy numeric/status fields from the last row but set `timestamp` to the **new** batch instant (e.g. the workflow’s `last_timestamp` / slice timestamp). This matches `assert_repeat`, which does not assert on `timestamp`.
|
||||
|
||||
### Tests affected
|
||||
|
||||
- `test_scenario_2_1_3_input_gate_triggers_repeat`
|
||||
- `test_scenario_2_2_3_transform_gate_triggers_repeat`
|
||||
- `test_scenario_2_3_3_predict_gate_triggers_repeat`
|
||||
|
||||
---
|
||||
|
||||
## Passing areas (sanity check)
|
||||
|
||||
All scenarios in `test_child_workflows_e2e.py`, `test_minimal_retrain.py`, `test_minio_offload.py`, `test_predictions_batch_format_export.py`, `test_predictions_batch_main_workflow.py` (except the three REPEAT cases above), and `test_simple_metrics.py` **passed** in this run.
|
||||
|
||||
---
|
||||
|
||||
## Suggested order of work
|
||||
|
||||
1. Fix **drift `value` NULL** — unblocks two high-value drift E2Es and may clarify the chunk-seconds scenario if exports start succeeding consistently.
|
||||
2. Fix **`repeat_last_prediction` uniqueness** — unblocks three prediction-process E2Es; implementation likely lives in **`sientia_do`** Postgres activities, not in this repo.
|
||||
3. Revisit **`test_drift_chunk_period_seconds`** data volume / expectations after drift export is stable.
|
||||
@@ -1,40 +0,0 @@
|
||||
# Two independent MinIO expiry mechanisms
|
||||
|
||||
This platform expires objects in MinIO through two mechanisms that were built independently, for
|
||||
different features, and remain independent today. This note exists so a reader who finds one does
|
||||
not assume it is the whole story.
|
||||
|
||||
## `ModelImport.ensure_import_bundle_retention`
|
||||
|
||||
`laborious/activities/model_import.py`. Applies a **bucket lifecycle rule** (a MinIO/S3-native
|
||||
`Expiration` rule, `ENABLED`, keyed by a fixed `rule_id`) to the import bundle prefix
|
||||
(`import_config['prefix']`), expiring objects after `import_config['retention_days']` days. The
|
||||
rule is declarative: once set, MinIO itself removes the object when it ages out — the workflow does
|
||||
not scan for or delete anything. The activity's job is only to ensure the rule exists and matches
|
||||
the configured `days`/`prefix`, which is what makes "lifecycle active" a criterion the code can
|
||||
verify rather than an assumption about a Helm value.
|
||||
|
||||
## `Storage.cleanup_minio_objects_expired`
|
||||
|
||||
`laborious/activities/storage.py`. An older, unrelated activity that **scans and deletes**: given a
|
||||
list of key prefixes and a retention window (`retention_hours`), it lists objects under each prefix,
|
||||
compares their age against `now() - retention_hours`, and deletes the ones that have aged out. There
|
||||
is no bucket-side rule — the deletion happens in this code, on whatever schedule calls the activity.
|
||||
|
||||
## Why they are not unified today
|
||||
|
||||
They expire different prefix families, on different units (days vs. hours), through different
|
||||
mechanisms (declarative rule vs. active scan), and were written for different features at different
|
||||
times. Neither mechanism depends on or interferes with the other; a bucket can carry an
|
||||
`ensure_import_bundle_retention` lifecycle rule for one prefix and still have
|
||||
`cleanup_minio_objects_expired` scanning a different prefix on its own cadence.
|
||||
|
||||
## What unifying them would require
|
||||
|
||||
Either retiring the scan-and-delete path in favor of a lifecycle rule for every prefix it currently
|
||||
handles (only possible if every one of its callers can tolerate a fixed number of retention *days*,
|
||||
since S3-style lifecycle rules do not support hour-granularity windows), or teaching the
|
||||
lifecycle-rule path to manage a per-prefix table of rules with arbitrary retention windows instead of
|
||||
the one fixed rule it manages today. Either direction touches callers and configuration this change
|
||||
does not, and neither is a refactor — this is tracked as a separate, larger change, not a task of
|
||||
this one.
|
||||
@@ -1,166 +0,0 @@
|
||||
# Model flavors
|
||||
|
||||
The *flavor* is how this runtime is told to deserialize a model. It is **configuration, never
|
||||
discovery**: the value comes from the model document in the `models` collection and the code never
|
||||
reads the artifact's `MLmodel` manifest to guess it.
|
||||
|
||||
```json
|
||||
{
|
||||
"model_config": {
|
||||
"transform_flavor": "sklearn",
|
||||
"predict_flavor": "joblib",
|
||||
"retention_minutes": 60,
|
||||
"target": "SE"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
`predict_flavor` drives `load_predict_model` (artifact `prediction_model`); `transform_flavor` drives
|
||||
`load_transform_model` (artifact `data_model`). Both default to `sklearn` when the key is absent.
|
||||
|
||||
## Other `model_config` flags
|
||||
|
||||
Ad-hoc behavioral overrides, read directly with `.get(key, default)` — not part of the model
|
||||
document written by import (`model-import.md` § 12), and not validated against a schema:
|
||||
|
||||
| Key | Default | Effect |
|
||||
|---|---|---|
|
||||
| `skip_transform` | `False` | `fit_models` (retrain) skips calling `transform` before fitting. |
|
||||
| `skip_predict_timestamp` | `False` | `request_predict` skips reinserting the `timestamp` column (from the index) into the transformed data before calling predict. |
|
||||
|
||||
## Accepted values
|
||||
|
||||
| Flavor | Read | Write (retrain) | How it loads |
|
||||
|---|---|---|---|
|
||||
| `sklearn` | yes | yes | `mlflow.sklearn.load_model(<uri>)` — the loader fetches what it needs |
|
||||
| `pyfunc` | yes | yes | `mlflow.pyfunc.load_model(<uri>)`; the wrapper path downloads the artifact instead and unwraps `._model_impl.python_model` |
|
||||
| `joblib` | yes | yes | downloads the same artifact and `joblib.load`s its `<artifact>/model.pkl` — **no MLflow flavor loader involved** |
|
||||
|
||||
Anything else raises `ValueError` with `INVALID_FLAVOR_MESSAGE`, **before** any artifact is resolved
|
||||
or downloaded. The same message and the same three values govern reading and writing.
|
||||
|
||||
`pytorch` was accepted until this change and is now rejected like any other unknown value: `torch` is
|
||||
not installed in the runtime, so `mlflow.pytorch.load_model` could never have loaded anything there.
|
||||
|
||||
**There is no fallback between flavors.** A failing loader propagates: one attempt, one cause. A
|
||||
model configured with the wrong flavor fails visibly instead of being rescued silently.
|
||||
|
||||
## Same artifact, three readers
|
||||
|
||||
The flavor changes the reader, never the path. All three resolve the same two artifact directories of
|
||||
the Production run:
|
||||
|
||||
```
|
||||
<run>/artifacts/
|
||||
├── data_model/ ← transform_flavor loads this one
|
||||
│ ├── MLmodel ← parsed by sklearn/pyfunc; ignored by joblib
|
||||
│ ├── model.pkl ← what joblib.load reads
|
||||
│ └── code/ ← put on sys.path by every flavor, when present
|
||||
├── prediction_model/ ← predict_flavor loads this one
|
||||
└── model_card/, *.csv
|
||||
```
|
||||
|
||||
The run's `artifacts/` root holds **no** `.pkl` — every model pickle lives one level down, inside its
|
||||
artifact directory. Predict resolves that directory through the registry URI
|
||||
(`models:/<name>/production`, which MLflow points at the same `prediction_model`); transform builds
|
||||
the run artifact URI directly with `get_model_uri`, since the run is already resolved. Observability
|
||||
is the same for the three: `MODEL_READ_*` on load, `MODEL_WRITE_*` on `log_model`, under the base
|
||||
`operation_type` — no per-flavor metric, operation type or label.
|
||||
|
||||
## `code/` on `sys.path`, for every flavor
|
||||
|
||||
An artifact may ship the modules its own classes live in, under `<artifact>/code/`. Without that
|
||||
directory on `sys.path`, unpickling raises `ModuleNotFoundError` — the artifact is intact, the class
|
||||
is simply not importable.
|
||||
|
||||
MLflow's loaders do add it, but **from the manifest**: `_add_code_from_conf_to_system_path` acts only
|
||||
when the `MLmodel` flavor config declares a `code` entry. An artifact carrying `code/` on disk
|
||||
without that entry — code logged with `log_artifacts`, a re-log that dropped the config, anything a
|
||||
`joblib` retrain wrote — loads with no module path at all. So this runtime reads the **directory**,
|
||||
which no manifest can misdescribe:
|
||||
|
||||
- `prepend_artifact_code_dir` inserts `<artifact>/code` at the front of `sys.path`, **only when that
|
||||
directory exists**, and only once (an entry already present is not duplicated).
|
||||
- `download_artifact_code_dir` is what feeds it on the download paths, and it is deliberately cheap:
|
||||
the remote listing (`list_artifacts(run_id, '<artifact>')`) decides first, so a run without `code`
|
||||
downloads **nothing**, and a run with it downloads only `<artifact>/code` — a handful of `.py`
|
||||
files, never the model's weights.
|
||||
- `download_model` runs that step first, for every flavor and both of its branches, *before* any
|
||||
model is loaded — the insert has to be in place before the loader unpickles.
|
||||
- The models themselves are not pre-downloaded: `mlflow.sklearn.load_model` and
|
||||
`mlflow.pyfunc.load_model` are handed the model URI and fetch what they need. `joblib` downloads
|
||||
its artifact because it must read the pickle, and adds the entry again inside `_load_model_pickle`
|
||||
through the same helper, so the flavor works when called directly too.
|
||||
- When the manifest *does* declare `code`, MLflow inserts the same directory a second time. Harmless.
|
||||
|
||||
The entry is never removed: a loaded model may import lazily, long after the load returned.
|
||||
|
||||
The directory `download_artifact_code_dir` downloaded into is also what `download_model` returns
|
||||
next to the model — `tuple[Any, str | None]`, the second element being the artifact directory the
|
||||
retrain path re-logs `code/` from, or `None` when the artifact has no `code/`. It falls back to
|
||||
whatever the loader downloaded on its own, which is a real artifact directory for `joblib` and for
|
||||
the wrapper branch. `load_predict_model` and `load_transform_model` stay the **single** place where
|
||||
the flavor is decided: `download_model` has no `joblib` branch, it delegates by `model_type`. Its
|
||||
only remaining branch is `load_wrapper`, which is derived from `flavor == 'pyfunc'` but is otherwise
|
||||
orthogonal to the flavor.
|
||||
|
||||
> `sys.path` is process-global: two models embedding a package with the same name (`utils/`) on the
|
||||
> same worker resolve to whichever entered first. Mitigation, if it bites: one worker per model.
|
||||
|
||||
## `sklearn`
|
||||
|
||||
Default. A standard MLflow sklearn artifact: `mlflow.sklearn.log_model` writes `MLmodel` +
|
||||
`model.pkl`, `mlflow.sklearn.load_model` reads it, and every outside consumer (MLflow UI, MLflow
|
||||
serving) recognizes it as a model.
|
||||
|
||||
## `pyfunc`
|
||||
|
||||
For models logged as `mlflow.pyfunc`, including the wrapper models this pipeline produces. Retrain
|
||||
writes through the wrapper's own `store_model(artifact_path=<artifact>, code_path=…)`. Two read
|
||||
paths, and the second is **orthogonal to the flavor**:
|
||||
|
||||
- Plain: `mlflow.pyfunc.load_model(<uri>)`, called as `predict(context, model_input)`.
|
||||
- Wrapper (`download_model(load_wrapper=True)`): downloads to `./tmp/artifacts/<model>/<artifact>`
|
||||
via `dowload_artifacts`, loads as pyfunc, then unwraps `._model_impl.python_model`. Inside such an
|
||||
artifact the real pickle is **nested** — `<artifact>/artifacts/stacking_model.pkl` for predict,
|
||||
`<artifact>/artifacts/training_transformer.pkl` for transform (`PREDICTION_COMPRESSED_PATH` /
|
||||
`TRANSFORMED_COMPRESSED_PATH`). That nesting is why the `joblib` listing does not recurse.
|
||||
|
||||
## `joblib`
|
||||
|
||||
For models whose `.pkl` is in the compressed form `joblib.dump` writes, which MLflow's `pickle.load`
|
||||
does not read — the artifact is intact, but no flavor loader opens it.
|
||||
|
||||
- **Read**: `client.download_artifacts(run_id, <artifact>)`, then
|
||||
`joblib.load(<artifact>/model.pkl)`. `model.pkl` is `MODEL_PICKLE_NAME`, the same constant
|
||||
`log_model` writes; another `.pkl` directly inside the artifact is a fallback, in name order.
|
||||
- **Write**: `joblib.dump(model, '<artifact>/model.pkl', compress=3)`, plus a re-log of the source
|
||||
artifact's `code/` to `<artifact>/code`.
|
||||
- **`code/` goes on `sys.path`** before the load, when present — see the section above; the joblib
|
||||
flavor is where that started, and it is now shared by all three.
|
||||
- **No `MLmodel` is written**, so only this runtime reads the artifact back —
|
||||
`mlflow.sklearn.load_model` and `models:/<name>/production` both fail on it. A model retrained in
|
||||
joblib stays joblib.
|
||||
- **The pickled object must expose `predict(data)`**, plus `fit(data)` if retrained. A pyfunc
|
||||
`PythonModel` wrapper expecting `predict(context, model_input)` raises `TypeError`.
|
||||
|
||||
## Switching a model to `joblib`
|
||||
|
||||
The runtime does not migrate documents; editing `models` is an operator step.
|
||||
|
||||
1. **Inventory before the deploy** — run predict/transform and note which models raise. There is no
|
||||
discovery afterwards: the failure names the loader, not the fix.
|
||||
2. **Edit the document** — `predict_flavor` and/or `transform_flavor` to `'joblib'`. Imported models
|
||||
arrive ready: `import_model` writes it itself, see [`model-import.md`](model-import.md) § 12.
|
||||
3. **Check the first retrain** — the new run must carry `prediction_model/model.pkl` and, when the
|
||||
source had one, `prediction_model/code/`.
|
||||
4. **Rollback is not just a revert** — artifacts a joblib retrain already wrote stop being loadable,
|
||||
so it also means promoting the previous version back and reverting the documents' flavors.
|
||||
|
||||
## Known gap: how the run id is resolved
|
||||
|
||||
`get_model_run_id` takes `source.split('/')[2]` of the registered version's `source`. **Not confirmed
|
||||
against the real tracking server**: no scheme reproduced locally puts the run id there (a file store
|
||||
yields `''`, S3 the bucket, `mlflow-artifacts:/` the experiment id), yet predict works in production.
|
||||
Every flavor that resolves a run goes through this line — confirm it with a **read-only** query
|
||||
before relying on it.
|
||||
@@ -1,417 +0,0 @@
|
||||
# Importing a model (`import_model`)
|
||||
|
||||
The `import_model` workflow turns an encrypted `.sientia` object in MinIO into a registered MLflow
|
||||
model plus one document in the MongoDB `models` collection — or into a row that says, in a sentence a
|
||||
person can read, why it did not.
|
||||
|
||||
This document is the contract for the two sides that talk to it: the **frontend**, which uploads the
|
||||
object and builds the request, and the **BFF**, which owns the import log row. Everything here is
|
||||
enforced by code in this repository; where a statement is about a column or a constraint, the
|
||||
authority is the migration named in §7.
|
||||
|
||||
Format of the file itself: [`sientia-bundle-format.md`](sientia-bundle-format.md).
|
||||
|
||||
---
|
||||
|
||||
## 1. What the caller does, in order
|
||||
|
||||
1. **Upload** the `.sientia` object to the import bucket, under the configured prefix (default
|
||||
`imported_models/`). The key must end in `.sientia`.
|
||||
2. **Compute `expected_digest`** — the SHA-256, lowercase hex, of the **exact bytes that were
|
||||
uploaded**. Not of a re-read, not of a re-serialisation: of what went over the wire.
|
||||
3. **Encrypt the bundle password** into a `password_envelope` (§3).
|
||||
4. **Insert the import log row** (§5) and take its `id`.
|
||||
5. **Start the workflow** on `import_model-{RUNTIME}-queue` with the input of §2.
|
||||
|
||||
The workflow claims the row, opens the bundle, provisions MLflow, writes the model document, cleans
|
||||
up and records one terminal verdict. It never inserts a row, never promotes a version and never
|
||||
activates a model.
|
||||
|
||||
## 2. The workflow input
|
||||
|
||||
```json
|
||||
{
|
||||
"import_run_id": 41,
|
||||
"object_key": "imported_models/sales_forecast_v3.sientia",
|
||||
"expected_digest": "9f86d081884c7d659a2feaa0c55ad015a3bf4f1b2b0b822cd15d6c15b0f00a08",
|
||||
"password_envelope": "Base64OfNonceCiphertextTag==",
|
||||
"bucket": "sientia"
|
||||
}
|
||||
```
|
||||
|
||||
| Field | Required | Meaning |
|
||||
|---|---|---|
|
||||
| `import_run_id` | yes | primary key of the `public.experiment_run` row the BFF inserted. Must be a positive integer |
|
||||
| `object_key` | yes | the uploaded object's key. Must match the row's `file_name` when the row carries one |
|
||||
| `expected_digest` | **yes** | SHA-256 of the uploaded bytes, 64 lowercase hex characters. There is no path that skips the comparison: a missing value fails at `received`, a mismatching one at `download`, before the header is parsed |
|
||||
| `password_envelope` | yes | the bundle password, encrypted (§3). A plaintext `password` field is **refused** |
|
||||
| `bucket` | no | defaults to `IMPORT_BUNDLE_BUCKET` |
|
||||
|
||||
**Fields that are refused rather than ignored**: `project_id`, `project_name`,
|
||||
`provisioning_target`, `schema`, `table_name`, `status_table`, `password`. An input carrying any of
|
||||
them fails at `received`, so no caller can believe it recorded something the import does not record
|
||||
(see §10).
|
||||
|
||||
## 3. The password envelope
|
||||
|
||||
AES-256-GCM, and the ciphertext layout is what the WebCrypto API returns natively:
|
||||
|
||||
```
|
||||
envelope = base64( nonce(12 bytes) || ciphertext || tag(16 bytes) )
|
||||
```
|
||||
|
||||
`crypto.subtle.encrypt({name: 'AES-GCM', iv: nonce}, key, plaintext)` returns `ciphertext || tag`
|
||||
already concatenated, so the browser side is: generate 12 random bytes, encrypt, concatenate the
|
||||
nonce in front, base64 the result.
|
||||
|
||||
The key is a 32-byte AES-256 key shared with the worker through `IMPORT_PASSWORD_KEY` (base64). The
|
||||
worker decrypts the envelope **inside the activity that uses the password**, so no plaintext password
|
||||
is ever an activity input, an activity result or a workflow variable — and therefore never appears in
|
||||
the Temporal event history. A test asserts that for every failure path.
|
||||
|
||||
A failure to open the envelope is recorded as a `decryption` failure, distinct in the logs from a
|
||||
wrong *bundle* password but sharing its step, because the step says where the import stopped, not who
|
||||
is at fault.
|
||||
|
||||
## 4. Environment
|
||||
|
||||
| Variable | Default | What it is |
|
||||
|---|---|---|
|
||||
| `IMPORT_MODEL_ENABLED` | unset | registers the fifth (import) worker. Truthy values: `1`, `true`, `yes`, `on` |
|
||||
| `IMPORT_PASSWORD_KEY` | — | base64 of the 32-byte AES-256-GCM key that wraps bundle passwords |
|
||||
| `IMPORT_BUNDLE_BUCKET` | `sientia` | bucket holding the uploaded bundles |
|
||||
| `IMPORT_BUNDLE_PREFIX` | `imported_models/` | key prefix the uploads must sit under |
|
||||
| `IMPORT_BUNDLE_RETENTION_DAYS` | `7` | lifecycle expiry the importer applies to that prefix |
|
||||
| `IMPORT_WORK_DIR` | `<tempdir>/sientia-import` | worker scratch directory |
|
||||
| `IMPORT_MAX_OBJECT_BYTES` | 1 GiB | gate 1 |
|
||||
| `IMPORT_MAX_ARCHIVE_ENTRIES` | `5000` | gate 4 |
|
||||
| `IMPORT_MAX_UNCOMPRESSED_BYTES` | 4 GiB | gate 4 |
|
||||
| `IMPORT_MAX_COMPRESSION_RATIO` | `200` | gate 4 |
|
||||
| `IMPORT_MODELS_COLLECTION` | `models` | the MongoDB collection that *is* the model listing |
|
||||
| `IMPORT_STATUS_DB_NAME` | `sientia-core-mlops-bff` | the BFF database holding the import log |
|
||||
| `IMPORT_STATUS_DB_HOST` | `POSTGRES_HOST` | override: the log is on another server |
|
||||
| `IMPORT_STATUS_DB_PORT` | `POSTGRES_PORT` | override |
|
||||
| `IMPORT_STATUS_DB_USER` | `POSTGRES_USER` | override |
|
||||
| `IMPORT_STATUS_DB_PASSWORD` | `POSTGRES_PASSWORD` | override |
|
||||
|
||||
The import log is a **second connection but the same server**: host, port, user and password come
|
||||
from `POSTGRES_*`, and only the database name is import configuration, because the log lives in the
|
||||
BFF's database rather than in the one Laborious holds for predictions. The import writes to that
|
||||
database and no other, and to one table in it. The four `IMPORT_STATUS_DB_*` connection variables
|
||||
are overrides for a deployment whose log genuinely sits elsewhere — the e2e suite, which runs the
|
||||
two databases as two containers, is exactly that case.
|
||||
|
||||
A missing or malformed `IMPORT_PASSWORD_KEY` is an import failure at the decryption step, not a boot
|
||||
failure: the worker starts either way.
|
||||
|
||||
## 5. The import log row — the BFF's half
|
||||
|
||||
The row is `public.experiment_run` in `sientia-core-mlops-bff`. **The workflow never inserts it.**
|
||||
The inserting side writes:
|
||||
|
||||
| Column | Value |
|
||||
|---|---|
|
||||
| `run_type` | `'IMPORT'` |
|
||||
| `status` | `'PENDING'` |
|
||||
| `username` | whoever asked for the import |
|
||||
| `bucket_name` | the bucket the object was uploaded to |
|
||||
| `file_name` | the object key |
|
||||
| `request_data` | the request payload. If it carries a digest under `expected_digest`, `digest`, `sha256` or `file_digest`, the workflow checks it against the input's |
|
||||
|
||||
Then it passes the returned `id` as `import_run_id`.
|
||||
|
||||
The workflow's first act is to **claim** that row:
|
||||
|
||||
```sql
|
||||
UPDATE public.experiment_run
|
||||
SET status = 'ORCHESTRATOR_WAITING_PROC', updated_at = :now,
|
||||
orchestrator_response_data = :detail
|
||||
WHERE id = :import_run_id
|
||||
AND run_type = 'IMPORT'
|
||||
AND (status = 'PENDING'
|
||||
OR (status = 'ORCHESTRATOR_WAITING_PROC'
|
||||
AND orchestrator_response_data -> 'import' ->> 'workflow_id' = :workflow_id))
|
||||
```
|
||||
|
||||
The second disjunct makes a re-execution of the same workflow idempotent instead of a conflict. Every
|
||||
later write keys on `id` alone, because ownership was settled here.
|
||||
|
||||
**Columns the workflow never writes**: `username`, `request_data`, `bucket_name`, `file_name`,
|
||||
`run_type`, `created_at`. **Columns it does write**: `status`, `error_message`, `updated_at`,
|
||||
`orchestrator_response_data`, and — as soon as the bundle is open, before anything exists in MLflow —
|
||||
`experiment_name` (the bundle's own, verbatim) and `run_name` (`import-{model_name}-v{model_version}`,
|
||||
truncated inside the name segment to fit 50 characters). Those two are **overwritten** by the
|
||||
workflow; whatever the inserting side put there is replaced.
|
||||
|
||||
### When the row cannot be claimed
|
||||
|
||||
If the claim matches nothing — no such id, a `run_type='TRAIN'` row, or a row already running under
|
||||
another workflow — the import **writes nothing anywhere**, because the row belongs to someone else.
|
||||
Instead it:
|
||||
|
||||
- fails the Temporal run with error type `ImportLogRowNotClaimableError` and code
|
||||
`IMPORT_LOG_ROW_NOT_CLAIMABLE`, non-retryable;
|
||||
- increments the import error metric and fires a notification;
|
||||
- logs the id and the observed state (which of the three causes it was).
|
||||
|
||||
So a row left `PENDING` with no `workflow_id` in its detail JSON means the workflow never owned it.
|
||||
**Closing that row out belongs to whichever side inserted it** — by timeout, by a sweeper, by a person
|
||||
— however that side chooses. The importer deliberately does not guess.
|
||||
|
||||
## 6. Status vocabulary
|
||||
|
||||
| Status | Written when | New in `V12`? |
|
||||
|---|---|---|
|
||||
| `ORCHESTRATOR_WAITING_PROC` | claimed, and while running | no — **reused** from the training flow |
|
||||
| `ORCHESTRATOR_VALIDATION_ERROR` | failure at `received` … `content_policy` — the user can act on it | no — **reused** |
|
||||
| `IMPORT_SUCCESS` | every step succeeded | yes |
|
||||
| `IMPORT_ERROR` | failure at `experiment_creation` … `model_document`, or unclassified — only support can act on it | yes |
|
||||
|
||||
The two reused values describe a *stage*, not a flow: what says the row is an import is `run_type`,
|
||||
not the status. There is deliberately no `IMPORT_RUNNING`; it would carry no information `run_type`
|
||||
does not already carry, and it would widen both the pairing `CHECK` and the frontend's switch for
|
||||
nothing.
|
||||
|
||||
`error_message` is the user-facing sentence, taken from the catalog in §8. On success it is written as
|
||||
`NULL`, never `''` — the column carries a `LENGTH(error_message) >= 3` check.
|
||||
|
||||
## 7. The migration this depends on
|
||||
|
||||
`sientia-dataops-database-migrations`,
|
||||
`postgres/sientia-core-mlops-bff/migrations/V12__add_import_statuses_to_experiment_run.sql`:
|
||||
|
||||
- adds `IMPORT_SUCCESS` and `IMPORT_ERROR` to `chk_experiment_run_status`;
|
||||
- adds `chk_experiment_run_terminal_status_run_type`, pairing only the **terminal** values with their
|
||||
`run_type` (`IMPORT_*` on `'IMPORT'`, `TRAINING_*` on `'TRAIN'`) and leaving the shared
|
||||
`ORCHESTRATOR_*` values out of the rule;
|
||||
- **adds no columns at all.**
|
||||
|
||||
`V12` must be merged and applied before `IMPORT_MODEL_ENABLED` is set anywhere. Without it, an import
|
||||
provisions the model and then cannot record its verdict: the claim still succeeds (no new columns are
|
||||
needed for it), the old `CHECK` refuses `IMPORT_SUCCESS`, and the failure surfaces as
|
||||
`IMPORT_STATUS_NOT_RECORDED` on the first attempt — no retry wait — leaving the row in
|
||||
`ORCHESTRATOR_WAITING_PROC` with a fully provisioned model behind it.
|
||||
|
||||
## 8. `orchestrator_response_data`, and the error catalog
|
||||
|
||||
There is no column for the error code and none for the workflow identifier (decision D21). Both live
|
||||
at fixed keys in the JSON column the workflow already owns, and the whole `import` object is rewritten
|
||||
on every write, so the identifier the claim wrote is never dropped:
|
||||
|
||||
```json
|
||||
{
|
||||
"import": {
|
||||
"workflow_id": "import_model-41-01HW…",
|
||||
"code": "IMPORT_BUNDLE_INCOMPLETE",
|
||||
"step": "structure_validation",
|
||||
"gate": 6,
|
||||
"source": { "model_name": "sales_forecast", "experiment_name": "…", "…": "…" },
|
||||
"cleanup_failed": false
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
- `workflow_id` — what a user quotes to support, and what the claim's idempotency disjunct reads.
|
||||
- `code` — the frontend's i18n key. **Absent on success.**
|
||||
- `step` — where the import stopped (the values of §9).
|
||||
- `gate` — which of the seven gates rejected it, when one did.
|
||||
- `source` — the bundle's own metadata block, once it is known.
|
||||
- `cleanup_failed` — temporary files were left on the worker. Present on a **successful** import too:
|
||||
it never changes the status, because the user's model was imported.
|
||||
|
||||
### Code → reason
|
||||
|
||||
`error_message` is exactly the sentence in this table; nothing is interpolated into it, and the
|
||||
exception's text — paths, table names, exception classes, tracebacks — reaches only the logs.
|
||||
|
||||
| Code | Step | Sentence written to `error_message` |
|
||||
|---|---|---|
|
||||
| `IMPORT_REQUEST_INVALID` | `received` | The import could not be started because the request was incomplete or contradictory. Nothing was created. Submit the import again. |
|
||||
| `IMPORT_FILE_UNREADABLE` | `download` | The uploaded file could not be read, or it is not the file that was sent. Upload it again. |
|
||||
| `IMPORT_BUNDLE_CANNOT_BE_OPENED` | `decryption` | This file could not be opened: the password is wrong, or the file is damaged or was altered. Confirm the password with whoever exported the model, then try again. |
|
||||
| `IMPORT_BUNDLE_REJECTED` | `archive_inspection` | The file's contents did not pass the platform's safety checks, so it was not opened. Export the model again from the origin platform. |
|
||||
| `IMPORT_BUNDLE_UNPACK_FAILED` | `extraction` | The file could not be unpacked safely and nothing from it was kept. Export the model again from the origin platform. |
|
||||
| `IMPORT_BUNDLE_INCOMPLETE` | `structure_validation` | This file is not a complete model export — part of what the platform needs is missing from it. Export the model again from the origin platform. |
|
||||
| `IMPORT_BUNDLE_UNEXPECTED_CONTENT` | `content_policy` | The file contains something a model export should not contain, so it was rejected. Export the model again from the origin platform. |
|
||||
| `IMPORT_STORAGE_PREPARATION_FAILED` | `experiment_creation` | The platform could not prepare a place to keep this model. The import stopped and no model was created. Contact support with this import's identifier. |
|
||||
| `IMPORT_MODEL_FILES_NOT_STORED` | `artifact_upload` | The model's files could not be stored on this platform. The import stopped and the model is not available. Contact support with this import's identifier. |
|
||||
| `IMPORT_MODEL_NOT_REGISTERED` | `registration` | The model's files were stored, but the model itself could not be registered, so it cannot be used. Contact support with this import's identifier. |
|
||||
| `IMPORT_MODEL_SETTINGS_NOT_SAVED` | `model_document` | The model was registered, but the settings that tell the platform how to run it could not be saved. Contact support with this import's identifier. |
|
||||
| `IMPORT_UNEXPECTED_ERROR` | — | The import stopped for an unexpected reason and no model was created. Contact support with this import's identifier. |
|
||||
|
||||
Three codes never reach a row's `error_message`, because there is no row to write or no failure to
|
||||
report to a user. They exist in the logs, the metric and the notification only:
|
||||
|
||||
| Code | When |
|
||||
|---|---|
|
||||
| `IMPORT_LOG_ROW_NOT_CLAIMABLE` | the row is not this workflow's (§5) |
|
||||
| `IMPORT_STATUS_NOT_RECORDED` | the verdict itself could not be written; the log line carries the verdict it could not record |
|
||||
| `IMPORT_CLEANUP_INCOMPLETE` | temporary files were left behind; also flagged as `cleanup_failed` in the detail |
|
||||
|
||||
## 9. Steps
|
||||
|
||||
`received`, `download`, `decryption`, `archive_inspection`, `extraction`, `structure_validation`,
|
||||
`content_policy`, `experiment_creation`, `artifact_upload`, `registration`, `model_document`,
|
||||
`cleanup`. The first seven are the gates (see the format document); a failure in any of them is a
|
||||
`ORCHESTRATOR_VALIDATION_ERROR`. The next four are provisioning; a failure there is an
|
||||
`IMPORT_ERROR`. `cleanup` is never a reported failure.
|
||||
|
||||
## 10. The import knows nothing about projects
|
||||
|
||||
The model listing **is** the MongoDB `models` document. No migrated relational database in this
|
||||
platform holds a project table or a model table, nothing in this repository reads a project-to-model
|
||||
link, and the workflow input carries no project target — which is why naming one is refused rather
|
||||
than accepted and ignored. Whatever the import screen does with projects is the frontend's business.
|
||||
|
||||
The bundle's own `metadata.model_project` is logged as an MLflow parameter (`source.model_project`),
|
||||
as a record of where the model came from. It is not a link.
|
||||
|
||||
## 11. What MLflow ends up holding
|
||||
|
||||
| Thing | Value |
|
||||
|---|---|
|
||||
| Experiment | `metadata.experiment_name` from the bundle, created if absent, reused if present |
|
||||
| Run | always **new**, created with the `run_name` already written on the record |
|
||||
| Artifacts | the bundle's whole `artifacts/` tree, with `prediction_model` and `data_model` at the run's artifact root |
|
||||
| Parameters | the bundle's `parameters`, verbatim, plus the origin metadata as `source.*` (including `source.run_id`, which is never dereferenced) |
|
||||
| Registered version | created on demand; source is the run's **resolved artifact URI** with `/prediction_model` appended, e.g. `mlflow-artifacts:/17/<run_id>/artifacts/prediction_model` |
|
||||
| Stage | **untouched** |
|
||||
|
||||
The source shape matters and is easy to get wrong: the legacy loader extracts the run id as
|
||||
`source.split('/')[2]`, which is the run id in the resolved artifact URI but the *artifact path* in a
|
||||
literal `runs:/<run_id>/prediction_model`. Registering with the unresolved URI breaks loading
|
||||
silently — nothing raises, the loader just returns the wrong string. This is the same value
|
||||
`mlflow.register_model` stores for the training flow, because it resolves the `runs:/` URI before
|
||||
writing it; `MlflowClient.create_model_version` does not, so the import resolves it itself.
|
||||
|
||||
## 12. The model document
|
||||
|
||||
Written **last**, into the collection named by `IMPORT_MODELS_COLLECTION` (default `models`):
|
||||
|
||||
```json
|
||||
{
|
||||
"id": "31",
|
||||
"name": "sales_forecast",
|
||||
"active": false,
|
||||
"model_config": {
|
||||
"transform_flavor": "sklearn",
|
||||
"predict_flavor": "joblib",
|
||||
"retention_minutes": 60,
|
||||
"target": "SE"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
- `id` is `str(max(int(id) for id in collection) + 1)` — computed in Python, from the ids that are
|
||||
numbers. Never a `count`, never a lexicographic sort (which would answer `"9" > "30"`). Ids that
|
||||
are not decimal numbers are ignored for the maximum and counted in the log.
|
||||
- `name` is `metadata.model_name` from the bundle, verbatim.
|
||||
- `active` is `false`, always. No code path in the import writes or updates a document with
|
||||
`active: true`.
|
||||
- `target` comes from `parameters.target_variable`; nothing else supplies it.
|
||||
- The `models` collection carries a **unique sparse index on `id`**, created by the MongoDB baseline
|
||||
migration. A duplicate-key insert is reported as a `model_document` failure — no retry, no id
|
||||
recomputation.
|
||||
|
||||
**A name already in the collection is left exactly as it is.** The import writes nothing, modifies
|
||||
nothing, logs the skip and succeeds: `active`, `target` and the flavors of a model already on the
|
||||
platform are operator-owned.
|
||||
|
||||
## 13. After a successful import — the manual steps
|
||||
|
||||
A `COMPLETED` import is **not** a servable model. Stated plainly:
|
||||
|
||||
- the registered version is **unstaged**, and `load_predict_model` only sees a version once it is in
|
||||
stage `Production`;
|
||||
- the model document is **inactive** (`active: false`).
|
||||
|
||||
So an operator must promote the version to `Production` and set `active: true`. This workflow never
|
||||
promotes and never activates — that runbook is QTZPOC-21's.
|
||||
|
||||
One more thing to know before serving an imported model: `predict_flavor: 'joblib'` is what the
|
||||
artifact actually is, and it loads, predicts **and** retrains — `log_model` writes that flavor too,
|
||||
as `model.pkl` + `code/`. The artifact it writes carries no `MLmodel` manifest, so an imported model
|
||||
stays readable only by this runtime, before and after a retrain. Full contract:
|
||||
[`model-flavors.md`](model-flavors.md).
|
||||
|
||||
## 14. Retry, and what does not retry
|
||||
|
||||
Only the status writes retry, and only because `UPDATE ... SET <fixed values> WHERE id` leaves the row
|
||||
the same however many times it runs:
|
||||
|
||||
| Call | Attempts | Why |
|
||||
|---|---|---|
|
||||
| claim, terminal write | 10, 1 s → 30 s backoff (≈151 s envelope) | the verdict must land; the database being briefly unavailable is exactly what a retry fixes |
|
||||
| progress writes, names | 3, capped at 2 s | a hint nobody reads — the terminal write rewrites it — must not stall a working import |
|
||||
| everything else | 1 | a repeated registry call or Mongo insert can write twice |
|
||||
|
||||
`ImportLogRowNotClaimableError` and `ImportStatusRejectedError` are non-retryable: the first would
|
||||
read the same rows and reach the same conclusion, and the second means the statement itself was
|
||||
refused (a `CHECK` violation, an undefined column, a value too long), which a second attempt cannot
|
||||
change.
|
||||
|
||||
There is no rollback, no compensation, no resume and no reconciliation. A failed import stays failed;
|
||||
trying again is a new upload and a new row.
|
||||
|
||||
## 15. What a failure looks like in the logs
|
||||
|
||||
Every failure — a row that cannot be claimed, a gate that rejected the bundle, an MLflow that refused
|
||||
the connection, a verdict that could not be written — goes through the one channel in
|
||||
`laborious/utils/models/import_errors.py`, which produces exactly two `logfmt` lines: the technical
|
||||
one, whose message
|
||||
carries the traceback (or the failure chain, root cause first, for anything that crossed an activity
|
||||
boundary), and the user-facing sentence the notification logs. **Both carry the same fields**, so one
|
||||
query returns the pair:
|
||||
|
||||
| Field | What it holds |
|
||||
|---|---|
|
||||
| `import_run_id` | the record's primary key — the id a user quotes to support |
|
||||
| `import_workflow_id` | the Temporal workflow id |
|
||||
| `import_step` | the step that failed, or `not_claimable` / `status_write` |
|
||||
| `import_code` | the code of § 8, the same one written to the row |
|
||||
| `error_type` | the failing exception's own type — `MlflowException`, not the `ApplicationError` wrapping it |
|
||||
| `error_summary` | one bounded line of what it said |
|
||||
| `import_gate` | the gate that rejected the bundle, when one did |
|
||||
| `import_verdict_not_recorded` | the verdict a failed terminal write could not store |
|
||||
|
||||
The workflow metadata every other line already carries (`pod_id`, `runtime`, `model_id`, `model_name`,
|
||||
`workflow_name`, `schedule_name`) is kept alongside them.
|
||||
|
||||
A failure at any pipeline step is reported from the workflow's `except`, before cleanup and before the
|
||||
terminal write, so the log reads in the order things happened and the line's step can never disagree
|
||||
with the row's. That report never changes the outcome: its own failure is swallowed, and the failure it
|
||||
was reporting is re-raised untouched. The Temporal SDK still prints its own traceback for each failed
|
||||
attempt — that is the per-attempt record, and it is not structured; the channel's line is the one to
|
||||
query.
|
||||
|
||||
## 16. How this contract is verified
|
||||
|
||||
Two suites, and they prove different things.
|
||||
|
||||
**Unit** — `pytest tests/`, part of the standard gate (`validate --project-name=laborious`). Covers
|
||||
the bundle reader against committed golden bytes, every gate, the failure catalog, the status SQL and
|
||||
the workflow's ordering, all against fakes.
|
||||
|
||||
**End to end** — `pytest e2e -m import_e2e`, opt-in and outside `testpaths`. Runs the real workflow on
|
||||
a real worker against real MinIO, a real Postgres holding this document's `public.experiment_run`
|
||||
(with and without `V12`), a real MongoDB with the unique index on `models.id`, and a real
|
||||
`mlflow server`. Nothing on the import path is mocked. This is what proves the claims a fake cannot:
|
||||
that the `CHECK` constraints of § 6 accept every status written, that the registry source of § 11
|
||||
resolves through the existing loader, that the `V12` failure shape of § 7 is what it says, and that no
|
||||
plaintext password reaches the Temporal event history.
|
||||
|
||||
Requires Docker plus one one-time setup per machine:
|
||||
|
||||
```bash
|
||||
echo "ryuk.disabled=true" >> ~/.testcontainers.properties
|
||||
```
|
||||
|
||||
Without it every e2e test fails in fixture setup with a Docker mount error — see
|
||||
[`README.md` § End-to-End Tests](../README.md#end-to-end-tests-e2e) for why and what it costs.
|
||||
|
||||
Scenario catalogue: [`e2e/scenarios.md`](../e2e/scenarios.md) § 4.
|
||||
Requirement-to-test mapping: `openspec/changes/import-model-e2e-suite/traceability.md`.
|
||||
|
||||
Doing one import by hand, standing in for the frontend that does not exist yet:
|
||||
[`scripts/import_model_manual_run.py`](../scripts/import_model_manual_run.py) — see
|
||||
[`README.md` § Running One Model Import by Hand](../README.md#running-one-model-import-by-hand).
|
||||
@@ -1,6 +1,8 @@
|
||||
# OPC UA communication (Laborious)
|
||||
|
||||
Laborious exports predictions to OPC UA servers through `OpcRepository` ([`laborious/utils/repository/opc_repository.py`](../laborious/utils/repository/opc_repository.py)) and the Temporal activity layer in [`laborious/activities/opc.py`](../laborious/activities/opc.py).
|
||||
Laborious exports predictions to OPC UA servers through `OpcRepository` ([`laborious/utils/repository/opc_repository.py`](../laborious/utils/repository/opc_repository.py)) and the synchronous Temporal activity layer in [`laborious/activities/opc.py`](../laborious/activities/opc.py). The repository uses `asyncua.sync.Client` (asyncio on a background thread) so activities remain blocking without `async def`.
|
||||
|
||||
OPC reconnect, write error classification (`opc_error_kind`), and activity confidence/comment behavior are converted from the **async** implementation on `main` at `fcc8920a8be4` (`asyncua.Client` + `asyncio` reconnect task → `threading` reconnect thread). Re-convert with `scripts/convert_opc_async_to_sync.py` when `main` OPC files change.
|
||||
|
||||
Implementation plan for session/channel recovery on Tier-1 `Bad*` errors: [`.cursor/plans/opc_bad_reconnect_ac4c6045.plan.md`](../.cursor/plans/opc_bad_reconnect_ac4c6045.plan.md).
|
||||
|
||||
@@ -12,7 +14,7 @@ Worker (long-lived)
|
||||
├── connect / disconnect / validate_connection (read-only)
|
||||
├── _connect_locked / _reconnect_locked (under _connection_lock)
|
||||
├── write_data (single attempt per call)
|
||||
└── background reconnect on Tier-1 Bad*, protocol closed, or stale session
|
||||
└── background reconnect on Tier-1 Bad*, closed protocol, or stale session
|
||||
|
||||
Temporal activity write_opc_data
|
||||
└── OPC.manage_output_tags → write_data per tag (sequential per activity)
|
||||
@@ -26,9 +28,9 @@ One worker process holds one `OpcRepository` instance per configured server. Mul
|
||||
|-------|----------|
|
||||
| Startup | `init_opc()` creates repositories and calls `connect()` → `_connect_locked()` |
|
||||
| Steady state | `validate_connection()` is read-only (`protocol.state` only); `_session_ready` is checked in `write_data` |
|
||||
| Tier-1 Bad* / protocol closed | `_start_reconnect` → `_run_reconnect` → `_reconnect_locked()` (respects `reconnection_interval`) |
|
||||
| Write | `write_data()` checks reconnect task, `_session_ready`, validates protocol, then one `get_node` + `write_value` |
|
||||
| Shutdown | `close()` disconnects all repositories |
|
||||
| Tier-1 Bad* / protocol closed / session not ready | `_start_reconnect(reason)` → `_run_reconnect` (thread) → `_reconnect_locked()` (respects `reconnection_interval`) |
|
||||
| Write | `write_data()` checks in-flight reconnect thread, `_session_ready`, validates, then one `get_node` + `write_value` |
|
||||
| Shutdown | `disconnect()` sets `_allow_reconnect = False`, then tears down session |
|
||||
|
||||
### Session and channel timeouts
|
||||
|
||||
@@ -36,7 +38,7 @@ Requested session and secure-channel lifetime: **10 minutes** (`OPC_UA_SESSION_A
|
||||
|
||||
### Reconnection interval
|
||||
|
||||
`OPC_RECONNECTION_INTERVAL` is in **seconds** (default `120`). It gates **background** reconnect after Tier-1 `Bad*`, closed protocol, or stale session (`last_reconnection_time` is updated only in `_reconnect_locked()`). It limits load on the OPC server when many workflows fail at once.
|
||||
`OPC_RECONNECTION_INTERVAL` is in **seconds** (default `120`). It gates **background** reconnect after Tier-1 `Bad*` (`last_reconnection_time` is updated only in `_reconnect_locked()`). It limits load on the OPC server when many workflows fail at once.
|
||||
|
||||
## Concurrency: connection lock and session readiness
|
||||
|
||||
@@ -44,8 +46,9 @@ To allow **multiple concurrent writes** when the session is healthy, but **block
|
||||
|
||||
| Primitive | Role |
|
||||
|-----------|------|
|
||||
| `_connection_lock` (`asyncio.Lock`) | Held for the entire `disconnect` → `connect` path. Only one connection-maintenance task at a time. |
|
||||
| `_session_ready` (`asyncio.Event`) | Set when a session is ready for writes; cleared before reconnect starts and set again after a successful connect. |
|
||||
| `_connection_lock` (`threading.Lock`) | Held for the entire `disconnect` → `connect` path. Only one connection-maintenance task at a time. |
|
||||
| `_session_ready` (`threading.Event`) | Set when a session is ready for writes; cleared before reconnect starts and set again after a successful connect. |
|
||||
| `_allow_reconnect` | Cleared in `disconnect()` so shutdown does not spawn reconnect threads |
|
||||
|
||||
**Connection methods (caller holds `_connection_lock` for `_*_locked` helpers):**
|
||||
|
||||
@@ -61,33 +64,24 @@ Public `connect()` / `disconnect()` acquire the lock and call `_connect_locked()
|
||||
|
||||
**Write path (`write_data`):**
|
||||
|
||||
1. If a reconnect task is **in flight** → **fail immediately** (`opc_error_kind=reconnect_in_progress`).
|
||||
2. If `_session_ready` is cleared and no task is running → schedule reconnect (`SessionNotReady`); fail with `connection_lost` or `reconnect_in_progress` if a task started.
|
||||
3. `validate_connection()` checks `protocol.state` only (read-only). If closed → schedule reconnect (`ProtocolClosed`) and fail with `opc_error_kind=connection_lost`.
|
||||
4. Single `get_node` + `write_value` (no retry). Tier-1 `Bad*` on write also schedules reconnect.
|
||||
1. If a reconnect **thread** is alive → `reconnect_in_progress`.
|
||||
2. If `_session_ready` is cleared → schedule `SessionNotReady` reconnect; return `reconnect_in_progress` or `connection_lost`.
|
||||
3. If `validate_connection()` fails (protocol closed) → schedule `ProtocolClosed` reconnect; return `connection_lost`.
|
||||
4. Single `get_node` + `write_value` (no retry in the same call).
|
||||
|
||||
**Reconnect path (`_run_reconnect`):**
|
||||
|
||||
1. `_start_reconnect` clears `_session_ready` and schedules the task when the interval allows and `_allow_reconnect` is true.
|
||||
2. `async with _connection_lock:` → `_reconnect_locked()`.
|
||||
1. `_start_reconnect` clears `_session_ready` and starts a daemon thread when the interval allows and `_allow_reconnect` is true.
|
||||
2. `with _connection_lock:` → `_reconnect_locked()`.
|
||||
3. `_session_ready` is set on successful `_open_session()`.
|
||||
4. `disconnect()` sets `_allow_reconnect=False` so shutdown does not respawn sessions.
|
||||
|
||||
A second `_connect_locked()` while a session is already open raises `OpcSessionAlreadyConnectedError` (disconnect first).
|
||||
|
||||
**asyncua note:** Concurrent `write_value` on the same session is only safe if the stack tolerates it. If production shows issues, serialize writes with an optional `asyncio.Semaphore(1)` while keeping the connection lock semantics above.
|
||||
**asyncua note:** Concurrent `write_value` on the same session is only safe if the stack tolerates it. If production shows issues, serialize writes while keeping the connection lock semantics above.
|
||||
|
||||
**Future threads:** replace `asyncio.Lock` / `Event` with `threading` primitives or route all OPC I/O through one dedicated loop.
|
||||
## Tier-1 `Bad*` errors and reconnect
|
||||
|
||||
## Reconnect triggers
|
||||
|
||||
Background reconnect is scheduled when:
|
||||
|
||||
- `validate_connection()` sees a closed or missing protocol (`ProtocolClosed`).
|
||||
- `_session_ready` is clear after a failed reconnect (`SessionNotReady`).
|
||||
- A write raises a Tier-1 `UaStatusCodeError` in `RECONNECTABLE_OPC_BAD_NAMES`.
|
||||
|
||||
For Tier-1 `Bad*` when the server invalidates the session (e.g. `BadSessionIdInvalid`) but the client still sees transport as open, `write_data` fails once, records the OPC status in metrics, and **schedules** reconnect if:
|
||||
When the server invalidates the session (e.g. `BadSessionIdInvalid`) but the client still sees transport as open, `write_data` fails once, records the OPC status in metrics, and **schedules** reconnect if:
|
||||
|
||||
- The exception is a `UaStatusCodeError` whose name is in `RECONNECTABLE_OPC_BAD_NAMES` (see plan), and
|
||||
- `reconnection_interval` has elapsed since `last_reconnection_time`, and
|
||||
|
||||
@@ -1,100 +0,0 @@
|
||||
# The `.sientia` bundle format
|
||||
|
||||
A `.sientia` file is one exported MLflow model: its whole artifact tree, its run parameters and its
|
||||
origin metadata, zipped and then encrypted under a password the exporter chose.
|
||||
|
||||
**It is written by the Streamlit platform, `sientia-projects-templates`** (`app/src/operations/
|
||||
model_export/`). That application owns the export — its screen, its flow, its code — and documents
|
||||
it. This page is the short version: the parameters both sides have to agree on, and what this
|
||||
repository's reader accepts.
|
||||
|
||||
## How it is generated
|
||||
|
||||
The producer downloads the run's artifact tree whole, copies `run.data.params` **verbatim** (with
|
||||
`target_variable` among them, or the import fails at gate 6), writes `metadata.json` first into a
|
||||
`ZIP_DEFLATED` archive with the artifacts following as `artifacts/<relpath>` and no directory
|
||||
entries, encrypts that zip, and offers it as `{model_name}_v{model_version}.sientia`.
|
||||
|
||||
`metadata.json` has exactly two keys: `parameters` — the run's parameters verbatim — and `metadata`,
|
||||
six fields: `model_name`, `experiment_name`, `run_id`, `model_version`, `model_project`,
|
||||
`export_timestamp` (UTC ISO 8601).
|
||||
|
||||
## Encryption parameters
|
||||
|
||||
The reader derives the key with the values the header carries, but refuses anything above the
|
||||
producer's own cost — a header declaring more is a resource-exhaustion attempt, refused before any
|
||||
memory is allocated.
|
||||
|
||||
| | Value |
|
||||
|---|---|
|
||||
| KDF | Argon2id (`crypto_pwhash_ALG_ARGON2ID13`) |
|
||||
| memlimit | `268435456` (256 MiB) — a **ceiling** on the reader's side |
|
||||
| opslimit | `3` — likewise a ceiling |
|
||||
| parallelism | `1` (reserved, not configurable) |
|
||||
| salt | 16 bytes, fresh per bundle |
|
||||
| key | 32 bytes |
|
||||
| AEAD | XChaCha20-Poly1305 `secretstream`, 64 KiB plaintext chunks, last tagged `FINAL` |
|
||||
| digest | SHA-256 of the **pushed chunks only** — on disk, `sha256(file[164:])`, past the 140-byte header and the 24-byte stream header. Hashing from 140 rejects every valid bundle |
|
||||
| password | normalised `NFC`, encoded UTF-8 — no trimming, no case folding |
|
||||
|
||||
The export-side password rule (12+ characters, mixed case, digit, symbol) is enforced there only: by
|
||||
the time a bundle exists, the password is whatever it was, and the reader cannot re-check it.
|
||||
|
||||
## Shape of the file
|
||||
|
||||
`HEADER_STRUCT` is `>8sBBB16sQQB32s64s`. Big-endian, no alignment padding, so the offsets are exact:
|
||||
|
||||
```
|
||||
offset width field constraint the reader enforces
|
||||
------ ----- -------------------- ------------------------------------------------------
|
||||
0 8 magic == b'SIENTIA1'
|
||||
8 1 format_version == 1
|
||||
9 1 aead_id == 1 (XChaCha20-Poly1305 secretstream)
|
||||
10 1 kdf_id == 1 (Argon2id)
|
||||
11 16 salt passed to the KDF as-is
|
||||
27 8 kdf_memlimit_bytes <= 268435456
|
||||
35 8 kdf_opslimit <= 3
|
||||
43 1 kdf_parallelism == 1
|
||||
44 32 ciphertext_digest == sha256(file[164:])
|
||||
76 64 reserved_signature all zero
|
||||
------ -----
|
||||
140 24 secretstream header libsodium HEADERBYTES
|
||||
164 ... chunk stream 64 KiB plaintext chunks, each +17 bytes of AEAD tag
|
||||
```
|
||||
|
||||
`format_version` is checked for equality, not for a range: an unknown version is refused at gate 2
|
||||
rather than parsed optimistically. `reserved_signature` must be all zero today, so a signed bundle
|
||||
will be a new format version and not a silent change of meaning.
|
||||
|
||||
`laborious/utils/bundle/format.py` is the normative layout on this side, and the golden fixture is
|
||||
what keeps it and the producer from drifting apart.
|
||||
|
||||
## How it is read
|
||||
|
||||
`BundleReader.open()` runs seven ordered gates, each mapped to an import step and to one message a
|
||||
person can read:
|
||||
|
||||
| Gate | Step | Refuses |
|
||||
|---|---|---|
|
||||
| 1 | `download` | object too large, wrong prefix or suffix, digest ≠ the one the uploader sent |
|
||||
| 2 | `download` | header that is not this format or not this version |
|
||||
| 3 | `decryption` | wrong password, or altered bytes — the AEAD cannot tell them apart, and the message says so |
|
||||
| 4 | `archive_inspection` | zip bomb, absolute or `..` entries, symlinks — read from the central directory, before extracting |
|
||||
| 5 | `extraction` | anything that escapes the fresh work directory, or overwrites a file |
|
||||
| 6 | `structure_validation` | missing `metadata.json` fields, missing `target_variable`, missing `prediction_model`/`data_model` |
|
||||
| 7 | `content_policy` | a file that is not a known model file by name or suffix |
|
||||
|
||||
Gates 1 and 2 share the `download` step on purpose: neither may borrow the `decryption` sentence,
|
||||
which implicates the password. A gate-6 rejection naming `target_variable` means the origin run never
|
||||
logged it — the fix is a re-export, on the producer's side.
|
||||
|
||||
Error codes, sentences and the ceilings the gates read from configuration:
|
||||
[`model-import.md`](model-import.md).
|
||||
|
||||
## Fixtures
|
||||
|
||||
- `tests/fixtures/bundle/` — a real bundle emitted by the producer, with its password and producer
|
||||
commit recorded in that directory's `README.md`. It pays the producer's real Argon2id cost, so it
|
||||
is used where the real bytes matter and not in every test.
|
||||
- `tests/helpers/bundle_factory.py` — the test-only writer: same header at the same offsets, cheap
|
||||
KDF parameters, and a knob for every hostile variant the gates above refuse.
|
||||
55
docs/sientia_model_drift_jensen_shannon_change.md
Normal file
55
docs/sientia_model_drift_jensen_shannon_change.md
Normal file
@@ -0,0 +1,55 @@
|
||||
# Specification: `sientia_model` — Jensen–Shannon drift (`DriftAnalysis`)
|
||||
|
||||
This document describes what **`sientia_model.analytics.drift_analysis.DriftAnalysis`** should change so downstream consumers (e.g. Laborious `calculate_drift` → Postgres `drift_metrics.value NOT NULL`) no longer receive **NaN** for Jensen–Shannon on valid finite data.
|
||||
|
||||
## Scope
|
||||
|
||||
- **File:** `sientia_model/analytics/drift_analysis.py`
|
||||
- **Method:** `_jensen_shannon_distance(self, ref: np.ndarray, cur: np.ndarray, bins: int = 20) -> float`
|
||||
- **Callers:** `detect_univariate_drift` uses this for the `jensen_shannon` method; results are written to `value` in the drift dataframe.
|
||||
|
||||
## Problem
|
||||
|
||||
The current implementation:
|
||||
|
||||
1. Builds bin edges from **`ref`** only: `np.histogram(ref, bins=bins, density=True)`.
|
||||
2. Builds the chunk histogram with **`density=True`** on the same edges: `np.histogram(cur, bins=edges, density=True)`.
|
||||
|
||||
When **every** value in **`cur`** falls **outside** the closed support implied by those edges (typical case: production chunk drifted above the reference max or below the reference min), NumPy yields **all-zero counts** for `cur`. With **`density=True`**, normalization does **0/0**, producing **NaN** for the whole histogram, which propagates to **`float('nan')`** in `detect_univariate_drift` → **SQL NULL** where `value` is `NOT NULL`.
|
||||
|
||||
This appears in Laborious when:
|
||||
|
||||
- `model_config.target` excludes the main target column from univariate features, so another feature (e.g. `sensor_2`) is compared chunk-by-chunk against the full reference series for that feature.
|
||||
- Reference and current ranges do not overlap for some chunks (strong drift or different scaling).
|
||||
|
||||
Other methods (`kolmogorov_smirnov`, `wasserstein`) do not use this histogram+density path, so they can stay finite while **Jensen–Shannon** alone becomes null.
|
||||
|
||||
## Required behavior
|
||||
|
||||
1. **Finite output** for finite `ref` and `cur` after removing non-finite values, whenever both sides have **at least one** usable sample.
|
||||
2. **Explicit handling of out-of-range chunk mass:** probability mass from `cur` that does not fall into any bin defined from `ref` must still be represented (so the chunk distribution sums to 1), analogous to NannyML’s continuous JS approach (tail / “leftover” mass).
|
||||
3. **Missing values:** drop `NaN` from `ref` and `cur` before computing. If either side is **empty** after that, return **`float('nan')`** (callers may filter or map; schema may still forbid null — product decision outside this spec).
|
||||
|
||||
## Recommended algorithm (replace current body)
|
||||
|
||||
1. `ref = np.asarray(ref, float); cur = np.asarray(cur, float)`.
|
||||
2. `ref = ref[~np.isnan(ref)]; cur = cur[~np.isnan(cur)]`.
|
||||
3. If `ref.size == 0` or `cur.size == 0`: return `float('nan')`.
|
||||
4. `hist_ref, edges = np.histogram(ref, bins=bins)` — **counts**, not `density=True`.
|
||||
5. `p = hist_ref.astype(float) / ref.size` (reference bin probabilities).
|
||||
6. `hist_cur, _ = np.histogram(cur, bins=edges)`; `q = hist_cur.astype(float) / cur.size`.
|
||||
7. `leftover = 1.0 - float(np.sum(q))`. If `leftover > 1e-15` (tolerance for float noise), append **`leftover`** to `q` and **`0.0`** to `p` so both remain proper discrete distributions over the same extended support.
|
||||
8. Apply small smoothing (existing module constant `EPSILON` is fine): add `EPSILON` to `p` and `q`, renormalize each to sum 1.
|
||||
9. `m = 0.5 * (p + q)`; compute symmetric JS via KL terms as today, e.g. `inner = 0.5 * (sum(p*log(p/m)) + sum(q*log(q/m)))`.
|
||||
10. Return `sqrt(max(inner, 0.0))` to guard against tiny negative `inner` from floating-point error.
|
||||
|
||||
## Non-goals / notes
|
||||
|
||||
- **Numerical parity** with the old `density=True` implementation is not required; parity with **NannyML** or **scipy** JS is desirable but optional. The priority is **finite, interpretable** drift when the chunk is outside the reference histogram range.
|
||||
- **Multivariate** drift in the same file is unchanged by this spec.
|
||||
- **Tests** in `sientia_model` should cover: (a) chunk entirely above reference max, (b) entirely below reference min, (c) overlapping range, (d) `ref` or `cur` all-NaN after cleaning.
|
||||
|
||||
## Reference (external)
|
||||
|
||||
- NannyML continuous JS uses count-based bin probabilities and a **leftover** mass bin; see `ContinuousJensenShannonDistance` in `nannyml/drift/univariate/methods.py` (`_calculate`, `leftover = 1 - np.sum(...)`).
|
||||
- Historical NaN-in-reference issue: [NannyML#339](https://github.com/NannyML/nannyml/issues/339) / [#340](https://github.com/NannyML/nannyml/pull/340) (orthogonal to out-of-range mass, but relevant for input cleaning).
|
||||
1117
e2e/conftest.py
1117
e2e/conftest.py
File diff suppressed because it is too large
Load Diff
110
e2e/db_schema.sql
Normal file
110
e2e/db_schema.sql
Normal file
@@ -0,0 +1,110 @@
|
||||
-- =============================================================================
|
||||
-- E2E test database schema for the ``sientia_data`` namespace.
|
||||
--
|
||||
-- Mirrors the production DDL one-to-one so any change in production can be
|
||||
-- pasted directly into this file. The conftest fixture loads this SQL into the
|
||||
-- testcontainers Postgres before each test run.
|
||||
--
|
||||
-- Notes on differences from production:
|
||||
-- * Tables that are partitioned in production (e.g. ``simple_metrics``,
|
||||
-- ``transformed_data``, ``drift_metrics``) are created as plain tables
|
||||
-- here because the test suite does not exercise partition pruning.
|
||||
-- * Indexes are intentionally omitted; tests rely on functional behavior,
|
||||
-- not query plans.
|
||||
-- =============================================================================
|
||||
|
||||
CREATE SCHEMA IF NOT EXISTS sientia_data;
|
||||
|
||||
-- -----------------------------------------------------------------------------
|
||||
-- sientia_data.laborious_data
|
||||
-- -----------------------------------------------------------------------------
|
||||
CREATE TABLE IF NOT EXISTS sientia_data.laborious_data (
|
||||
model_id int4 NOT NULL,
|
||||
variable text NOT NULL,
|
||||
value numeric NULL,
|
||||
"timestamp" timestamptz NOT NULL,
|
||||
created_at timestamptz DEFAULT CURRENT_TIMESTAMP NOT NULL,
|
||||
CONSTRAINT unique_timestamp_variable
|
||||
UNIQUE (model_id, "timestamp", variable)
|
||||
);
|
||||
|
||||
-- -----------------------------------------------------------------------------
|
||||
-- sientia_data.predictions
|
||||
-- -----------------------------------------------------------------------------
|
||||
CREATE TABLE IF NOT EXISTS sientia_data.predictions (
|
||||
model_id int4 NOT NULL,
|
||||
prediction numeric NULL,
|
||||
prediction_confidence numeric NOT NULL,
|
||||
response_time numeric NOT NULL,
|
||||
prediction_status text NOT NULL,
|
||||
"timestamp" timestamptz NOT NULL,
|
||||
created_at timestamptz DEFAULT CURRENT_TIMESTAMP NOT NULL,
|
||||
"comments" text NULL,
|
||||
metadata_id varchar(24) NULL,
|
||||
CONSTRAINT unique_model_id_timestamp
|
||||
UNIQUE (model_id, "timestamp")
|
||||
);
|
||||
|
||||
-- -----------------------------------------------------------------------------
|
||||
-- sientia_data.transformed_data
|
||||
-- Production: PARTITION BY RANGE (created_at). Tests use a plain table.
|
||||
-- -----------------------------------------------------------------------------
|
||||
CREATE TABLE IF NOT EXISTS sientia_data.transformed_data (
|
||||
id SERIAL NOT NULL,
|
||||
model_id int4 NOT NULL,
|
||||
variable text NOT NULL,
|
||||
value numeric NULL,
|
||||
"timestamp" timestamptz NOT NULL,
|
||||
created_at timestamptz DEFAULT CURRENT_TIMESTAMP NOT NULL,
|
||||
PRIMARY KEY (id, created_at)
|
||||
);
|
||||
|
||||
-- -----------------------------------------------------------------------------
|
||||
-- sientia_data.drift_metrics
|
||||
-- Production: PARTITION BY RANGE (created_at). Tests use a plain table.
|
||||
-- -----------------------------------------------------------------------------
|
||||
CREATE TABLE IF NOT EXISTS sientia_data.drift_metrics (
|
||||
id SERIAL NOT NULL,
|
||||
model_id text NOT NULL,
|
||||
feature text NULL,
|
||||
method text NOT NULL,
|
||||
value numeric NOT NULL,
|
||||
alert bool NOT NULL,
|
||||
chunk_index int4 NOT NULL,
|
||||
chunk_start_date text NOT NULL,
|
||||
chunk_end_date text NOT NULL,
|
||||
accurate bool NOT NULL,
|
||||
"timestamp" timestamptz NULL,
|
||||
created_at timestamptz DEFAULT CURRENT_TIMESTAMP NOT NULL,
|
||||
PRIMARY KEY (id, created_at)
|
||||
);
|
||||
|
||||
-- -----------------------------------------------------------------------------
|
||||
-- sientia_data.simple_metrics
|
||||
-- Production: PARTITION BY RANGE (created_at). Tests use a plain table.
|
||||
-- -----------------------------------------------------------------------------
|
||||
CREATE TABLE IF NOT EXISTS sientia_data.simple_metrics (
|
||||
id SERIAL NOT NULL,
|
||||
model_id text NOT NULL,
|
||||
metric text NOT NULL,
|
||||
value numeric NOT NULL,
|
||||
"timestamp" timestamptz NULL,
|
||||
data_size int4 NOT NULL,
|
||||
interval_minutes int4 NOT NULL,
|
||||
created_at timestamptz DEFAULT CURRENT_TIMESTAMP NOT NULL,
|
||||
PRIMARY KEY (id, created_at)
|
||||
);
|
||||
|
||||
-- -----------------------------------------------------------------------------
|
||||
-- sientia_data.log_retrain
|
||||
-- No primary key in production; all columns nullable.
|
||||
-- -----------------------------------------------------------------------------
|
||||
CREATE TABLE IF NOT EXISTS sientia_data.log_retrain (
|
||||
mlflow_experiment_id int8 NULL,
|
||||
mlflow_run_id text NULL,
|
||||
model_id text NULL,
|
||||
model_name text NULL,
|
||||
status text NULL,
|
||||
"timestamp" timestamptz NULL,
|
||||
"version" text NULL
|
||||
);
|
||||
@@ -1,65 +0,0 @@
|
||||
-- public.experiment_run as migrations V1 … V11 leave it, for the e2e Postgres container.
|
||||
--
|
||||
-- Source of truth: sientia-dataops-database-migrations,
|
||||
-- postgres/sientia-core-mlops-bff/migrations/ — this file mirrors the *final state* of:
|
||||
--
|
||||
-- V1 create_table_experiment_run (13 columns, one index)
|
||||
-- V2 update_check_constraint_on_status (named the status CHECK)
|
||||
-- V3 update_check_constraints (named every other CHECK)
|
||||
-- V4 update_status_column_size_and_add_orchestrator_validation_error (status -> VARCHAR(50))
|
||||
-- V5 migrate_mage_statuses_to_orchestrator (data only)
|
||||
-- V6 rename_mage_response_data_column (-> orchestrator_response_data)
|
||||
-- V7 migrate_mlflow_statuses_to_tracking (data only)
|
||||
-- V8 remove_orchestrator_request_sent_status (status vocabulary)
|
||||
-- V9 remove_file_upload_statuses (status vocabulary)
|
||||
-- V10 remove_unused_statuses (status vocabulary)
|
||||
-- V11 add_run_type_to_experiment_run (run_type + reindexed)
|
||||
--
|
||||
-- V12 (add_import_statuses_to_experiment_run) lives in `experiment_run_v12.sql`, applied on top of
|
||||
-- this file. It is separate because one e2e test — the evidence for why the migration gates
|
||||
-- enabling the import worker — has to run against a table that does *not* have it, and a second
|
||||
-- whole-schema copy with V12 removed would rot the moment either copy were edited.
|
||||
--
|
||||
-- The same drift trade already accepted for the bundle fixtures: this is a copy, and the tests
|
||||
-- that read it are what make the copy auditable. If a migration lands there, update this file and
|
||||
-- the version list above. Nothing in this repository owns the DDL — no import ever executes it
|
||||
-- outside the test suite.
|
||||
|
||||
CREATE TABLE IF NOT EXISTS public.experiment_run (
|
||||
id BIGSERIAL CONSTRAINT pk_experiment_run PRIMARY KEY,
|
||||
experiment_name VARCHAR(50) NOT NULL,
|
||||
run_name VARCHAR(50),
|
||||
username VARCHAR(320) NOT NULL,
|
||||
status VARCHAR(50) NOT NULL,
|
||||
error_message VARCHAR(1024),
|
||||
created_at TIMESTAMP WITHOUT TIME ZONE NOT NULL,
|
||||
updated_at TIMESTAMP WITHOUT TIME ZONE NOT NULL,
|
||||
bucket_name VARCHAR(63),
|
||||
file_name VARCHAR(1024),
|
||||
request_data JSONB NOT NULL,
|
||||
orchestrator_response_data JSONB,
|
||||
run_type VARCHAR(20) NOT NULL DEFAULT 'TRAIN',
|
||||
|
||||
CONSTRAINT chk_experiment_run_experiment_name_min_length CHECK (LENGTH(experiment_name) >= 3),
|
||||
CONSTRAINT chk_experiment_run_run_name_min_length CHECK (LENGTH(run_name) >= 3),
|
||||
CONSTRAINT chk_experiment_run_username_min_length CHECK (LENGTH(username) >= 3),
|
||||
CONSTRAINT chk_experiment_run_username_pattern
|
||||
CHECK (username ~* '^[a-zA-Z0-9_!#$%&''*+/=?`{|}~^.-]+@[a-zA-Z0-9.-]+$'),
|
||||
CONSTRAINT chk_experiment_run_error_message_min_length CHECK (LENGTH(error_message) >= 3),
|
||||
CONSTRAINT chk_experiment_run_bucket_name_min_length CHECK (LENGTH(bucket_name) >= 3),
|
||||
CONSTRAINT chk_experiment_run_bucket_name_pattern CHECK (bucket_name ~* '^[a-z0-9-]+$'),
|
||||
CONSTRAINT chk_experiment_run_file_name_min_length CHECK (LENGTH(file_name) >= 3),
|
||||
CONSTRAINT chk_experiment_run_file_name_pattern
|
||||
CHECK (file_name ~* '^[a-zA-Z0-9 /._-]+$'),
|
||||
|
||||
-- The V1…V11 status vocabulary: the two shared stage values plus the training terminals.
|
||||
-- V12 adds IMPORT_SUCCESS and IMPORT_ERROR to this list.
|
||||
CONSTRAINT chk_experiment_run_status CHECK (status IN ('PENDING', 'ORCHESTRATOR_REQUEST_ERROR',
|
||||
'ORCHESTRATOR_WAITING_PROC',
|
||||
'ORCHESTRATOR_VALIDATION_ERROR',
|
||||
'TRAINING_SUCCESS', 'TRAINING_ERROR')),
|
||||
CONSTRAINT chk_experiment_run_run_type CHECK (run_type IN ('TRAIN', 'IMPORT'))
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_experiment_run_user_exp_status_date_type
|
||||
ON public.experiment_run (username, experiment_name, status, created_at, run_type);
|
||||
@@ -1,36 +0,0 @@
|
||||
-- V12 add_import_statuses_to_experiment_run, as a delta applied on top of `experiment_run.sql`.
|
||||
--
|
||||
-- Source of truth: sientia-dataops-database-migrations,
|
||||
-- postgres/sientia-core-mlops-bff/migrations/V12__add_import_statuses_to_experiment_run.sql
|
||||
--
|
||||
-- Two things, and nothing else:
|
||||
--
|
||||
-- 1. the status vocabulary gains IMPORT_SUCCESS and IMPORT_ERROR;
|
||||
-- 2. chk_experiment_run_terminal_status_run_type pairs each terminal value with its own
|
||||
-- run_type, so a TRAIN row can never end IMPORT_SUCCESS and an IMPORT row can never end
|
||||
-- TRAINING_ERROR. The two shared stage values (PENDING, ORCHESTRATOR_WAITING_PROC) stay out
|
||||
-- of the pairing — both run types legitimately pass through them.
|
||||
--
|
||||
-- This file is applied separately because one e2e test runs an import against a table *without*
|
||||
-- it: that failure shape — provisioned, then unrecordable, refused on the first attempt — is the
|
||||
-- evidence for why this migration has to land before IMPORT_MODEL_ENABLED is turned on anywhere.
|
||||
|
||||
ALTER TABLE public.experiment_run
|
||||
DROP CONSTRAINT IF EXISTS chk_experiment_run_status;
|
||||
|
||||
ALTER TABLE public.experiment_run
|
||||
ADD CONSTRAINT chk_experiment_run_status CHECK (status IN ('PENDING',
|
||||
'ORCHESTRATOR_REQUEST_ERROR',
|
||||
'ORCHESTRATOR_WAITING_PROC',
|
||||
'ORCHESTRATOR_VALIDATION_ERROR',
|
||||
'TRAINING_SUCCESS', 'TRAINING_ERROR',
|
||||
'IMPORT_SUCCESS', 'IMPORT_ERROR'));
|
||||
|
||||
ALTER TABLE public.experiment_run
|
||||
DROP CONSTRAINT IF EXISTS chk_experiment_run_terminal_status_run_type;
|
||||
|
||||
ALTER TABLE public.experiment_run
|
||||
ADD CONSTRAINT chk_experiment_run_terminal_status_run_type CHECK (
|
||||
(run_type = 'IMPORT' OR status NOT IN ('IMPORT_SUCCESS', 'IMPORT_ERROR'))
|
||||
AND (run_type = 'TRAIN' OR status NOT IN ('TRAINING_SUCCESS', 'TRAINING_ERROR'))
|
||||
);
|
||||
236
e2e/helpers.py
236
e2e/helpers.py
@@ -3,17 +3,61 @@ Shared helpers for E2E tests (Temporal workflows + PostgreSQL).
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from datetime import datetime
|
||||
from decimal import Decimal
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.engine import Engine
|
||||
|
||||
SCENARIO_INPUTS_DIR = Path(__file__).parent / 'scenario_inputs'
|
||||
|
||||
async def start_and_await_workflow(
|
||||
client, workflow_run, input_data: dict, workflow_id: str, timeout: float = 60.0
|
||||
):
|
||||
|
||||
def _replace_template_values(payload: Any, model_id: int) -> Any:
|
||||
"""
|
||||
Replace string placeholders in scenario payloads with the concrete model id.
|
||||
|
||||
Args:
|
||||
payload: JSON-like structure loaded from scenario input file.
|
||||
model_id: Model id used to render template placeholders.
|
||||
|
||||
Return:
|
||||
Any: Payload with ``{{MODEL_ID}}`` replaced where applicable.
|
||||
"""
|
||||
if isinstance(payload, dict):
|
||||
return {key: _replace_template_values(value, model_id) for key, value in payload.items()}
|
||||
if isinstance(payload, list):
|
||||
return [_replace_template_values(item, model_id) for item in payload]
|
||||
if isinstance(payload, str):
|
||||
if payload == '{{MODEL_ID}}':
|
||||
return model_id
|
||||
return payload.replace('{{MODEL_ID}}', str(model_id))
|
||||
return payload
|
||||
|
||||
|
||||
def load_scenario_input(file_name: str, model_id: int | None = None) -> dict[str, Any]:
|
||||
"""
|
||||
Load a scenario input JSON from ``e2e/scenario_inputs``.
|
||||
|
||||
Args:
|
||||
file_name: JSON file name inside ``e2e/scenario_inputs``.
|
||||
model_id: Optional model id used to render ``{{MODEL_ID}}`` placeholders.
|
||||
|
||||
Return:
|
||||
dict[str, Any]: Input payload ready to be passed to workflow/activity calls.
|
||||
"""
|
||||
file_path = SCENARIO_INPUTS_DIR / file_name
|
||||
with file_path.open('r', encoding='utf-8') as f:
|
||||
payload = json.load(f)
|
||||
|
||||
if model_id is not None:
|
||||
return _replace_template_values(payload, model_id)
|
||||
return payload
|
||||
|
||||
|
||||
async def start_and_await_workflow(client, workflow_run, input_data: dict, workflow_id: str, timeout: float = 60.0):
|
||||
"""
|
||||
Start a workflow and wait for its result.
|
||||
|
||||
@@ -36,7 +80,17 @@ async def start_and_await_workflow(
|
||||
return await asyncio.wait_for(handle.result(), timeout=timeout)
|
||||
|
||||
|
||||
def insert_sample_data(postgres_engine: Engine, model_id: int, values: list[Any]) -> None:
|
||||
DEFAULT_BATCH_TIMESTAMP = '2024-01-01 12:00:00+00:00'
|
||||
DEFAULT_PREDICTION_HISTORY_TIMESTAMP = '2024-01-01 12:00:00+00:00'
|
||||
|
||||
|
||||
def insert_sample_data(
|
||||
postgres_engine: Engine,
|
||||
model_id: int,
|
||||
values: list[Any],
|
||||
*,
|
||||
data_timestamp: str = DEFAULT_BATCH_TIMESTAMP,
|
||||
) -> None:
|
||||
"""
|
||||
Replace laborious_data rows for a model_id with one row per value (sensor_1..n).
|
||||
|
||||
@@ -44,31 +98,113 @@ def insert_sample_data(postgres_engine: Engine, model_id: int, values: list[Any]
|
||||
postgres_engine: SQLAlchemy engine.
|
||||
model_id: Model id column value.
|
||||
values: Per-sensor values; use string 'NULL' for SQL NULL.
|
||||
data_timestamp: Timestamp and created_at for every inserted row; drives
|
||||
``last_timestamp`` on the MinIO/query payload (max row time).
|
||||
"""
|
||||
with postgres_engine.begin() as conn:
|
||||
conn.execute(
|
||||
text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}')
|
||||
)
|
||||
conn.execute(text(f'DELETE FROM sientia_data.laborious_data WHERE model_id = {model_id}'))
|
||||
values_sql = []
|
||||
for i, value in enumerate(values):
|
||||
values_sql.append(f"""
|
||||
({model_id}, 'sensor_{i + 1}', {value}, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
|
||||
({model_id}, 'sensor_{i + 1}', {value}, '{data_timestamp}', '{data_timestamp}')
|
||||
""")
|
||||
insert_sql = f"""
|
||||
INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
|
||||
INSERT INTO sientia_data.laborious_data (model_id, variable, value, timestamp, created_at)
|
||||
VALUES
|
||||
{', '.join(values_sql)}
|
||||
"""
|
||||
conn.execute(text(insert_sql))
|
||||
|
||||
|
||||
def insert_sample_prediction(
|
||||
postgres_engine: Engine,
|
||||
model_id: int,
|
||||
*,
|
||||
prediction_timestamp: str = DEFAULT_PREDICTION_HISTORY_TIMESTAMP,
|
||||
) -> tuple[int, Decimal, Decimal, str]:
|
||||
"""
|
||||
Insert a single historical prediction row for REPEAT scenarios.
|
||||
|
||||
Args:
|
||||
postgres_engine: SQLAlchemy engine.
|
||||
model_id: Model id.
|
||||
prediction_timestamp: Row ``timestamp`` (unique with model_id in tests).
|
||||
|
||||
Return:
|
||||
tuple: (model_id, prediction, prediction_confidence, prediction_status) for assertions.
|
||||
"""
|
||||
with postgres_engine.begin() as conn:
|
||||
conn.execute(text(f'DELETE FROM sientia_data.predictions WHERE model_id = {model_id}'))
|
||||
insert_sql = f"""
|
||||
INSERT INTO sientia_data.predictions (
|
||||
model_id, timestamp, prediction, prediction_confidence, prediction_status, comments, response_time
|
||||
)
|
||||
VALUES (
|
||||
{model_id}, '{prediction_timestamp}', 10, 0, 'Good', '', 0.1
|
||||
)
|
||||
"""
|
||||
conn.execute(text(insert_sql))
|
||||
return (model_id, Decimal(10), Decimal(0), 'Good')
|
||||
|
||||
|
||||
def workflow_failure_message_chain(exc: BaseException) -> list[str]:
|
||||
"""
|
||||
Collect ``str()`` / ``message`` from an exception and its ``__cause__`` chain.
|
||||
|
||||
Args:
|
||||
exc: Root exception (e.g. from ``pytest.raises``).
|
||||
|
||||
Return:
|
||||
list[str]: Messages from root to innermost cause.
|
||||
"""
|
||||
messages: list[str] = []
|
||||
current: BaseException | None = exc
|
||||
while current is not None:
|
||||
messages.append(getattr(current, 'message', None) or str(current) or repr(current))
|
||||
current = current.__cause__
|
||||
return messages
|
||||
|
||||
|
||||
def assert_postgres_unique_violation_in_chain(exc: BaseException) -> None:
|
||||
"""
|
||||
Assert the exception chain mentions Postgres unique-constraint violation.
|
||||
|
||||
Args:
|
||||
exc: Workflow or activity error from Temporal.
|
||||
|
||||
Raises:
|
||||
AssertionError: If no link in the chain looks like UniqueViolation.
|
||||
"""
|
||||
chain = ' | '.join(workflow_failure_message_chain(exc))
|
||||
assert 'UniqueViolation' in chain or 'unique_model_id_timestamp' in chain, (
|
||||
f'Expected unique constraint violation in error chain, got: {chain}'
|
||||
)
|
||||
|
||||
|
||||
def assert_prediction_row_count(postgres_engine: Engine, model_id: int, expected: int) -> None:
|
||||
"""
|
||||
Assert how many prediction rows exist for a model_id.
|
||||
|
||||
Args:
|
||||
postgres_engine: SQLAlchemy engine.
|
||||
model_id: Model id filter.
|
||||
expected: Expected row count.
|
||||
"""
|
||||
with postgres_engine.connect() as conn:
|
||||
n = conn.execute(
|
||||
text('SELECT COUNT(*) FROM sientia_data.predictions WHERE model_id = :m'),
|
||||
{'m': model_id},
|
||||
).scalar()
|
||||
assert n == expected, f'Expected {expected} prediction rows, got {n}'
|
||||
|
||||
|
||||
def assert_prediction(
|
||||
postgres_engine: Engine,
|
||||
model_id: int,
|
||||
prediction: float = 0.5,
|
||||
prediction_confidence: int | Decimal = 0,
|
||||
prediction_status: str = 'Good',
|
||||
comments: str | None = None,
|
||||
comments: str = '',
|
||||
comments_contains: str | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
@@ -80,22 +216,19 @@ def assert_prediction(
|
||||
prediction: Expected prediction value.
|
||||
prediction_confidence: Expected confidence (int or Decimal for numeric column).
|
||||
prediction_status: Expected status string.
|
||||
comments: Expected exact comments string (optional).
|
||||
comments_contains: Substring expected in comments when queued (optional).
|
||||
comments: Expected exact comments string (ignored when ``comments_contains`` is set).
|
||||
comments_contains: When set, assert this substring appears in comments.
|
||||
"""
|
||||
|
||||
with postgres_engine.connect() as conn:
|
||||
result_query = conn.execute(
|
||||
text(
|
||||
f'SELECT model_id, prediction, prediction_confidence, prediction_status, comments '
|
||||
f'FROM predictions_schema.predictions WHERE model_id = {model_id} '
|
||||
f'FROM sientia_data.predictions WHERE model_id = {model_id} '
|
||||
f'ORDER BY created_at ASC'
|
||||
)
|
||||
)
|
||||
prediction_rows = result_query.fetchall()
|
||||
assert len(prediction_rows) == 1, (
|
||||
f'Expected one prediction record, got {len(prediction_rows)}'
|
||||
)
|
||||
assert len(prediction_rows) == 1, f'Expected one prediction record, got {len(prediction_rows)}'
|
||||
row = prediction_rows[0]
|
||||
assert row[0] == model_id, f'Expected model_id={model_id}, got {row[0]}'
|
||||
assert row[1] == prediction or Decimal(str(row[1])) == Decimal(str(prediction)), (
|
||||
@@ -104,15 +237,14 @@ def assert_prediction(
|
||||
assert row[2] == prediction_confidence or Decimal(str(row[2])) == Decimal(
|
||||
str(prediction_confidence)
|
||||
), f'Expected prediction_confidence={prediction_confidence}, got {row[2]}'
|
||||
assert row[3] == prediction_status, (
|
||||
f"Expected prediction_status='{prediction_status}', got {row[3]}"
|
||||
)
|
||||
if comments is not None:
|
||||
assert row[4] == comments, f"Expected comments='{comments}', got {row[4]}"
|
||||
assert row[3] == prediction_status, f"Expected prediction_status='{prediction_status}', got {row[3]}"
|
||||
actual_comments = row[4] or ''
|
||||
if comments_contains is not None:
|
||||
assert comments_contains in row[4], (
|
||||
f"Expected comments to contain '{comments_contains}', got {row[4]}"
|
||||
assert comments_contains in actual_comments, (
|
||||
f"Expected comments to contain '{comments_contains}', got '{actual_comments}'"
|
||||
)
|
||||
else:
|
||||
assert actual_comments == comments, f"Expected comments='{comments}', got '{actual_comments}'"
|
||||
|
||||
|
||||
def assert_continue(
|
||||
@@ -126,7 +258,7 @@ def assert_continue(
|
||||
result_query = conn.execute(
|
||||
text(
|
||||
f'SELECT model_id, prediction, prediction_confidence, prediction_status, comments '
|
||||
f'FROM predictions_schema.predictions WHERE model_id = {model_id}'
|
||||
f'FROM sientia_data.predictions WHERE model_id = {model_id}'
|
||||
)
|
||||
)
|
||||
prediction_rows = result_query.fetchall()
|
||||
@@ -142,10 +274,11 @@ def assert_continue(
|
||||
|
||||
def assert_stop(postgres_engine: Engine, model_id: int) -> None:
|
||||
"""Assert no prediction rows for model_id."""
|
||||
import pytest
|
||||
|
||||
with postgres_engine.connect() as conn:
|
||||
result_query = conn.execute(
|
||||
text(f'SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = {model_id}')
|
||||
text(f'SELECT COUNT(*) FROM sientia_data.predictions WHERE model_id = {model_id}')
|
||||
)
|
||||
count = result_query.scalar()
|
||||
assert count == 0, f'Expected no predictions, but found {count} records'
|
||||
@@ -162,12 +295,13 @@ def assert_repeat(postgres_engine: Engine, model_id: int, last_prediction: tuple
|
||||
model_id: Model id.
|
||||
last_prediction: Tuple (model_id, prediction, confidence, status) to match both rows.
|
||||
"""
|
||||
import pytest
|
||||
|
||||
with postgres_engine.connect() as conn:
|
||||
result_query = conn.execute(
|
||||
text(
|
||||
f'SELECT model_id, prediction, prediction_confidence, prediction_status '
|
||||
f'FROM predictions_schema.predictions WHERE model_id = {model_id} '
|
||||
f'FROM sientia_data.predictions WHERE model_id = {model_id} '
|
||||
f'ORDER BY created_at ASC'
|
||||
)
|
||||
)
|
||||
@@ -184,3 +318,51 @@ def assert_repeat(postgres_engine: Engine, model_id: int, last_prediction: tuple
|
||||
def make_workflow_id(prefix: str) -> str:
|
||||
"""Build a unique workflow id using a prefix and current timestamp."""
|
||||
return f'{prefix}-{datetime.now().timestamp()}'
|
||||
|
||||
|
||||
def insert_target_data_for_drift(
|
||||
postgres_engine: Engine,
|
||||
model_id: int,
|
||||
timestamps: list[str],
|
||||
variables_values: dict[str, list[float]],
|
||||
) -> None:
|
||||
"""
|
||||
Insert one row per (timestamp, variable) pair into ``laborious_data``.
|
||||
|
||||
Used by drift scenarios that need wide-format input where the pivot keeps a
|
||||
full row for every timestamp.
|
||||
|
||||
Args:
|
||||
- postgres_engine: SQLAlchemy engine bound to the test container.
|
||||
- model_id: Model id stamped on every row.
|
||||
- timestamps: ISO-8601 strings used both as ``timestamp`` and ``created_at``.
|
||||
- variables_values: Mapping of variable name to a list of values; each list
|
||||
must be the same length as ``timestamps``.
|
||||
"""
|
||||
for var_name, values in variables_values.items():
|
||||
if len(values) != len(timestamps):
|
||||
raise ValueError(
|
||||
f"Variable '{var_name}' has {len(values)} values but {len(timestamps)} timestamps"
|
||||
)
|
||||
|
||||
rows_sql = []
|
||||
for index, ts in enumerate(timestamps):
|
||||
for var_name, values in variables_values.items():
|
||||
rows_sql.append(
|
||||
f"({model_id}, '{var_name}', {values[index]}, '{ts}', '{ts}')"
|
||||
)
|
||||
|
||||
with postgres_engine.begin() as conn:
|
||||
conn.execute(
|
||||
text(f'DELETE FROM sientia_data.laborious_data WHERE model_id = {model_id}')
|
||||
)
|
||||
if rows_sql:
|
||||
conn.execute(
|
||||
text(
|
||||
'INSERT INTO sientia_data.laborious_data '
|
||||
'(model_id, variable, value, "timestamp", created_at) VALUES '
|
||||
+ ', '.join(rows_sql)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
14
e2e/scenario_inputs/drift_base.json
Normal file
14
e2e/scenario_inputs/drift_base.json
Normal file
@@ -0,0 +1,14 @@
|
||||
{
|
||||
"schedule_name": "test-schedule",
|
||||
"model_name": "test_model",
|
||||
"model_id": "{{MODEL_ID}}",
|
||||
"schema": "sientia_data",
|
||||
"source_table_name": "laborious_data",
|
||||
"target_table_name": "drift_metrics",
|
||||
"interval": 60,
|
||||
"drift_metrics": ["kolmogorov_smirnov", "jensen_shannon", "wasserstein"],
|
||||
"chunk_period": "min",
|
||||
"model_config": {
|
||||
"target": "sensor_1"
|
||||
}
|
||||
}
|
||||
37
e2e/scenario_inputs/format_export_base.json
Normal file
37
e2e/scenario_inputs/format_export_base.json
Normal file
@@ -0,0 +1,37 @@
|
||||
{
|
||||
"schedule_name": "test-schedule",
|
||||
"model_name": "test_model",
|
||||
"model_id": "{{MODEL_ID}}",
|
||||
"query": "SELECT timestamp, variable, value, created_at FROM sientia_data.laborious_data WHERE model_id = {{MODEL_ID}}",
|
||||
"schema": "sientia_data",
|
||||
"table_name": "predictions",
|
||||
"transform_table_name": "transformed_data",
|
||||
"input_filters": {
|
||||
"EMPTY_DATA": {
|
||||
"POLICY": "STOP",
|
||||
"CONFIG": {}
|
||||
}
|
||||
},
|
||||
"mlflow_transform_filters": {
|
||||
"API_ERROR": {
|
||||
"POLICY": "STOP",
|
||||
"CONFIG": {}
|
||||
}
|
||||
},
|
||||
"mlflow_predict_filters": {
|
||||
"API_ERROR": {
|
||||
"POLICY": "STOP",
|
||||
"CONFIG": {}
|
||||
}
|
||||
},
|
||||
"path_priority": ["STOP", "CONTINUE", "REPEAT"],
|
||||
"opc_output_config": {},
|
||||
"pi_web_api_output_config": {},
|
||||
"save_transform": true,
|
||||
"prediction_store_policy": "lts:1",
|
||||
"model_config": {
|
||||
"retention_minutes": 0,
|
||||
"target": "sensor_1"
|
||||
},
|
||||
"datetime_columns": ["timestamp", "created_at"]
|
||||
}
|
||||
45
e2e/scenario_inputs/main_happy_path.json
Normal file
45
e2e/scenario_inputs/main_happy_path.json
Normal file
@@ -0,0 +1,45 @@
|
||||
{
|
||||
"metadata": {
|
||||
"metadata": {
|
||||
"model_id": "{{MODEL_ID}}",
|
||||
"model_name": "test_model",
|
||||
"schedule_name": "test-schedule",
|
||||
"workflow_name": "predictions_batch"
|
||||
}
|
||||
},
|
||||
"schedule_name": "test-schedule",
|
||||
"model_name": "test_model",
|
||||
"model_id": "{{MODEL_ID}}",
|
||||
"query": "SELECT timestamp, variable, value, created_at FROM sientia_data.laborious_data WHERE model_id = {{MODEL_ID}}",
|
||||
"schema": "sientia_data",
|
||||
"table_name": "predictions",
|
||||
"transform_table_name": "transformed_data",
|
||||
"input_filters": {
|
||||
"EMPTY_DATA": {
|
||||
"POLICY": "STOP",
|
||||
"CONFIG": {}
|
||||
}
|
||||
},
|
||||
"mlflow_transform_filters": {
|
||||
"API_ERROR": {
|
||||
"POLICY": "STOP",
|
||||
"CONFIG": {}
|
||||
}
|
||||
},
|
||||
"mlflow_predict_filters": {
|
||||
"API_ERROR": {
|
||||
"POLICY": "STOP",
|
||||
"CONFIG": {}
|
||||
}
|
||||
},
|
||||
"path_priority": ["STOP", "CONTINUE", "REPEAT"],
|
||||
"opc_output_config": {},
|
||||
"pi_web_api_output_config": {},
|
||||
"save_transform": true,
|
||||
"prediction_store_policy": "lts:1",
|
||||
"model_config": {
|
||||
"target": "sensor_1",
|
||||
"retention_minutes": 0
|
||||
},
|
||||
"datetime_columns": ["timestamp", "created_at"]
|
||||
}
|
||||
45
e2e/scenario_inputs/main_invalid_datetime.json
Normal file
45
e2e/scenario_inputs/main_invalid_datetime.json
Normal file
@@ -0,0 +1,45 @@
|
||||
{
|
||||
"metadata": {
|
||||
"metadata": {
|
||||
"model_id": "{{MODEL_ID}}",
|
||||
"model_name": "test_model",
|
||||
"schedule_name": "test-schedule",
|
||||
"workflow_name": "predictions_batch"
|
||||
}
|
||||
},
|
||||
"schedule_name": "test-schedule",
|
||||
"model_name": "test_model",
|
||||
"model_id": "{{MODEL_ID}}",
|
||||
"query": "SELECT timestamp, variable, value FROM sientia_data.laborious_data WHERE model_id = {{MODEL_ID}}",
|
||||
"schema": "sientia_data",
|
||||
"table_name": "predictions",
|
||||
"transform_table_name": "transformed_data",
|
||||
"input_filters": {
|
||||
"EMPTY_DATA": {
|
||||
"POLICY": "STOP",
|
||||
"CONFIG": {}
|
||||
}
|
||||
},
|
||||
"mlflow_transform_filters": {
|
||||
"API_ERROR": {
|
||||
"POLICY": "STOP",
|
||||
"CONFIG": {}
|
||||
}
|
||||
},
|
||||
"mlflow_predict_filters": {
|
||||
"API_ERROR": {
|
||||
"POLICY": "STOP",
|
||||
"CONFIG": {}
|
||||
}
|
||||
},
|
||||
"path_priority": ["STOP", "CONTINUE", "REPEAT"],
|
||||
"opc_output_config": {},
|
||||
"pi_web_api_output_config": {},
|
||||
"save_transform": true,
|
||||
"prediction_store_policy": "lts:1",
|
||||
"model_config": {
|
||||
"target": "sensor_1",
|
||||
"retention_minutes": 0
|
||||
},
|
||||
"datetime_columns": ["nonexistent_column"]
|
||||
}
|
||||
16
e2e/scenario_inputs/main_missing_required.json
Normal file
16
e2e/scenario_inputs/main_missing_required.json
Normal file
@@ -0,0 +1,16 @@
|
||||
{
|
||||
"metadata": {
|
||||
"metadata": {
|
||||
"model_id": "{{MODEL_ID}}",
|
||||
"model_name": "test_model",
|
||||
"schedule_name": "test-schedule",
|
||||
"workflow_name": "predictions_batch"
|
||||
}
|
||||
},
|
||||
"schedule_name": "test-schedule",
|
||||
"model_name": "test_model",
|
||||
"model_id": "{{MODEL_ID}}",
|
||||
"schema": "sientia_data",
|
||||
"table_name": "predictions",
|
||||
"transform_table_name": "transformed_data"
|
||||
}
|
||||
44
e2e/scenario_inputs/main_sql_error.json
Normal file
44
e2e/scenario_inputs/main_sql_error.json
Normal file
@@ -0,0 +1,44 @@
|
||||
{
|
||||
"metadata": {
|
||||
"metadata": {
|
||||
"model_id": "{{MODEL_ID}}",
|
||||
"model_name": "test_model",
|
||||
"schedule_name": "test-schedule",
|
||||
"workflow_name": "predictions_batch"
|
||||
}
|
||||
},
|
||||
"schedule_name": "test-schedule",
|
||||
"model_name": "test_model",
|
||||
"model_id": "{{MODEL_ID}}",
|
||||
"query": "SELECT * FROM nonexistent_table WHERE invalid_syntax =",
|
||||
"schema": "sientia_data",
|
||||
"table_name": "predictions",
|
||||
"transform_table_name": "transformed_data",
|
||||
"input_filters": {
|
||||
"EMPTY_DATA": {
|
||||
"POLICY": "STOP",
|
||||
"CONFIG": {}
|
||||
}
|
||||
},
|
||||
"mlflow_transform_filters": {
|
||||
"API_ERROR": {
|
||||
"POLICY": "STOP",
|
||||
"CONFIG": {}
|
||||
}
|
||||
},
|
||||
"mlflow_predict_filters": {
|
||||
"API_ERROR": {
|
||||
"POLICY": "STOP",
|
||||
"CONFIG": {}
|
||||
}
|
||||
},
|
||||
"path_priority": ["STOP", "CONTINUE", "REPEAT"],
|
||||
"opc_output_config": {},
|
||||
"pi_web_api_output_config": {},
|
||||
"save_transform": true,
|
||||
"prediction_store_policy": "lts:1",
|
||||
"model_config": {
|
||||
"target": "sensor_1",
|
||||
"retention_minutes": 0
|
||||
}
|
||||
}
|
||||
12
e2e/scenario_inputs/minimal_retrain_base.json
Normal file
12
e2e/scenario_inputs/minimal_retrain_base.json
Normal file
@@ -0,0 +1,12 @@
|
||||
{
|
||||
"schedule_name": "test-schedule",
|
||||
"model_name": "test_model",
|
||||
"model_id": "{{MODEL_ID}}",
|
||||
"query": "SELECT timestamp, variable, value, created_at FROM sientia_data.laborious_data WHERE model_id = {{MODEL_ID}}",
|
||||
"schema": "sientia_data",
|
||||
"table_name": "log_retrain",
|
||||
"datetime_columns": ["timestamp", "created_at"],
|
||||
"model_config": {
|
||||
"target": "sensor_1"
|
||||
}
|
||||
}
|
||||
11
e2e/scenario_inputs/minio_offload_load_query.json
Normal file
11
e2e/scenario_inputs/minio_offload_load_query.json
Normal file
@@ -0,0 +1,11 @@
|
||||
{
|
||||
"metadata": {
|
||||
"schedule_name": "test-schedule",
|
||||
"model_name": "test_model",
|
||||
"model_id": "{{MODEL_ID}}",
|
||||
"workflow_name": "predictions_batch"
|
||||
},
|
||||
"query": "SELECT timestamp, variable, value, created_at FROM sientia_data.laborious_data WHERE model_id = {{MODEL_ID}}",
|
||||
"model_name": "test_model",
|
||||
"datetime_columns": ["timestamp", "created_at"]
|
||||
}
|
||||
37
e2e/scenario_inputs/minio_offload_workflow.json
Normal file
37
e2e/scenario_inputs/minio_offload_workflow.json
Normal file
@@ -0,0 +1,37 @@
|
||||
{
|
||||
"schedule_name": "test-schedule",
|
||||
"model_name": "test_model",
|
||||
"model_id": "{{MODEL_ID}}",
|
||||
"query": "SELECT timestamp, variable, value, created_at FROM sientia_data.laborious_data WHERE model_id = {{MODEL_ID}}",
|
||||
"schema": "sientia_data",
|
||||
"table_name": "predictions",
|
||||
"transform_table_name": "transformed_data",
|
||||
"input_filters": {
|
||||
"EMPTY_DATA": {
|
||||
"POLICY": "STOP",
|
||||
"CONFIG": {}
|
||||
}
|
||||
},
|
||||
"mlflow_transform_filters": {
|
||||
"API_ERROR": {
|
||||
"POLICY": "STOP",
|
||||
"CONFIG": {}
|
||||
}
|
||||
},
|
||||
"mlflow_predict_filters": {
|
||||
"API_ERROR": {
|
||||
"POLICY": "STOP",
|
||||
"CONFIG": {}
|
||||
}
|
||||
},
|
||||
"path_priority": ["STOP", "CONTINUE", "REPEAT"],
|
||||
"opc_output_config": {},
|
||||
"pi_web_api_output_config": {},
|
||||
"save_transform": false,
|
||||
"prediction_store_policy": "lts:1",
|
||||
"model_config": {
|
||||
"retention_minutes": 0,
|
||||
"target": "sensor_1"
|
||||
},
|
||||
"datetime_columns": ["timestamp", "created_at"]
|
||||
}
|
||||
39
e2e/scenario_inputs/prediction_process_base.json
Normal file
39
e2e/scenario_inputs/prediction_process_base.json
Normal file
@@ -0,0 +1,39 @@
|
||||
{
|
||||
"schedule_name": "test-schedule",
|
||||
"model_name": "test_model",
|
||||
"model_id": "{{MODEL_ID}}",
|
||||
"query": "SELECT timestamp, variable, value, created_at FROM sientia_data.laborious_data WHERE model_id = {{MODEL_ID}}",
|
||||
"schema": "sientia_data",
|
||||
"table_name": "predictions",
|
||||
"transform_table_name": "transformed_data",
|
||||
"input_filters": {
|
||||
"SPECIFIC_VARIABLES_NULL_VALUES": {
|
||||
"POLICY": "CONTINUE",
|
||||
"CONFIG": {
|
||||
"variables": ["sensor_1"]
|
||||
}
|
||||
}
|
||||
},
|
||||
"mlflow_transform_filters": {
|
||||
"API_ERROR": {
|
||||
"POLICY": "STOP",
|
||||
"CONFIG": {}
|
||||
}
|
||||
},
|
||||
"mlflow_predict_filters": {
|
||||
"API_ERROR": {
|
||||
"POLICY": "STOP",
|
||||
"CONFIG": {}
|
||||
}
|
||||
},
|
||||
"path_priority": ["STOP", "CONTINUE", "REPEAT"],
|
||||
"opc_output_config": {},
|
||||
"pi_web_api_output_config": {},
|
||||
"save_transform": true,
|
||||
"prediction_store_policy": "lts:1",
|
||||
"model_config": {
|
||||
"retention_minutes": 0,
|
||||
"target": "sensor_1"
|
||||
},
|
||||
"datetime_columns": ["timestamp", "created_at"]
|
||||
}
|
||||
14
e2e/scenario_inputs/simple_metrics_base.json
Normal file
14
e2e/scenario_inputs/simple_metrics_base.json
Normal file
@@ -0,0 +1,14 @@
|
||||
{
|
||||
"schedule_name": "test-schedule",
|
||||
"model_name": "test_model",
|
||||
"model_id": "{{MODEL_ID}}",
|
||||
"schema": "sientia_data",
|
||||
"predictions_table_name": "predictions",
|
||||
"data_table_name": "laborious_data",
|
||||
"target_table_name": "simple_metrics",
|
||||
"interval_minutes": 60,
|
||||
"metrics": ["rmse", "mse", "mae", "r2"],
|
||||
"model_config": {
|
||||
"target": "sensor_target"
|
||||
}
|
||||
}
|
||||
1375
e2e/scenarios.md
1375
e2e/scenarios.md
File diff suppressed because it is too large
Load Diff
@@ -26,9 +26,7 @@ async def test_format_and_export_prediction_default_path_e2e(
|
||||
client = temporal_test_env.client
|
||||
model_id = 401
|
||||
with postgres_engine.begin() as conn:
|
||||
conn.execute(
|
||||
text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')
|
||||
)
|
||||
conn.execute(text(f'DELETE FROM sientia_data.predictions WHERE model_id = {model_id}'))
|
||||
|
||||
metadata = {
|
||||
'metadata': {
|
||||
@@ -46,7 +44,7 @@ async def test_format_and_export_prediction_default_path_e2e(
|
||||
'timestamp': '2024-01-01 12:00:00+00:00',
|
||||
'model_id': model_id,
|
||||
'model_name': 'test_model',
|
||||
'schema': 'predictions_schema',
|
||||
'schema': 'sientia_data',
|
||||
'table_name': 'predictions',
|
||||
'transform_table_name': 'transformed_data',
|
||||
'comment': 'e2e child workflow default path',
|
||||
@@ -66,7 +64,7 @@ async def test_format_and_export_prediction_default_path_e2e(
|
||||
row = conn.execute(
|
||||
text(
|
||||
f'SELECT prediction, prediction_confidence, prediction_status, comments '
|
||||
f'FROM predictions_schema.predictions WHERE model_id = {model_id}'
|
||||
f'FROM sientia_data.predictions WHERE model_id = {model_id}'
|
||||
)
|
||||
).fetchone()
|
||||
assert row is not None
|
||||
|
||||
601
e2e/test_drift.py
Normal file
601
e2e/test_drift.py
Normal file
@@ -0,0 +1,601 @@
|
||||
"""
|
||||
End-to-end tests for the Drift workflow.
|
||||
|
||||
The drift suite drives the **real** ``sientia_model.analytics.drift_analysis.DriftAnalysis``
|
||||
analyzer (no mocking). Each scenario exercises the full pipeline:
|
||||
|
||||
laborious_data (Postgres)
|
||||
-> load_custom_query
|
||||
-> calculate_drift (DriftAnalysis univariate + multivariate)
|
||||
-> export_data_to_postgres (sientia_data.drift_metrics)
|
||||
|
||||
Coverage focus:
|
||||
|
||||
- Happy path persists every column required by ``sientia_data.drift_metrics``
|
||||
with a valid reference dataset downloaded from MLflow.
|
||||
- 30% fallback path activates when the MLflow reference is unavailable and
|
||||
emits the ``MODEL_METRICS_REFERENCE_DATA_WARNING`` notification.
|
||||
- Empty target data short-circuits the workflow without persisting anything.
|
||||
- Invalid ``chunk_period`` is rejected by ``calculate_drift``.
|
||||
- ``chunk_period='s'`` preserves second-level precision in
|
||||
``chunk_start_date``.
|
||||
|
||||
Tests assert behavioral / structural properties (column presence, NOT NULL
|
||||
constraints, business-key invariants) rather than exact numeric values, since
|
||||
those depend on the real analyzer implementation and synthetic data.
|
||||
"""
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pandas as pd
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
from temporalio.testing import WorkflowEnvironment
|
||||
from temporalio.worker import Worker
|
||||
|
||||
from e2e.helpers import (
|
||||
insert_target_data_for_drift,
|
||||
load_scenario_input,
|
||||
make_workflow_id,
|
||||
start_and_await_workflow,
|
||||
)
|
||||
from laborious.activities.activities import Activities
|
||||
from laborious.activities.model_metrics import ModelMetrics
|
||||
from laborious.workflows.drift import Drift
|
||||
from sientia_model.analytics.drift_analysis import DriftAnalysis
|
||||
|
||||
# Drift columns persisted on every row in ``sientia_data.drift_metrics`` —
|
||||
# mirrors the production DDL.
|
||||
EXPECTED_DRIFT_COLUMNS = [
|
||||
'id',
|
||||
'model_id',
|
||||
'feature',
|
||||
'method',
|
||||
'value',
|
||||
'alert',
|
||||
'chunk_index',
|
||||
'chunk_start_date',
|
||||
'chunk_end_date',
|
||||
'accurate',
|
||||
'timestamp',
|
||||
'created_at',
|
||||
]
|
||||
|
||||
# Columns the DDL marks as NOT NULL. ``feature`` and ``timestamp`` are
|
||||
# nullable in the production schema (multivariate rows do not bind to a
|
||||
# single feature; ``timestamp`` is allowed to be empty when upstream data has
|
||||
# no usable instant).
|
||||
NON_NULL_DRIFT_COLUMNS = {
|
||||
'id',
|
||||
'model_id',
|
||||
'method',
|
||||
'value',
|
||||
'alert',
|
||||
'chunk_index',
|
||||
'chunk_start_date',
|
||||
'chunk_end_date',
|
||||
'accurate',
|
||||
'created_at',
|
||||
}
|
||||
|
||||
DEFAULT_DRIFT_METHODS = ['kolmogorov_smirnov', 'jensen_shannon', 'wasserstein']
|
||||
|
||||
|
||||
def _chunk_dataframe_skip_empty_groups(
|
||||
self: DriftAnalysis,
|
||||
df: pd.DataFrame,
|
||||
timestamp_col: str,
|
||||
chunk_period: str,
|
||||
) -> list[tuple[int, pd.DataFrame]]:
|
||||
"""
|
||||
Same as ``DriftAnalysis._chunk_dataframe`` but omit empty time buckets.
|
||||
|
||||
``pd.Grouper(freq='s')`` yields every second between min and max timestamp;
|
||||
empty buckets still appear in the groupby iterator and produce invalid
|
||||
drift rows (e.g. NaT timestamps) that ``calculate_drift`` later filters out
|
||||
entirely. Production fix belongs in ``sientia_model``; this shim keeps the
|
||||
e2e honest about second-level chunk boundaries with sparse samples.
|
||||
"""
|
||||
grouped = df.groupby(pd.Grouper(key=timestamp_col, freq=chunk_period), dropna=True)
|
||||
chunks: list[tuple[int, pd.DataFrame]] = []
|
||||
idx = 0
|
||||
for _, chunk in grouped:
|
||||
if chunk.empty:
|
||||
continue
|
||||
chunks.append((idx, chunk.copy()))
|
||||
idx += 1
|
||||
return chunks
|
||||
|
||||
|
||||
def _drift_input(model_id: int, **overrides) -> dict:
|
||||
"""Load the base drift scenario JSON and apply ad-hoc overrides."""
|
||||
input_data = load_scenario_input('drift_base.json', model_id=model_id)
|
||||
input_data.update(overrides)
|
||||
return input_data
|
||||
|
||||
|
||||
def _recent_minute_timestamps(count: int, offset_minutes: int = 6) -> list[str]:
|
||||
"""
|
||||
Build ``count`` consecutive UTC minute timestamps placed in the recent past.
|
||||
|
||||
The Drift workflow filters target rows with ``timestamp > NOW() - INTERVAL``,
|
||||
so timestamps must be recent for tests to retrieve any data. Snapping to
|
||||
minute precision keeps the helper deterministic regardless of clock skew.
|
||||
|
||||
Args:
|
||||
- count (int): How many consecutive minute timestamps to generate.
|
||||
- offset_minutes (int): Minutes ago for the EARLIEST generated timestamp.
|
||||
|
||||
Return:
|
||||
list[str]: ISO strings with ``+0000`` offset, one per minute.
|
||||
"""
|
||||
base = datetime.now(timezone.utc).replace(second=0, microsecond=0) - timedelta(
|
||||
minutes=offset_minutes
|
||||
)
|
||||
return [
|
||||
(base + timedelta(minutes=i)).strftime('%Y-%m-%d %H:%M:%S%z') for i in range(count)
|
||||
]
|
||||
|
||||
|
||||
def _configure_reference_csv(mlflow_repository_stub, reference_rows: pd.DataFrame) -> None:
|
||||
"""
|
||||
Wire ``mlflow_repository_stub`` so ``get_reference_data`` returns
|
||||
``reference_rows`` by writing them to ``dst_path/evaluation_data.csv``.
|
||||
|
||||
Args:
|
||||
- mlflow_repository_stub: External MLflow repository fixture.
|
||||
- reference_rows (pd.DataFrame): Rows to expose as the production reference.
|
||||
"""
|
||||
|
||||
def _download(run_id: str, artifact_path: str, dst_path: str, metadata=None):
|
||||
target = Path(dst_path) / artifact_path
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
reference_rows.to_csv(target, index=False)
|
||||
|
||||
mlflow_repository_stub._client.get_model_version_by_alias.return_value = MagicMock(
|
||||
run_id='fake-reference-run'
|
||||
)
|
||||
file_info = MagicMock()
|
||||
# Matches `MLFlow._REFERENCE_ARTIFACT_CANDIDATES` (SIENTIAPDE-1646).
|
||||
file_info.path = 'evaluation_data.csv'
|
||||
mlflow_repository_stub._client.list_artifacts.return_value = [file_info]
|
||||
mlflow_repository_stub.download_artifacts.side_effect = _download
|
||||
|
||||
|
||||
def _force_reference_unavailable(mlflow_repository_stub) -> None:
|
||||
"""Make ``get_reference_data`` return ``None`` by failing alias resolution."""
|
||||
mlflow_repository_stub._client.get_model_version_by_alias.side_effect = Exception(
|
||||
'no production alias registered'
|
||||
)
|
||||
|
||||
|
||||
def _select_drift_rows(postgres_engine, model_id: int) -> list[dict]:
|
||||
"""Read every persisted drift row for ``model_id`` ordered by chunk/feature/method."""
|
||||
with postgres_engine.connect() as conn:
|
||||
rows = (
|
||||
conn.execute(
|
||||
text(
|
||||
'SELECT * FROM sientia_data.drift_metrics '
|
||||
'WHERE model_id = :m '
|
||||
'ORDER BY chunk_index, feature, method'
|
||||
),
|
||||
{'m': str(model_id)},
|
||||
)
|
||||
.mappings()
|
||||
.all()
|
||||
)
|
||||
return [dict(row) for row in rows]
|
||||
|
||||
|
||||
def _assert_required_columns_populated(rows: list[dict]) -> None:
|
||||
"""Validate column presence and NOT NULL constraints on every row."""
|
||||
assert rows, 'expected at least one drift row to be persisted'
|
||||
seen_columns = set(rows[0].keys())
|
||||
for column in EXPECTED_DRIFT_COLUMNS:
|
||||
assert column in seen_columns, f'Missing drift column in postgres: {column}'
|
||||
for row in rows:
|
||||
for column in NON_NULL_DRIFT_COLUMNS:
|
||||
assert row[column] is not None, f"Column '{column}' is NULL in {row}"
|
||||
assert 'p_value' not in row, 'p_value must not be persisted to drift_metrics'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_drift_happy_path_persists_all_columns_with_reference_data(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker_drift: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
mlflow_repository_stub,
|
||||
):
|
||||
"""
|
||||
Scenario D.1.1: Happy path with reference data downloaded from MLflow.
|
||||
|
||||
Drives the full pipeline against the real ``DriftAnalysis``. Asserts:
|
||||
|
||||
- One row is persisted per ``(chunk_index, feature, method)`` combination
|
||||
plus the multivariate row block, with every column required by
|
||||
``sientia_data.drift_metrics`` populated.
|
||||
- The three default univariate methods are forwarded to the analyzer.
|
||||
- ``model_id`` and ``timestamp`` are stamped by the activity (not by the
|
||||
analyzer); ``timestamp`` equals ``max(target_data.timestamp)`` and is
|
||||
identical on every persisted row.
|
||||
- ``chunk_start_date`` / ``chunk_end_date`` are persisted as ISO text so
|
||||
the analyzer's nanosecond-precision boundaries survive the ``text``
|
||||
column type.
|
||||
- ``accurate=True`` because the reference dataset was available.
|
||||
"""
|
||||
client = temporal_test_env.client
|
||||
model_id = 411
|
||||
|
||||
target_timestamps = _recent_minute_timestamps(count=10)
|
||||
insert_target_data_for_drift(
|
||||
postgres_engine,
|
||||
model_id=model_id,
|
||||
timestamps=target_timestamps,
|
||||
variables_values={
|
||||
'sensor_1': [10.0 + i * 0.1 for i in range(10)],
|
||||
'sensor_2': [20.0 + i * 0.5 for i in range(10)],
|
||||
},
|
||||
)
|
||||
|
||||
reference_df = pd.DataFrame(
|
||||
{
|
||||
'timestamp': [
|
||||
f'2023-12-31 11:{minute:02d}:00+00:00' for minute in range(10)
|
||||
],
|
||||
'sensor_1': [9.0 + i * 0.05 for i in range(10)],
|
||||
'sensor_2': [18.0 + i * 0.25 for i in range(10)],
|
||||
}
|
||||
)
|
||||
_configure_reference_csv(mlflow_repository_stub, reference_df)
|
||||
|
||||
input_data = _drift_input(model_id)
|
||||
await start_and_await_workflow(
|
||||
client, Drift.run, input_data, make_workflow_id('test-drift-happy-path')
|
||||
)
|
||||
|
||||
rows = _select_drift_rows(postgres_engine, model_id)
|
||||
|
||||
exported_csv_path = '/tmp/test_drift_happy_path_exported.csv'
|
||||
pd.DataFrame(rows).to_csv(exported_csv_path, index=False)
|
||||
print(
|
||||
f'\n[test_drift_happy_path] Exported drift dataframe '
|
||||
f'({len(rows)} rows) -> {exported_csv_path}'
|
||||
)
|
||||
|
||||
_assert_required_columns_populated(rows)
|
||||
|
||||
# The activity drops the target column from the feature list, so only
|
||||
# ``sensor_2`` participates in univariate analysis (``sensor_1`` is the
|
||||
# configured target). Multivariate produces one row per chunk regardless.
|
||||
univariate_rows = [r for r in rows if r['feature'] != 'multivariate']
|
||||
multivariate_rows = [r for r in rows if r['feature'] == 'multivariate']
|
||||
assert univariate_rows, 'expected univariate drift rows for non-target features'
|
||||
assert multivariate_rows, 'expected one multivariate drift row per chunk'
|
||||
|
||||
# All three default methods must reach the analyzer.
|
||||
assert {r['method'] for r in univariate_rows} == set(DEFAULT_DRIFT_METHODS)
|
||||
assert all(r['method'] == 'multivariate' for r in multivariate_rows)
|
||||
assert {r['feature'] for r in univariate_rows} == {'sensor_2'}
|
||||
|
||||
# ``timestamp`` is stamped uniformly with ``max(target_data.timestamp)``.
|
||||
expected_timestamp = pd.to_datetime(max(target_timestamps), utc=True)
|
||||
persisted_timestamps = {pd.to_datetime(r['timestamp'], utc=True) for r in rows}
|
||||
assert len(persisted_timestamps) == 1, (
|
||||
'timestamp must be uniform across all drift rows '
|
||||
f'(got {len(persisted_timestamps)} distinct values)'
|
||||
)
|
||||
assert pd.Timestamp(persisted_timestamps.pop()) == expected_timestamp, (
|
||||
'timestamp must equal max(target_data.timestamp)'
|
||||
)
|
||||
|
||||
# ``model_id`` is stamped by ``calculate_drift`` (not produced by the analyzer).
|
||||
assert all(r['model_id'] == str(model_id) for r in rows), (
|
||||
'model_id must be stamped on every drift row'
|
||||
)
|
||||
|
||||
# Reference path → accurate=True.
|
||||
assert all(r['accurate'] is True for r in rows), (
|
||||
'reference path should mark all rows as accurate'
|
||||
)
|
||||
|
||||
# ISO text serialization preserves ordering between start/end of each chunk.
|
||||
for row in rows:
|
||||
assert 'T' in row['chunk_start_date'], (
|
||||
f"chunk_start_date should be ISO text, got {row['chunk_start_date']!r}"
|
||||
)
|
||||
assert 'T' in row['chunk_end_date'], (
|
||||
f"chunk_end_date should be ISO text, got {row['chunk_end_date']!r}"
|
||||
)
|
||||
assert row['chunk_start_date'] <= row['chunk_end_date'], (
|
||||
f'chunk_start_date must precede chunk_end_date '
|
||||
f"(start={row['chunk_start_date']}, end={row['chunk_end_date']})"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_drift_uses_30pct_fallback_when_reference_unavailable(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker_drift: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
mlflow_repository_stub,
|
||||
notification_inserts,
|
||||
):
|
||||
"""
|
||||
Scenario D.1.2: ``get_reference_data`` returns ``None`` (production alias
|
||||
missing), so ``calculate_drift`` falls back to using the first 30% of
|
||||
target rows as reference. Persisted rows must report ``accurate=False``
|
||||
and a ``MODEL_METRICS_REFERENCE_DATA_WARNING`` notification must be
|
||||
emitted to mongo.
|
||||
"""
|
||||
client = temporal_test_env.client
|
||||
model_id = 412
|
||||
|
||||
target_timestamps = _recent_minute_timestamps(count=10)
|
||||
insert_target_data_for_drift(
|
||||
postgres_engine,
|
||||
model_id=model_id,
|
||||
timestamps=target_timestamps,
|
||||
variables_values={
|
||||
'sensor_1': [10.0 + i * 0.1 for i in range(10)],
|
||||
'sensor_2': [20.0 + i * 0.5 for i in range(10)],
|
||||
},
|
||||
)
|
||||
|
||||
_force_reference_unavailable(mlflow_repository_stub)
|
||||
|
||||
input_data = _drift_input(model_id)
|
||||
await start_and_await_workflow(
|
||||
client, Drift.run, input_data, make_workflow_id('test-drift-fallback')
|
||||
)
|
||||
|
||||
rows = _select_drift_rows(postgres_engine, model_id)
|
||||
_assert_required_columns_populated(rows)
|
||||
|
||||
assert all(r['accurate'] is False for r in rows), (
|
||||
'fallback path must mark all rows as inaccurate'
|
||||
)
|
||||
|
||||
fallback_warnings = [
|
||||
call
|
||||
for call in notification_inserts.call_args_list
|
||||
if call.args
|
||||
and isinstance(call.args[0], dict)
|
||||
and call.args[0].get('notification_id') == 'MODEL_METRICS_REFERENCE_DATA_WARNING'
|
||||
]
|
||||
assert len(fallback_warnings) >= 1, 'expected reference fallback warning notification'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_drift_empty_target_data_short_circuits_workflow(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker_drift: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
mlflow_repository_stub,
|
||||
):
|
||||
"""
|
||||
Scenario D.3.1: When ``load_custom_query`` returns no rows the workflow
|
||||
must return early without invoking the analyzer or writing any drift rows.
|
||||
"""
|
||||
client = temporal_test_env.client
|
||||
model_id = 431
|
||||
|
||||
with postgres_engine.begin() as conn:
|
||||
conn.execute(text(f'DELETE FROM sientia_data.laborious_data WHERE model_id = {model_id}'))
|
||||
|
||||
_force_reference_unavailable(mlflow_repository_stub)
|
||||
|
||||
input_data = _drift_input(model_id)
|
||||
await start_and_await_workflow(
|
||||
client, Drift.run, input_data, make_workflow_id('test-drift-empty-target')
|
||||
)
|
||||
|
||||
with postgres_engine.connect() as conn:
|
||||
count = conn.execute(
|
||||
text('SELECT COUNT(*) FROM sientia_data.drift_metrics WHERE model_id = :m'),
|
||||
{'m': str(model_id)},
|
||||
).scalar()
|
||||
assert count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_drift_invalid_chunk_period_raises_value_error(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker_drift: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
mlflow_repository_stub,
|
||||
):
|
||||
"""
|
||||
Scenario D.4.2: ``calculate_drift`` validates ``chunk_period`` and rejects
|
||||
anything other than ``min`` / ``s``. The workflow must surface the
|
||||
``ValueError`` and persist nothing.
|
||||
"""
|
||||
client = temporal_test_env.client
|
||||
model_id = 442
|
||||
|
||||
target_timestamps = _recent_minute_timestamps(count=5)
|
||||
insert_target_data_for_drift(
|
||||
postgres_engine,
|
||||
model_id=model_id,
|
||||
timestamps=target_timestamps,
|
||||
variables_values={
|
||||
'sensor_1': [1.0, 2.0, 3.0, 4.0, 5.0],
|
||||
'sensor_2': [10.0, 20.0, 30.0, 40.0, 50.0],
|
||||
},
|
||||
)
|
||||
_force_reference_unavailable(mlflow_repository_stub)
|
||||
|
||||
input_data = _drift_input(model_id, chunk_period='hour')
|
||||
|
||||
with pytest.raises(Exception) as excinfo:
|
||||
await start_and_await_workflow(
|
||||
client,
|
||||
Drift.run,
|
||||
input_data,
|
||||
make_workflow_id('test-drift-bad-chunk-period'),
|
||||
)
|
||||
|
||||
# Temporal wraps the activity ValueError in WorkflowFailureError; the
|
||||
# message may live on ``.message`` or ``str(exc)`` depending on the SDK
|
||||
# error class, so walk the cause chain looking for the guard text.
|
||||
cause_descriptions = []
|
||||
current: BaseException | None = excinfo.value
|
||||
while current is not None:
|
||||
cause_descriptions.append(
|
||||
getattr(current, 'message', None) or str(current) or repr(current)
|
||||
)
|
||||
current = current.__cause__
|
||||
assert any('Invalid chunk period' in msg for msg in cause_descriptions), (
|
||||
f'Expected ValueError about chunk period in chain, got: {cause_descriptions}'
|
||||
)
|
||||
|
||||
with postgres_engine.connect() as conn:
|
||||
count = conn.execute(
|
||||
text('SELECT COUNT(*) FROM sientia_data.drift_metrics WHERE model_id = :m'),
|
||||
{'m': str(model_id)},
|
||||
).scalar()
|
||||
assert count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_drift_empty_merge_skips_export_without_insufficient_notification(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker_drift: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
mlflow_repository_stub,
|
||||
):
|
||||
"""
|
||||
Scenario D.4.3a: When the analyzer returns an empty merged frame (no metric rows),
|
||||
``calculate_drift`` yields ``[]``; the workflow skips export. Real insufficient-data
|
||||
cases are signaled by ``DriftInsufficientDataError`` inside ``sientia_model``, not by
|
||||
empty output alone.
|
||||
"""
|
||||
client = temporal_test_env.client
|
||||
model_id = 444
|
||||
|
||||
target_timestamps = _recent_minute_timestamps(count=5)
|
||||
insert_target_data_for_drift(
|
||||
postgres_engine,
|
||||
model_id=model_id,
|
||||
timestamps=target_timestamps,
|
||||
variables_values={
|
||||
'sensor_1': [1.0 + i * 0.1 for i in range(5)],
|
||||
'sensor_2': [10.0 + i * 0.5 for i in range(5)],
|
||||
},
|
||||
)
|
||||
_force_reference_unavailable(mlflow_repository_stub)
|
||||
|
||||
input_data = _drift_input(model_id, chunk_period='min')
|
||||
empty_merge = pd.DataFrame(
|
||||
columns=[
|
||||
'timestamp',
|
||||
'feature',
|
||||
'method',
|
||||
'value',
|
||||
'alert',
|
||||
'chunk_index',
|
||||
'chunk_start_date',
|
||||
'chunk_end_date',
|
||||
'threshold',
|
||||
'drift_type',
|
||||
]
|
||||
)
|
||||
with patch.object(ModelMetrics, 'get_drift_metrics', return_value=empty_merge):
|
||||
await start_and_await_workflow(
|
||||
client,
|
||||
Drift.run,
|
||||
input_data,
|
||||
make_workflow_id('test-drift-empty-merge'),
|
||||
)
|
||||
|
||||
with postgres_engine.connect() as conn:
|
||||
count = conn.execute(
|
||||
text('SELECT COUNT(*) FROM sientia_data.drift_metrics WHERE model_id = :m'),
|
||||
{'m': str(model_id)},
|
||||
).scalar()
|
||||
assert count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_drift_chunk_period_seconds_sufficient_data_preserves_seconds_in_chunk_start_date(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker_drift: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
mlflow_repository_stub,
|
||||
):
|
||||
"""
|
||||
Scenario D.4.3b: With enough sub-minute samples and ``chunk_period='s'``, drift rows
|
||||
persist and ``chunk_start_date`` keeps second-level precision (incl. second=30).
|
||||
|
||||
``DriftAnalysis._chunk_dataframe`` is patched to skip empty ``pd.Grouper(freq='s')``
|
||||
buckets so sparse seconds between samples do not flood the pipeline with NaT rows;
|
||||
the durable fix belongs in ``sientia_model``.
|
||||
"""
|
||||
client = temporal_test_env.client
|
||||
model_id = 443
|
||||
|
||||
base = datetime.now(timezone.utc).replace(second=0, microsecond=0) - timedelta(minutes=6)
|
||||
target_timestamps = []
|
||||
sensor_1_vals = []
|
||||
sensor_2_vals = []
|
||||
for minute_offset in range(6):
|
||||
t0 = base + timedelta(minutes=minute_offset)
|
||||
t1 = t0 + timedelta(seconds=30)
|
||||
target_timestamps.append(t0.strftime('%Y-%m-%d %H:%M:%S%z'))
|
||||
target_timestamps.append(t1.strftime('%Y-%m-%d %H:%M:%S%z'))
|
||||
v0 = 1.0 + minute_offset * 0.1
|
||||
v1 = v0 + 0.05
|
||||
sensor_1_vals.extend([v0, v1])
|
||||
sensor_2_vals.extend([10.0 + v0, 10.0 + v1])
|
||||
|
||||
insert_target_data_for_drift(
|
||||
postgres_engine,
|
||||
model_id=model_id,
|
||||
timestamps=target_timestamps,
|
||||
variables_values={
|
||||
'sensor_1': sensor_1_vals,
|
||||
'sensor_2': sensor_2_vals,
|
||||
},
|
||||
)
|
||||
_force_reference_unavailable(mlflow_repository_stub)
|
||||
|
||||
input_data = _drift_input(model_id, chunk_period='s')
|
||||
with patch.object(DriftAnalysis, '_chunk_dataframe', _chunk_dataframe_skip_empty_groups):
|
||||
await start_and_await_workflow(
|
||||
client,
|
||||
Drift.run,
|
||||
input_data,
|
||||
make_workflow_id('test-drift-chunk-seconds-sufficient'),
|
||||
)
|
||||
|
||||
with postgres_engine.connect() as conn:
|
||||
rows = (
|
||||
conn.execute(
|
||||
text(
|
||||
'SELECT chunk_start_date FROM sientia_data.drift_metrics '
|
||||
'WHERE model_id = :m ORDER BY chunk_index'
|
||||
),
|
||||
{'m': str(model_id)},
|
||||
)
|
||||
.mappings()
|
||||
.all()
|
||||
)
|
||||
|
||||
assert rows, 'expected at least one drift row to be persisted'
|
||||
seconds_present = {pd.Timestamp(r['chunk_start_date']).second for r in rows}
|
||||
assert 30 in seconds_present, (
|
||||
f'expected at least one chunk_start_date with seconds=30, got {seconds_present}'
|
||||
)
|
||||
@@ -1,196 +0,0 @@
|
||||
"""Does the harness itself come up? Diagnosable separately from the suite that runs on it.
|
||||
|
||||
Four containers, an `mlflow server` subprocess and a Temporal worker is a lot of moving parts, and
|
||||
when one of them fails to start the first *real* test fails somewhere deep inside an activity with
|
||||
a message about something else. These tests exist so that a broken harness says so in its own
|
||||
words: each one checks one piece is alive and does nothing with the import path.
|
||||
|
||||
The last test is the exception — it runs a real workflow, but with an input designed to be rejected
|
||||
at the very first step, so what it proves is only that the wiring reaches the first activity.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import mlflow
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
from temporalio.client import WorkflowFailureError
|
||||
|
||||
from e2e.import_harness import (
|
||||
BFF_DATABASE,
|
||||
IMPORT_ACTIVITY_NAMES,
|
||||
IMPORT_BUCKET,
|
||||
IMPORT_PREFIX,
|
||||
IMPORT_TASK_QUEUE,
|
||||
SEEDED_MODEL_IDS,
|
||||
failure_text,
|
||||
failure_types,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.import_e2e, pytest.mark.asyncio]
|
||||
|
||||
|
||||
class TestContainers:
|
||||
"""Each dependency is real, reachable, and the one the harness says it is."""
|
||||
|
||||
async def test_bff_container_holds_the_status_table(self, import_status_table: Any) -> None:
|
||||
with import_status_table.begin() as connection:
|
||||
database = connection.execute(text('SELECT current_database()')).scalar_one()
|
||||
columns = {
|
||||
row[0]
|
||||
for row in connection.execute(
|
||||
text(
|
||||
'SELECT column_name FROM information_schema.columns '
|
||||
"WHERE table_schema = 'public' AND table_name = 'experiment_run'"
|
||||
)
|
||||
)
|
||||
}
|
||||
assert database == BFF_DATABASE
|
||||
assert {'id', 'status', 'run_type', 'request_data', 'orchestrator_response_data'} <= columns
|
||||
|
||||
async def test_bff_container_logs_every_statement(self, import_bff_container: Any) -> None:
|
||||
"""`log_statement=all` — without it the absence assertions have no vantage point."""
|
||||
result = import_bff_container.exec(
|
||||
['psql', '-U', 'bff', '-d', BFF_DATABASE, '-tAc', 'SHOW log_statement']
|
||||
)
|
||||
assert b'all' in result.output
|
||||
|
||||
async def test_the_v12_delta_is_what_separates_the_two_tables(
|
||||
self, import_status_table_pre_v12: Any
|
||||
) -> None:
|
||||
"""The pre-`V12` table has `run_type` and refuses the import terminal statuses."""
|
||||
with import_status_table_pre_v12.begin() as connection:
|
||||
constraints = {
|
||||
row[0]
|
||||
for row in connection.execute(
|
||||
text(
|
||||
'SELECT conname FROM pg_constraint '
|
||||
"WHERE conrelid = 'public.experiment_run'::regclass"
|
||||
)
|
||||
)
|
||||
}
|
||||
columns = {
|
||||
row[0]
|
||||
for row in connection.execute(
|
||||
text(
|
||||
'SELECT column_name FROM information_schema.columns '
|
||||
"WHERE table_schema = 'public' AND table_name = 'experiment_run'"
|
||||
)
|
||||
)
|
||||
}
|
||||
assert 'run_type' in columns
|
||||
assert 'chk_experiment_run_terminal_status_run_type' not in constraints
|
||||
|
||||
async def test_the_seeded_collection_has_gapped_ids_and_a_unique_index(
|
||||
self, import_models_collection: Any
|
||||
) -> None:
|
||||
ids = {
|
||||
str(document['id']) for document in import_models_collection.find({}, {'_id': False})
|
||||
}
|
||||
assert ids == set(SEEDED_MODEL_IDS)
|
||||
unique_keys = {
|
||||
tuple(index['key'].keys())
|
||||
for index in import_models_collection.list_indexes()
|
||||
if index.get('unique')
|
||||
}
|
||||
assert ('id',) in unique_keys
|
||||
|
||||
async def test_the_import_bucket_exists_and_its_prefix_is_empty(
|
||||
self, import_bucket: Any
|
||||
) -> None:
|
||||
assert import_bucket.bucket_exists(IMPORT_BUCKET)
|
||||
assert not list(
|
||||
import_bucket.list_objects(IMPORT_BUCKET, prefix=IMPORT_PREFIX, recursive=True)
|
||||
)
|
||||
|
||||
async def test_uploading_a_bundle_returns_a_key_and_its_digest(
|
||||
self, import_upload_bundle: Any, import_bucket: Any
|
||||
) -> None:
|
||||
object_key, digest = import_upload_bundle(b'not a real bundle, just bytes')
|
||||
assert object_key.startswith(IMPORT_PREFIX)
|
||||
assert len(digest) == 64
|
||||
assert import_bucket.stat_object(IMPORT_BUCKET, object_key).size > 0
|
||||
|
||||
async def test_the_mlflow_server_answers_and_serves_a_registry(
|
||||
self, import_mlflow_server: dict[str, Any]
|
||||
) -> None:
|
||||
client = mlflow.tracking.MlflowClient(tracking_uri=import_mlflow_server['tracking_uri'])
|
||||
experiment_id = client.create_experiment('harness_smoke_experiment')
|
||||
run = client.create_run(experiment_id, run_name='harness-smoke-run')
|
||||
client.create_registered_model('harness_smoke_model')
|
||||
version = client.create_model_version(
|
||||
'harness_smoke_model', f'runs:/{run.info.run_id}/prediction_model', run.info.run_id
|
||||
)
|
||||
assert version.version == '1'
|
||||
assert version.current_stage in (None, 'None')
|
||||
|
||||
|
||||
class TestActivitiesAndWorker:
|
||||
"""The object under test is built from the real dependencies, and the worker registers it."""
|
||||
|
||||
async def test_activities_point_at_the_real_dependencies(
|
||||
self, import_activities: Any, import_mlflow_server: dict[str, Any]
|
||||
) -> None:
|
||||
assert mlflow.get_tracking_uri() == import_mlflow_server['tracking_uri']
|
||||
assert import_activities.import_config['bucket'] == IMPORT_BUCKET
|
||||
assert import_activities.import_config['prefix'] == IMPORT_PREFIX
|
||||
assert import_activities.status_config['dbname'] == BFF_DATABASE
|
||||
assert import_activities.mongo_repository is not None
|
||||
# Not a mock: a MagicMock would satisfy any attribute access, so the type is the assertion.
|
||||
assert type(import_activities.import_minio_repository).__name__ == 'MinioRepository'
|
||||
|
||||
async def test_every_import_activity_is_registered(self, import_activities: Any) -> None:
|
||||
for name in IMPORT_ACTIVITY_NAMES:
|
||||
assert callable(getattr(import_activities, name)), f'{name} is missing'
|
||||
|
||||
async def test_the_worker_polls_the_import_queue(self, import_worker: Any) -> None:
|
||||
assert import_worker.task_queue == IMPORT_TASK_QUEUE
|
||||
|
||||
async def test_a_workflow_input_reaches_the_first_activity(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_execute: Any,
|
||||
import_build_input: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
"""The claim runs, and the row it could not claim proves the wiring is live.
|
||||
|
||||
An id no row carries is the cheapest input that gets past workflow validation and into
|
||||
`claim_import_status`, so a run that ends failed *there* means the client, the worker, the
|
||||
workflow, the activity and the BFF connection are all in place. Nothing else is exercised.
|
||||
"""
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=987654321,
|
||||
object_key=f'{IMPORT_PREFIX}nothing.sientia',
|
||||
expected_digest='0' * 64,
|
||||
password_envelope='not-a-real-envelope',
|
||||
)
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError) as failure:
|
||||
await handle.result()
|
||||
text_of_failure = failure_text(failure.value)
|
||||
assert 'claim' in text_of_failure.lower()
|
||||
assert '987654321' in text_of_failure
|
||||
assert 'ImportLogRowNotClaimableError' in failure_types(failure.value)
|
||||
assert import_read_status_row(987654321) is None
|
||||
|
||||
async def test_the_insert_helper_writes_a_claimable_row(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
object_key = f'{IMPORT_PREFIX}claimable.sientia'
|
||||
row_id = import_insert_status_row(object_key=object_key, expected_digest='a' * 64)
|
||||
row = import_read_status_row(row_id)
|
||||
assert row is not None
|
||||
assert row['status'] == 'PENDING'
|
||||
assert row['run_type'] == 'IMPORT'
|
||||
assert row['file_name'] == object_key
|
||||
request_data = row['request_data']
|
||||
payload = json.loads(request_data) if isinstance(request_data, str) else request_data
|
||||
assert payload['expected_digest'] == 'a' * 64
|
||||
@@ -1,184 +0,0 @@
|
||||
"""One table, one database — proven from the servers' own logs, not from the application's side.
|
||||
|
||||
Decision D22 says an import writes to `public.experiment_run` and to nothing else relational: no
|
||||
project table, no model table, no row in the predictions database. That claim is only worth making
|
||||
if it is checked from somewhere that sees *every* client, and there is exactly one such place.
|
||||
|
||||
Why not the obvious alternatives:
|
||||
|
||||
- a SQLAlchemy `before_cursor_execute` listener sees the status engine and is blind to the psycopg
|
||||
pool `Storage` holds — which is precisely the client that would betray the claim;
|
||||
- `pg_stat_statements` is aggregated, so it cannot be scoped to the window of one import.
|
||||
|
||||
So both containers run with `log_statement=all` and the window is bracketed by a marker statement on
|
||||
each server. The assertions are about *which relation names appear*, not about exact statement text,
|
||||
because the text is the driver's business and changes with it.
|
||||
|
||||
The predictions database is deliberately real and reachable here: `import_activities` is built with
|
||||
working credentials for it, so "no statement reached it" is a claim about the import's behaviour and
|
||||
not about a missing connection.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from temporalio.client import WorkflowFailureError
|
||||
|
||||
from e2e.import_harness import names_no_relation, strip_literals
|
||||
|
||||
pytestmark = [pytest.mark.import_e2e, pytest.mark.asyncio]
|
||||
|
||||
FACTORY_PASSWORD = 'Factory1Model!Test'
|
||||
|
||||
# Relations that exist in the predictions database and in nobody's import path.
|
||||
FOREIGN_RELATIONS = (
|
||||
'predictions_schema',
|
||||
'laborious_data',
|
||||
'transformed_data',
|
||||
'predictions',
|
||||
)
|
||||
|
||||
# Tables a reader might expect an import to want, and which decision D22 says do not exist.
|
||||
RELATIONS_D22_DENIES = ('project', 'model_registry', 'models', 'import_run', 'model_project')
|
||||
|
||||
|
||||
class TestOneTableOneDatabase:
|
||||
"""A full import, from claim to terminal write, watched from both servers."""
|
||||
|
||||
async def test_every_statement_the_import_issues_names_experiment_run(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_log_window: Any,
|
||||
) -> None:
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
|
||||
with import_log_window() as captured:
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
await handle.result()
|
||||
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None
|
||||
assert row['status'] == 'IMPORT_SUCCESS', (
|
||||
'the window has to bracket a *working* import, or it proves nothing'
|
||||
)
|
||||
|
||||
bff_statements = captured['bff']
|
||||
assert bff_statements, 'the status database logged nothing at all — the window is wrong'
|
||||
|
||||
touching_a_relation = [
|
||||
statement for statement in bff_statements if not names_no_relation(statement)
|
||||
]
|
||||
assert touching_a_relation, 'no relational statement was logged, so nothing was checked'
|
||||
for statement in touching_a_relation:
|
||||
assert 'experiment_run' in statement, (
|
||||
f'a statement naming another relation reached the status database: {statement!r}'
|
||||
)
|
||||
|
||||
# And no relation a reader might expect the feature to want appears anywhere in the window.
|
||||
# Literals are stripped first: the import writes the bundle's own `model_project` value into
|
||||
# the JSON detail, and `factory_project` is data, not a reference to a `project` table.
|
||||
window_text = '\n'.join(strip_literals(s) for s in bff_statements).lower()
|
||||
for relation in RELATIONS_D22_DENIES:
|
||||
assert relation not in window_text, (
|
||||
f'{relation!r} appears in the import window: D22 says no such table is written'
|
||||
)
|
||||
for relation in FOREIGN_RELATIONS:
|
||||
assert relation not in window_text, f'{relation!r} appears in the import window'
|
||||
|
||||
async def test_the_predictions_database_sees_no_statement_from_an_import(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_log_window: Any,
|
||||
) -> None:
|
||||
"""Not "no statement against our tables" — no statement at all.
|
||||
|
||||
The import holds a working connection to this server (`Activities` is built with one), so
|
||||
this is about what the import chooses to do, not about what it is able to do.
|
||||
"""
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
|
||||
with import_log_window() as captured:
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
await handle.result()
|
||||
|
||||
assert import_read_status_row(import_run_id)['status'] == 'IMPORT_SUCCESS'
|
||||
|
||||
predictions_statements = [
|
||||
statement for statement in captured['predictions'] if not names_no_relation(statement)
|
||||
]
|
||||
assert not predictions_statements, (
|
||||
f'the import issued statements against the predictions database: '
|
||||
f'{predictions_statements}'
|
||||
)
|
||||
|
||||
async def test_a_rejected_import_also_touches_nothing_else(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_log_window: Any,
|
||||
) -> None:
|
||||
"""The failure path writes more status rows than the success path, so it is worth its own case.
|
||||
|
||||
A rejected import runs the claim, the progress hints, cleanup and the terminal write — more
|
||||
statements than a success, and every one of them still has to name the same table.
|
||||
"""
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory(include_metadata=False))
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
|
||||
with import_log_window() as captured:
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError):
|
||||
await handle.result()
|
||||
|
||||
for statement in captured['bff']:
|
||||
if names_no_relation(statement):
|
||||
continue
|
||||
assert 'experiment_run' in statement, statement
|
||||
assert not [s for s in captured['predictions'] if not names_no_relation(s)]
|
||||
@@ -1,247 +0,0 @@
|
||||
"""What the worker is left holding, what the bucket is left promising, and what the history reveals.
|
||||
|
||||
Three things an import must be true about after it finishes, whichever way it finished:
|
||||
|
||||
- **The worker's disk is empty of it.** No extracted tree, no downloaded object, no decrypted zip.
|
||||
The bundle is decrypted onto local disk, so a leak here is a plaintext model sitting in a scratch
|
||||
directory for the next import to trip over — and after a *failure* is exactly when that would go
|
||||
unnoticed.
|
||||
- **The bucket carries the retention rule.** The import never deletes the uploaded object: it is the
|
||||
record of an attempt, and the lifecycle rule is what eventually removes it. That rule is applied
|
||||
by the importer rather than assumed from a Helm value, so it is checkable from here.
|
||||
- **The status record outlives the object**, because it is the audit trail and the object is not.
|
||||
|
||||
And the one claim that can only be made against a real Temporal history: no plaintext password
|
||||
anywhere in it. The workflow carries an envelope it cannot open, and the activity that opens it
|
||||
never returns what it found. That is the design; this reads the history and checks it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from temporalio.client import WorkflowFailureError
|
||||
|
||||
from e2e.import_harness import IMPORT_BUCKET, IMPORT_PREFIX
|
||||
from tests.helpers.bundle_factory import DEFAULT_ENVELOPE_KEY, envelope_key_b64
|
||||
|
||||
pytestmark = [pytest.mark.import_e2e, pytest.mark.asyncio]
|
||||
|
||||
FACTORY_PASSWORD = 'Factory1Model!Test'
|
||||
RETENTION_RULE_ID = 'sientia-import-bundle-expiry'
|
||||
|
||||
|
||||
async def _run(
|
||||
*,
|
||||
payload: bytes,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
password: str = FACTORY_PASSWORD,
|
||||
) -> tuple[Any, int, str]:
|
||||
object_key, digest = import_upload_bundle(payload)
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(password),
|
||||
)
|
||||
)
|
||||
return handle, import_run_id, object_key
|
||||
|
||||
|
||||
class TestWorkerFilesystem:
|
||||
"""Nothing of the import is left on the worker — on the success path or on either failure path."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
('case', 'factory_kwargs', 'expect_success'),
|
||||
[
|
||||
('success', {}, True),
|
||||
('rejected at decryption', {'drop_final_tag': True}, False),
|
||||
('rejected at structure validation', {'include_metadata': False}, False),
|
||||
],
|
||||
)
|
||||
async def test_the_work_directory_is_empty_afterwards(
|
||||
self,
|
||||
case: str,
|
||||
factory_kwargs: dict[str, Any],
|
||||
expect_success: bool,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
tmp_path: Any,
|
||||
) -> None:
|
||||
"""Two different failure steps, because they clean up at different points.
|
||||
|
||||
A decryption failure leaves a partially written zip in the staging directory; a structure
|
||||
failure has already extracted a whole tree. The reader unwinds both, and the cleanup activity
|
||||
removes the scratch directory after either — but they are separate code paths, so asserting
|
||||
one would not cover the other.
|
||||
"""
|
||||
handle, _, _ = await _run(
|
||||
payload=import_bundle_factory(**factory_kwargs),
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
if expect_success:
|
||||
await handle.result()
|
||||
else:
|
||||
with pytest.raises(WorkflowFailureError):
|
||||
await handle.result()
|
||||
|
||||
work_dir = tmp_path / 'import-work'
|
||||
leftovers = sorted(str(path) for path in work_dir.rglob('*')) if work_dir.exists() else []
|
||||
assert not leftovers, f'{case}: the worker still holds {leftovers}'
|
||||
|
||||
|
||||
class TestRetention:
|
||||
"""The rule the uploaded object's eventual removal depends on."""
|
||||
|
||||
async def test_the_import_bucket_carries_the_retention_rule(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_bundle_factory: Any,
|
||||
import_bucket: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
) -> None:
|
||||
handle, _, object_key = await _run(
|
||||
payload=import_bundle_factory(),
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
await handle.result()
|
||||
|
||||
lifecycle = import_bucket.get_bucket_lifecycle(IMPORT_BUCKET)
|
||||
assert lifecycle is not None, 'the bucket carries no lifecycle configuration at all'
|
||||
rules = {rule.rule_id: rule for rule in lifecycle.rules}
|
||||
assert RETENTION_RULE_ID in rules, (
|
||||
f'no {RETENTION_RULE_ID!r} rule; the bucket has {sorted(rules)}'
|
||||
)
|
||||
rule = rules[RETENTION_RULE_ID]
|
||||
assert rule.status == 'Enabled'
|
||||
assert rule.rule_filter.prefix == IMPORT_PREFIX
|
||||
assert rule.expiration.days == 7
|
||||
|
||||
# The import did not delete the uploaded object: it is the record of an attempt, and the
|
||||
# rule above is what removes it later.
|
||||
assert import_bucket.stat_object(IMPORT_BUCKET, object_key).size > 0
|
||||
|
||||
async def test_the_status_record_outlives_the_uploaded_object(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_bundle_factory: Any,
|
||||
import_bucket: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
"""Deleting the object is what the retention rule will eventually do; the verdict must survive.
|
||||
|
||||
The record is the audit trail. If reading a finished import's outcome depended on the bundle
|
||||
still being in the bucket, every import would become unreadable after seven days.
|
||||
"""
|
||||
handle, import_run_id, object_key = await _run(
|
||||
payload=import_bundle_factory(),
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
await handle.result()
|
||||
|
||||
import_bucket.remove_object(IMPORT_BUCKET, object_key)
|
||||
with pytest.raises(Exception): # noqa: B017 - MinIO's own not-found type is not the subject
|
||||
import_bucket.stat_object(IMPORT_BUCKET, object_key)
|
||||
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None
|
||||
assert row['status'] == 'IMPORT_SUCCESS'
|
||||
assert row['file_name'] == object_key
|
||||
assert row['run_name'] is not None
|
||||
|
||||
|
||||
class TestEventHistorySecrecy:
|
||||
"""The claim that only a real event history can settle."""
|
||||
|
||||
async def test_no_plaintext_password_or_key_appears_in_the_history(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_history_events: Any,
|
||||
) -> None:
|
||||
"""Dump every event of a completed import and search the whole thing.
|
||||
|
||||
Three values must be absent, and they are absent for three different reasons:
|
||||
|
||||
- the **bundle password** never crosses an activity boundary — `open_import_bundle` decrypts
|
||||
the envelope inside itself and returns metadata, never the password;
|
||||
- the **envelope key** is configuration, read from the environment inside the activity, so it
|
||||
is never an argument to anything;
|
||||
- the **decrypted password** is the first of these two by another name, and is searched for
|
||||
separately because a future refactor that returned it "just for logging" would pass a test
|
||||
that only looked for the literal input.
|
||||
|
||||
The envelope itself *is* in the history, and must be: it is how the password reaches the
|
||||
activity. That it is there is asserted too, so this test cannot pass by finding nothing at
|
||||
all — which is the way an assertion like this usually rots.
|
||||
"""
|
||||
password = 'Secret1Password!InHistory'
|
||||
handle, _, _ = await _run(
|
||||
payload=import_bundle_factory(password=password),
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
password=password,
|
||||
)
|
||||
await handle.result()
|
||||
|
||||
events, blob = await import_history_events(handle)
|
||||
assert events, 'the history is empty'
|
||||
|
||||
assert password not in blob, 'the bundle password appears in the event history'
|
||||
assert envelope_key_b64(DEFAULT_ENVELOPE_KEY) not in blob, (
|
||||
'the envelope key appears in the event history'
|
||||
)
|
||||
assert DEFAULT_ENVELOPE_KEY.hex() not in blob.lower(), (
|
||||
'the envelope key appears in the event history, hex-encoded'
|
||||
)
|
||||
|
||||
# The negative assertions above are only meaningful if the history is actually readable and
|
||||
# actually carries the envelope. Without this, an empty or opaque history would "pass".
|
||||
assert 'password_envelope' in blob, (
|
||||
'the history does not even carry the envelope, so the search above proved nothing'
|
||||
)
|
||||
@@ -1,214 +0,0 @@
|
||||
"""The digest is a promise about the bytes in the bucket, and it is kept against the bucket.
|
||||
|
||||
Three ways the promise can be broken, and where each is caught:
|
||||
|
||||
- the uploader computed the digest over *other* bytes — caught at `download`, after the object is
|
||||
fetched and before its header is parsed, so a wrong file never costs an Argon2id pass;
|
||||
- the object was replaced between the upload and the run — the same failure, which is the point:
|
||||
the check is against what is stored now, not against what was stored when the row was written;
|
||||
- no digest at all — caught at `received`, before anything is downloaded, because a request with no
|
||||
digest is not a request this workflow can honour.
|
||||
|
||||
The first two matter because MinIO is real here. A mocked object store would return whatever the
|
||||
test told it to, and "the digest is enforced against the bytes that were uploaded" would be a
|
||||
statement about the mock.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from temporalio.client import WorkflowFailureError
|
||||
|
||||
from e2e.import_harness import IMPORT_BUCKET, failure_text
|
||||
from laborious.utils.bundle.steps import ImportStep
|
||||
from laborious.utils.models.import_errors import IMPORT_FAILURE_CATALOG
|
||||
|
||||
pytestmark = [pytest.mark.import_e2e, pytest.mark.asyncio]
|
||||
|
||||
VALIDATION_ERROR_STATUS = 'ORCHESTRATOR_VALIDATION_ERROR'
|
||||
FACTORY_PASSWORD = 'Factory1Model!Test'
|
||||
|
||||
|
||||
def _as_dict(value: Any) -> dict[str, Any]:
|
||||
return json.loads(value) if isinstance(value, str) else (value or {})
|
||||
|
||||
|
||||
def _assert_step(row: dict[str, Any], step: ImportStep, workflow_id: str) -> None:
|
||||
code, reason, _responsibility = IMPORT_FAILURE_CATALOG[step]
|
||||
detail = _as_dict(row['orchestrator_response_data'])['import']
|
||||
assert row['status'] == VALIDATION_ERROR_STATUS, row['status']
|
||||
assert detail['step'] == step.value, detail
|
||||
assert detail['code'] == code.value, detail
|
||||
assert detail['workflow_id'] == workflow_id
|
||||
assert row['error_message'] == reason
|
||||
|
||||
|
||||
class TestWrongDigest:
|
||||
"""A digest computed over bytes that are not the ones in the bucket."""
|
||||
|
||||
async def test_a_digest_from_another_file_fails_at_download(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
"""The bundle is perfectly valid; only the digest is another file's.
|
||||
|
||||
The row carries the same wrong digest as the input, so the request and the record agree and
|
||||
the `received` guard passes — otherwise this would be a `received` failure about a
|
||||
contradictory request instead of the download-step digest check it is meant to exercise.
|
||||
"""
|
||||
payload = import_bundle_factory()
|
||||
other_digest = hashlib.sha256(
|
||||
import_bundle_factory(metadata={'run_id': 'other'})
|
||||
).hexdigest()
|
||||
assert other_digest != hashlib.sha256(payload).hexdigest()
|
||||
|
||||
object_key, _ = import_upload_bundle(payload)
|
||||
import_run_id = import_insert_status_row(
|
||||
object_key=object_key, expected_digest=other_digest
|
||||
)
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=other_digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError) as failure:
|
||||
await handle.result()
|
||||
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None
|
||||
_assert_step(row, ImportStep.DOWNLOAD, handle.id)
|
||||
|
||||
# No header was parsed: the failure names the digest, not the format. If the digest check
|
||||
# ran *after* the header the message here would be about magic bytes or a version.
|
||||
text_of_failure = failure_text(failure.value).lower()
|
||||
assert 'digest' in text_of_failure, text_of_failure
|
||||
assert 'marker' not in text_of_failure and 'version' not in text_of_failure
|
||||
|
||||
# The names are still unwritten, which is the other half of "no header was parsed": the
|
||||
# bundle was never opened, so nothing knew what to call it.
|
||||
assert row['run_name'] is None
|
||||
assert row['experiment_name'] == 'import-e2e-placeholder'
|
||||
|
||||
import_assert_nothing_provisioned(
|
||||
experiment_name=import_names['experiment_name'],
|
||||
model_name=import_names['model_name'],
|
||||
object_key=object_key,
|
||||
)
|
||||
|
||||
|
||||
class TestOverwrittenObject:
|
||||
"""The bytes changed after the row was written. Only a real object store can show this."""
|
||||
|
||||
async def test_replacing_the_object_after_upload_fails_at_download(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_bucket: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
original = import_bundle_factory()
|
||||
object_key, digest = import_upload_bundle(original)
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
|
||||
# Same key, different bytes — also a valid bundle, so nothing but the digest can tell.
|
||||
replacement = import_bundle_factory(metadata={'run_id': 'ff00ff00ff00ff00ff00ff00ff00ff00'})
|
||||
assert hashlib.sha256(replacement).hexdigest() != digest
|
||||
import_bucket.put_object(
|
||||
IMPORT_BUCKET,
|
||||
object_key,
|
||||
io.BytesIO(replacement),
|
||||
length=len(replacement),
|
||||
content_type='application/octet-stream',
|
||||
)
|
||||
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError):
|
||||
await handle.result()
|
||||
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None
|
||||
_assert_step(row, ImportStep.DOWNLOAD, handle.id)
|
||||
import_assert_nothing_provisioned(
|
||||
experiment_name=import_names['experiment_name'],
|
||||
model_name=import_names['model_name'],
|
||||
object_key=object_key,
|
||||
)
|
||||
|
||||
|
||||
class TestMissingDigest:
|
||||
"""No digest is refused at the door, before a single byte is fetched."""
|
||||
|
||||
async def test_an_absent_digest_fails_at_received_with_nothing_downloaded(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
tmp_path: Any,
|
||||
) -> None:
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=None,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError) as failure:
|
||||
await handle.result()
|
||||
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None
|
||||
_assert_step(row, ImportStep.RECEIVED, handle.id)
|
||||
assert 'expected_digest' in failure_text(failure.value)
|
||||
|
||||
# Nothing downloaded: the work directory was never even created.
|
||||
work_dir = tmp_path / 'import-work'
|
||||
assert not work_dir.exists() or not list(work_dir.rglob('*'))
|
||||
import_assert_nothing_provisioned(
|
||||
experiment_name=import_names['experiment_name'],
|
||||
model_name=import_names['model_name'],
|
||||
object_key=object_key,
|
||||
)
|
||||
@@ -1,679 +0,0 @@
|
||||
"""One rejection per gate, against a real object, a real archive and a real filesystem.
|
||||
|
||||
The gates are unit-tested already. What only a real run can show is that a rejection is *recorded*
|
||||
correctly and that it provisions nothing: the row ends `ORCHESTRATOR_VALIDATION_ERROR` naming that
|
||||
gate's step, its catalog code and its catalog sentence, and no experiment, run, artifact, registry
|
||||
version, model document or extracted file survives.
|
||||
|
||||
Every case here is hostile **input** — design E7's first preference. Nothing is patched, no
|
||||
dependency is broken; the bundles are built by the shared factory and the gates do the rest. The
|
||||
one exception to "input only" is the object-size ceiling, where the ceiling itself is narrowed
|
||||
rather than a gigabyte being uploaded; that is a configuration change, not a patched code path, and
|
||||
it is scoped to the one class that needs it.
|
||||
|
||||
`_assert_rejected` is where the shared claim of task 5.8 lives — terminal status, recorded step,
|
||||
recorded code, catalog sentence, and `assert_nothing_provisioned` — so no case can quietly assert
|
||||
four fifths of it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import tempfile
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from temporalio.client import WorkflowFailureError
|
||||
|
||||
from e2e.import_harness import (
|
||||
IMPORT_BUCKET,
|
||||
IMPORT_PREFIX,
|
||||
failure_text,
|
||||
)
|
||||
from laborious.utils.bundle.steps import ImportStep
|
||||
from laborious.utils.models.import_errors import IMPORT_FAILURE_CATALOG
|
||||
|
||||
pytestmark = [pytest.mark.import_e2e, pytest.mark.asyncio]
|
||||
|
||||
VALIDATION_ERROR_STATUS = 'ORCHESTRATOR_VALIDATION_ERROR'
|
||||
FACTORY_PASSWORD = 'Factory1Model!Test'
|
||||
|
||||
|
||||
def _as_dict(value: Any) -> dict[str, Any]:
|
||||
return json.loads(value) if isinstance(value, str) else (value or {})
|
||||
|
||||
|
||||
async def _run_rejected_import(
|
||||
*,
|
||||
payload: bytes,
|
||||
password: str,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
envelope_password: str | None = None,
|
||||
) -> tuple[Any, int, str]:
|
||||
"""Upload a hostile bundle, insert the row, run the import and expect it to fail."""
|
||||
object_key, digest = import_upload_bundle(payload)
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(
|
||||
envelope_password if envelope_password is not None else password
|
||||
),
|
||||
)
|
||||
)
|
||||
return handle, import_run_id, object_key
|
||||
|
||||
|
||||
async def _assert_rejected(
|
||||
*,
|
||||
handle: Any,
|
||||
import_run_id: int,
|
||||
object_key: str,
|
||||
step: ImportStep,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
import_names: dict[str, str],
|
||||
experiment_name: str | None = None,
|
||||
model_name: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""The whole of task 5.8, in one call. Returns the row so a case can add to it."""
|
||||
with pytest.raises(WorkflowFailureError):
|
||||
await handle.result()
|
||||
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None, 'the rejection lost the row it was recorded on'
|
||||
assert row['status'] == VALIDATION_ERROR_STATUS, (
|
||||
f'terminal status is {row["status"]!r}, expected {VALIDATION_ERROR_STATUS!r}'
|
||||
)
|
||||
|
||||
expected_code, expected_reason, _expected_responsibility = IMPORT_FAILURE_CATALOG[step]
|
||||
detail = _as_dict(row['orchestrator_response_data'])['import']
|
||||
assert detail['step'] == step.value, (
|
||||
f'recorded step is {detail["step"]!r}, expected {step.value!r}'
|
||||
)
|
||||
assert detail['code'] == expected_code.value, (
|
||||
f'recorded code is {detail["code"]!r}, expected {expected_code.value!r}'
|
||||
)
|
||||
assert detail['workflow_id'] == handle.id
|
||||
assert row['error_message'] == expected_reason, (
|
||||
'error_message is not the catalog sentence for this step'
|
||||
)
|
||||
|
||||
import_assert_nothing_provisioned(
|
||||
experiment_name=experiment_name or import_names['experiment_name'],
|
||||
model_name=model_name or import_names['model_name'],
|
||||
object_key=object_key,
|
||||
)
|
||||
return row
|
||||
|
||||
|
||||
class TestGate1ObjectLimits:
|
||||
"""Gate 1 rejects before the object's body is fetched at all."""
|
||||
|
||||
@pytest.fixture
|
||||
def import_limits_override(self) -> dict[str, str]:
|
||||
"""A 1 KiB ceiling, so a factory bundle is over it.
|
||||
|
||||
The production ceiling is 1 GiB. Narrowing it is the only way to cross gate 1 without
|
||||
uploading a gigabyte, and it changes a configured limit rather than a code path — the same
|
||||
comparison runs, against a different number.
|
||||
"""
|
||||
return {'IMPORT_MAX_OBJECT_BYTES': '1024'}
|
||||
|
||||
async def test_an_oversized_object_is_rejected_at_download(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
payload = import_bundle_factory()
|
||||
assert len(payload) > 1024, 'the fixture bundle must exceed the narrowed ceiling'
|
||||
handle, import_run_id, object_key = await _run_rejected_import(
|
||||
payload=payload,
|
||||
password=FACTORY_PASSWORD,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
await _assert_rejected(
|
||||
handle=handle,
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
step=ImportStep.DOWNLOAD,
|
||||
import_read_status_row=import_read_status_row,
|
||||
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
|
||||
import_names=import_names,
|
||||
)
|
||||
|
||||
|
||||
class TestGate2Header:
|
||||
"""Four header rejections, all landing on `download` — the step that says "not these bytes"."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
('case', 'header_kwargs'),
|
||||
[
|
||||
('wrong magic', {'magic': b'NOTSIENT'}),
|
||||
('unsupported version', {'format_version': 9}),
|
||||
('kdf memory above the ceiling', {'declared_memlimit': 512 * 1024 * 1024}),
|
||||
('kdf passes above the ceiling', {'declared_opslimit': 9}),
|
||||
('digest not of the chunk stream', {'digest_override': bytes(32)}),
|
||||
],
|
||||
)
|
||||
async def test_a_malformed_header_is_rejected_at_download(
|
||||
self,
|
||||
case: str,
|
||||
header_kwargs: dict[str, Any],
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
handle, import_run_id, object_key = await _run_rejected_import(
|
||||
payload=import_bundle_factory(**header_kwargs),
|
||||
password=FACTORY_PASSWORD,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
row = await _assert_rejected(
|
||||
handle=handle,
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
step=ImportStep.DOWNLOAD,
|
||||
import_read_status_row=import_read_status_row,
|
||||
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
|
||||
import_names=import_names,
|
||||
)
|
||||
# None of these may borrow the decryption sentence: the password is not implicated.
|
||||
assert 'password' not in (row['error_message'] or '').lower(), case
|
||||
|
||||
|
||||
class TestGate3Decryption:
|
||||
"""The password, the tag and the framing — the three ways a stream refuses to open."""
|
||||
|
||||
async def test_a_wrong_password_is_rejected_at_decryption(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
"""The envelope opens — it is the *bundle* password inside it that is wrong.
|
||||
|
||||
Worth separating: an envelope the worker cannot open is also a decryption-step failure, but
|
||||
for a different reason, and only this one proves the KDF actually ran and produced a key
|
||||
the stream rejected.
|
||||
"""
|
||||
handle, import_run_id, object_key = await _run_rejected_import(
|
||||
payload=import_bundle_factory(password=FACTORY_PASSWORD),
|
||||
password=FACTORY_PASSWORD,
|
||||
envelope_password='Wrong1Password!Entirely',
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
await _assert_rejected(
|
||||
handle=handle,
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
step=ImportStep.DECRYPTION,
|
||||
import_read_status_row=import_read_status_row,
|
||||
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
|
||||
import_names=import_names,
|
||||
)
|
||||
|
||||
async def test_a_tampered_chunk_is_rejected_at_decryption(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
"""One flipped byte inside the stream, with the header digest repaired to match.
|
||||
|
||||
Repairing the digest is what makes this gate 3 rather than gate 2: the file is
|
||||
self-consistent, so only the AEAD can tell it was altered.
|
||||
"""
|
||||
from tests.helpers.bundle_factory import tamper_ciphertext
|
||||
|
||||
handle, import_run_id, object_key = await _run_rejected_import(
|
||||
payload=tamper_ciphertext(import_bundle_factory()),
|
||||
password=FACTORY_PASSWORD,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
await _assert_rejected(
|
||||
handle=handle,
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
step=ImportStep.DECRYPTION,
|
||||
import_read_status_row=import_read_status_row,
|
||||
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
|
||||
import_names=import_names,
|
||||
)
|
||||
|
||||
async def test_a_stream_without_its_final_tag_is_rejected_at_decryption(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
"""Every chunk framed as MESSAGE, so the stream claims more is coming and never ends.
|
||||
|
||||
The digest cannot see this — the bytes on disk are exactly the ones it covers. Only the
|
||||
`secretstream` framing knows the file is incomplete.
|
||||
"""
|
||||
handle, import_run_id, object_key = await _run_rejected_import(
|
||||
payload=import_bundle_factory(drop_final_tag=True),
|
||||
password=FACTORY_PASSWORD,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
await _assert_rejected(
|
||||
handle=handle,
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
step=ImportStep.DECRYPTION,
|
||||
import_read_status_row=import_read_status_row,
|
||||
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
|
||||
import_names=import_names,
|
||||
)
|
||||
|
||||
|
||||
class TestGate4ArchiveInspection:
|
||||
"""The central directory alone decides — nothing has been written to disk yet."""
|
||||
|
||||
@pytest.fixture
|
||||
def import_limits_override(self) -> dict[str, str]:
|
||||
"""A 1 MiB declared-expansion ceiling, for the oversized-expansion case in this class."""
|
||||
return {'IMPORT_MAX_UNCOMPRESSED_BYTES': str(1024 * 1024)}
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
('case', 'kwargs'),
|
||||
[
|
||||
('entry escaping the tree', {'traversal_entry': 'artifacts/../../escaped.csv'}),
|
||||
('absolute entry path', {'absolute_entry': '/opt/absolute.csv'}),
|
||||
('symbolic link', {'symlinks': {'artifacts/link.csv': 'artifacts/train_data.csv'}}),
|
||||
('declared expansion over the ceiling', {'bomb_bytes': 4 * 1024 * 1024}),
|
||||
],
|
||||
)
|
||||
async def test_a_hostile_archive_is_rejected_before_anything_is_extracted(
|
||||
self,
|
||||
case: str,
|
||||
kwargs: dict[str, Any],
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
tmp_path: Any,
|
||||
) -> None:
|
||||
handle, import_run_id, object_key = await _run_rejected_import(
|
||||
payload=import_bundle_factory(**kwargs),
|
||||
password=FACTORY_PASSWORD,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
await _assert_rejected(
|
||||
handle=handle,
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
step=ImportStep.ARCHIVE_INSPECTION,
|
||||
import_read_status_row=import_read_status_row,
|
||||
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
|
||||
import_names=import_names,
|
||||
)
|
||||
# "Nothing extracted" is stronger than "nothing provisioned": not one member reached disk.
|
||||
work_dir = tmp_path / 'import-work'
|
||||
assert not list(work_dir.rglob('*')) if work_dir.exists() else True, case
|
||||
|
||||
|
||||
class TestGate5Extraction:
|
||||
"""A member that gate 4 passed and extraction still cannot place.
|
||||
|
||||
Gate 5 is about *placement*, not size: containment on the resolved path, and never overwriting
|
||||
a file already written. A zip carrying the same name twice passes gate 4 — which checks names
|
||||
and modes, not duplicates — and is refused here on the second member.
|
||||
"""
|
||||
|
||||
# `zipfile` warns while *writing* the duplicate entry. That is the fixture doing exactly what
|
||||
# this test needs, so it is silenced here rather than left to look like a problem — and scoped to
|
||||
# this one test, so a duplicate name appearing anywhere else still gets noticed.
|
||||
@pytest.mark.filterwarnings('ignore:Duplicate name:UserWarning')
|
||||
async def test_a_duplicate_member_is_rejected_at_extraction(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
from tests.helpers.bundle_factory import default_tree
|
||||
|
||||
# `bomb_bytes` writes `artifacts/bomb.csv`, and `files` already carries that name, so the
|
||||
# archive declares it twice. Small enough that no size ceiling is involved.
|
||||
files = dict(default_tree())
|
||||
files['artifacts/bomb.csv'] = b'first\n'
|
||||
handle, import_run_id, object_key = await _run_rejected_import(
|
||||
payload=import_bundle_factory(files=files, bomb_bytes=64),
|
||||
password=FACTORY_PASSWORD,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
await _assert_rejected(
|
||||
handle=handle,
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
step=ImportStep.EXTRACTION,
|
||||
import_read_status_row=import_read_status_row,
|
||||
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
|
||||
import_names=import_names,
|
||||
)
|
||||
|
||||
|
||||
class TestGate6Structure:
|
||||
"""The layout and the metadata schema, checked on the extracted tree."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
('case', 'kwargs'),
|
||||
[
|
||||
('no metadata.json', {'include_metadata': False}),
|
||||
('metadata missing run_id', {'drop_metadata_field': 'run_id'}),
|
||||
('parameters missing target_variable', {'drop_parameter': 'target_variable'}),
|
||||
],
|
||||
)
|
||||
async def test_an_incomplete_bundle_is_rejected_at_structure_validation(
|
||||
self,
|
||||
case: str,
|
||||
kwargs: dict[str, Any],
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
handle, import_run_id, object_key = await _run_rejected_import(
|
||||
payload=import_bundle_factory(**kwargs),
|
||||
password=FACTORY_PASSWORD,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
await _assert_rejected(
|
||||
handle=handle,
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
step=ImportStep.STRUCTURE_VALIDATION,
|
||||
import_read_status_row=import_read_status_row,
|
||||
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
|
||||
import_names=import_names,
|
||||
)
|
||||
|
||||
async def test_a_bundle_without_prediction_model_is_rejected(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
"""Without `prediction_model` there is nothing to register a version against."""
|
||||
from tests.helpers.bundle_factory import default_tree
|
||||
|
||||
files = {
|
||||
name: content
|
||||
for name, content in default_tree().items()
|
||||
if not name.startswith('artifacts/prediction_model/')
|
||||
}
|
||||
handle, import_run_id, object_key = await _run_rejected_import(
|
||||
payload=import_bundle_factory(files=files),
|
||||
password=FACTORY_PASSWORD,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
await _assert_rejected(
|
||||
handle=handle,
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
step=ImportStep.STRUCTURE_VALIDATION,
|
||||
import_read_status_row=import_read_status_row,
|
||||
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
|
||||
import_names=import_names,
|
||||
)
|
||||
|
||||
async def test_an_over_long_experiment_name_never_reaches_the_varchar_column(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
import_bff_engine: Any,
|
||||
) -> None:
|
||||
"""A 60-character experiment name is refused at gate 6, not by Postgres three steps later.
|
||||
|
||||
`experiment_name` is `VARCHAR(50)` and the import writes the bundle's own name onto the row
|
||||
*before* provisioning. If the gate did not enforce the column's limit, the names write would
|
||||
be a `DataError` — non-retryable, recorded as a status failure, on a row whose import was
|
||||
otherwise fine. Asserting the row's `experiment_name` is untouched is what proves the write
|
||||
never happened.
|
||||
"""
|
||||
from sqlalchemy import text
|
||||
|
||||
long_name = 'x' * 60
|
||||
handle, import_run_id, object_key = await _run_rejected_import(
|
||||
payload=import_bundle_factory(metadata={'experiment_name': long_name}),
|
||||
password=FACTORY_PASSWORD,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
row = await _assert_rejected(
|
||||
handle=handle,
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
step=ImportStep.STRUCTURE_VALIDATION,
|
||||
import_read_status_row=import_read_status_row,
|
||||
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
|
||||
import_names=import_names,
|
||||
experiment_name=long_name,
|
||||
)
|
||||
assert row['experiment_name'] == 'import-e2e-placeholder', (
|
||||
'the over-long name reached the row, so the gate did not hold the column'
|
||||
)
|
||||
with import_bff_engine.begin() as connection:
|
||||
width = connection.execute(
|
||||
text(
|
||||
'SELECT character_maximum_length FROM information_schema.columns '
|
||||
"WHERE table_name = 'experiment_run' AND column_name = 'experiment_name'"
|
||||
)
|
||||
).scalar_one()
|
||||
assert width == 50, 'the column this gate exists to protect is no longer VARCHAR(50)'
|
||||
|
||||
|
||||
class TestGate7ContentPolicy:
|
||||
"""A file kind an MLflow artifact tree has no reason to carry."""
|
||||
|
||||
async def test_a_forbidden_file_type_is_rejected_at_content_policy(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
handle, import_run_id, object_key = await _run_rejected_import(
|
||||
payload=import_bundle_factory(extra_files={'artifacts/payload.exe': b'MZ\x00\x00'}),
|
||||
password=FACTORY_PASSWORD,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
await _assert_rejected(
|
||||
handle=handle,
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
step=ImportStep.CONTENT_POLICY,
|
||||
import_read_status_row=import_read_status_row,
|
||||
import_assert_nothing_provisioned=import_assert_nothing_provisioned,
|
||||
import_names=import_names,
|
||||
)
|
||||
|
||||
|
||||
class TestTheRejectionSaysNothingTechnical:
|
||||
"""`error_message` is a sentence for a person: no path, no object key, no traceback.
|
||||
|
||||
One case is enough to make the point structurally — the catalog is keyed by step and
|
||||
`classify_import_failure` discards the exception's text entirely — and the per-gate cases above
|
||||
each assert the sentence is *exactly* the catalog's, which is the same guarantee stated
|
||||
positively.
|
||||
"""
|
||||
|
||||
async def test_no_technical_detail_reaches_the_recorded_reason(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
handle, import_run_id, object_key = await _run_rejected_import(
|
||||
payload=import_bundle_factory(include_metadata=False),
|
||||
password=FACTORY_PASSWORD,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError) as failure:
|
||||
await handle.result()
|
||||
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None
|
||||
message = row['error_message'] or ''
|
||||
for forbidden in (
|
||||
object_key,
|
||||
IMPORT_BUCKET,
|
||||
IMPORT_PREFIX,
|
||||
'experiment_run',
|
||||
'Traceback',
|
||||
'BundleGateError',
|
||||
tempfile.gettempdir(),
|
||||
'gate',
|
||||
):
|
||||
assert forbidden not in message, f'{forbidden!r} leaked into error_message: {message!r}'
|
||||
|
||||
# The technical truth is not lost — it is on the Temporal error, where support can read it.
|
||||
assert 'metadata' in failure_text(failure.value).lower()
|
||||
@@ -1,315 +0,0 @@
|
||||
"""An import that works, twice over: once from the producer's own bytes and once from the factory.
|
||||
|
||||
Two origins, because they prove different things. The **golden** bundle is bytes the producer's code
|
||||
emitted — it decrypts at the producer's real Argon2id cost (256 MiB, 3 passes) and carries the tree
|
||||
the producer really writes, so it is the only case that would notice the producer's layout moving.
|
||||
The **factory** bundle is synthesised here, cheaply, with names unique to the test; it is what the
|
||||
rest of the suite runs on, so a happy path over it is what makes the negative cases' "and nothing
|
||||
was provisioned" meaningful.
|
||||
|
||||
What "it worked" means, in full: the experiment exists, a run exists inside it carrying the derived
|
||||
run name, the artifact tree is under the run with `prediction_model` and `data_model` at its root,
|
||||
every parameter the bundle declared is logged plus the origin's identity under `source.*`, a
|
||||
registry version exists, the MongoDB document exists, and the status row ends `IMPORT_SUCCESS`.
|
||||
|
||||
Two things a successful import must **not** have done are asserted here too, because they are the
|
||||
difference between "imported" and "live": no version is in `Production`, and the document is
|
||||
`active: false`. Promotion and activation are manual operator steps (QTZPOC-21).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import mlflow
|
||||
import pytest
|
||||
|
||||
from e2e.import_harness import (
|
||||
IMPORT_BUCKET,
|
||||
IMPORT_PREFIX,
|
||||
SEEDED_MODEL_IDS,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.import_e2e, pytest.mark.asyncio]
|
||||
|
||||
GOLDEN_MODEL_NAME = 'fixture_model'
|
||||
GOLDEN_EXPERIMENT_NAME = 'fixture_experiment'
|
||||
GOLDEN_MODEL_VERSION = '3'
|
||||
GOLDEN_TARGET = 'SE'
|
||||
|
||||
REQUIRED_ARTIFACT_ROOTS = {'prediction_model', 'data_model'}
|
||||
|
||||
|
||||
def _mlflow_client(server: dict[str, Any]) -> Any:
|
||||
return mlflow.tracking.MlflowClient(tracking_uri=server['tracking_uri'])
|
||||
|
||||
|
||||
def _as_dict(value: Any) -> dict[str, Any]:
|
||||
"""`request_data` / `orchestrator_response_data` come back as a dict or as JSON text."""
|
||||
return json.loads(value) if isinstance(value, str) else (value or {})
|
||||
|
||||
|
||||
async def _run_import(
|
||||
*,
|
||||
payload: bytes,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
password: str,
|
||||
import_envelope: Any,
|
||||
) -> tuple[Any, int, str]:
|
||||
"""Upload, insert the row the BFF would have inserted, run the workflow. Return the outcome."""
|
||||
object_key, digest = import_upload_bundle(payload)
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(password),
|
||||
)
|
||||
)
|
||||
return handle, import_run_id, object_key
|
||||
|
||||
|
||||
def _assert_provisioned(
|
||||
*,
|
||||
client: Any,
|
||||
experiment_name: str,
|
||||
model_name: str,
|
||||
run_name: str,
|
||||
expected_parameters: dict[str, str],
|
||||
expected_source_fields: dict[str, str],
|
||||
) -> tuple[str, Any]:
|
||||
"""The whole positive claim about MLflow, in one place. Returns `(run_id, version)`."""
|
||||
experiment = client.get_experiment_by_name(experiment_name)
|
||||
assert experiment is not None, f'no experiment named {experiment_name!r}'
|
||||
|
||||
runs = [
|
||||
run
|
||||
for run in client.search_runs([experiment.experiment_id])
|
||||
if run.info.run_name == run_name
|
||||
]
|
||||
assert len(runs) == 1, f'expected exactly one run named {run_name!r}, found {len(runs)}'
|
||||
run = runs[0]
|
||||
assert run.info.status == 'FINISHED'
|
||||
|
||||
artifact_roots = {artifact.path for artifact in client.list_artifacts(run.info.run_id)}
|
||||
assert REQUIRED_ARTIFACT_ROOTS <= artifact_roots, (
|
||||
f'the run holds {sorted(artifact_roots)}, missing {sorted(REQUIRED_ARTIFACT_ROOTS - artifact_roots)}'
|
||||
)
|
||||
# `prediction_model` is a directory the legacy loader walks, not an opaque blob.
|
||||
prediction_files = {
|
||||
artifact.path.rsplit('/', 1)[-1]
|
||||
for artifact in client.list_artifacts(run.info.run_id, 'prediction_model')
|
||||
}
|
||||
assert {'MLmodel', 'model.pkl'} <= prediction_files, prediction_files
|
||||
|
||||
for name, value in expected_parameters.items():
|
||||
assert run.data.params.get(name) == value, (
|
||||
f'parameter {name!r} is {run.data.params.get(name)!r}, expected {value!r}'
|
||||
)
|
||||
for name, value in expected_source_fields.items():
|
||||
assert run.data.params.get(f'source.{name}') == value, (
|
||||
f'source.{name} is {run.data.params.get(f"source.{name}")!r}, expected {value!r}'
|
||||
)
|
||||
|
||||
versions = list(client.search_model_versions(f"name='{model_name}'"))
|
||||
assert len(versions) == 1, f'expected one registered version, found {len(versions)}'
|
||||
version = versions[0]
|
||||
assert version.run_id == run.info.run_id
|
||||
return run.info.run_id, version
|
||||
|
||||
|
||||
class TestGoldenBundle:
|
||||
"""The producer's own bytes, at the producer's own KDF cost."""
|
||||
|
||||
async def test_the_producer_artefact_imports_end_to_end(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_golden_bundle: dict[str, Any],
|
||||
import_mlflow_server: dict[str, Any],
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
handle, import_run_id, object_key = await _run_import(
|
||||
payload=import_golden_bundle['bytes'],
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
password=import_golden_bundle['password'],
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
result = await handle.result()
|
||||
|
||||
assert result['model_name'] == GOLDEN_MODEL_NAME
|
||||
assert result['import_run_id'] == import_run_id
|
||||
assert result['version'] == '1', 'the destination numbering starts at 1, not at the bundle'
|
||||
|
||||
run_id, version = _assert_provisioned(
|
||||
client=_mlflow_client(import_mlflow_server),
|
||||
experiment_name=GOLDEN_EXPERIMENT_NAME,
|
||||
model_name=GOLDEN_MODEL_NAME,
|
||||
run_name=f'import-{GOLDEN_MODEL_NAME}-v{GOLDEN_MODEL_VERSION}',
|
||||
expected_parameters={'target_variable': GOLDEN_TARGET, 'nan_treatment': 'drop'},
|
||||
expected_source_fields={
|
||||
'model_name': GOLDEN_MODEL_NAME,
|
||||
'experiment_name': GOLDEN_EXPERIMENT_NAME,
|
||||
'model_version': GOLDEN_MODEL_VERSION,
|
||||
'model_project': 'fixture_project',
|
||||
},
|
||||
)
|
||||
assert result['run_id'] == run_id
|
||||
|
||||
document = import_models_collection.find_one({'name': GOLDEN_MODEL_NAME}, {'_id': False})
|
||||
assert document is not None, 'no model document was written'
|
||||
assert document['id'] == str(max(int(i) for i in SEEDED_MODEL_IDS) + 1)
|
||||
assert document['model_config']['target'] == GOLDEN_TARGET
|
||||
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None
|
||||
assert row['status'] == 'IMPORT_SUCCESS'
|
||||
assert row['error_message'] is None
|
||||
assert row['experiment_name'] == GOLDEN_EXPERIMENT_NAME
|
||||
assert row['run_name'] == f'import-{GOLDEN_MODEL_NAME}-v{GOLDEN_MODEL_VERSION}'
|
||||
assert row['file_name'] == object_key
|
||||
assert row['bucket_name'] == IMPORT_BUCKET
|
||||
assert object_key.startswith(IMPORT_PREFIX)
|
||||
|
||||
# Imported, not live.
|
||||
assert version.current_stage in (None, 'None'), version.current_stage
|
||||
assert document['active'] is False
|
||||
|
||||
|
||||
class TestFactoryBundle:
|
||||
"""A synthesised tree with names unique to this test — the shape the rest of the suite runs on."""
|
||||
|
||||
async def test_a_factory_bundle_imports_end_to_end(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_mlflow_server: dict[str, Any],
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
payload = import_bundle_factory(parameters={'target_variable': 'SE'})
|
||||
handle, import_run_id, _ = await _run_import(
|
||||
payload=payload,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
password='Factory1Model!Test',
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
result = await handle.result()
|
||||
|
||||
assert result['model_name'] == import_names['model_name']
|
||||
run_id, version = _assert_provisioned(
|
||||
client=_mlflow_client(import_mlflow_server),
|
||||
experiment_name=import_names['experiment_name'],
|
||||
model_name=import_names['model_name'],
|
||||
# DEFAULT_METADATA declares model_version '2'.
|
||||
run_name=f'import-{import_names["model_name"]}-v2',
|
||||
expected_parameters={'target_variable': 'SE', 'retrain': 'False'},
|
||||
expected_source_fields={
|
||||
'model_name': import_names['model_name'],
|
||||
'experiment_name': import_names['experiment_name'],
|
||||
'model_version': '2',
|
||||
},
|
||||
)
|
||||
assert result['run_id'] == run_id
|
||||
assert result['version'] == '1'
|
||||
|
||||
document = import_models_collection.find_one(
|
||||
{'name': import_names['model_name']}, {'_id': False}
|
||||
)
|
||||
assert document is not None
|
||||
assert document['model_config'] == {
|
||||
'transform_flavor': 'sklearn',
|
||||
'predict_flavor': 'joblib',
|
||||
'retention_minutes': 60,
|
||||
'target': 'SE',
|
||||
'skip_predict_timestamp': True,
|
||||
}
|
||||
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None
|
||||
assert row['status'] == 'IMPORT_SUCCESS'
|
||||
detail = _as_dict(row['orchestrator_response_data'])
|
||||
assert detail['import']['workflow_id'] == handle.id
|
||||
|
||||
assert version.current_stage in (None, 'None')
|
||||
assert document['active'] is False
|
||||
|
||||
|
||||
class TestLegacyLoaderResolution:
|
||||
"""The registry source an import writes has to be one the *existing* loader can parse.
|
||||
|
||||
This is the assertion that could only ever be made against a real registry: the importer writes
|
||||
the run's resolved artifact URI with `prediction_model` appended, because `get_model_run_id`
|
||||
extracts the run id as `source.split('/')[2]` — an index that is the run id only in the resolved
|
||||
form. A literal `runs:/<run_id>/prediction_model` would put `prediction_model` at index 2, and a
|
||||
`models:/m-*` source would have none of the parts at all.
|
||||
|
||||
It is called with the stage the version is actually in rather than `Production`, because an
|
||||
import deliberately promotes nothing — passing `Production` here would test the promotion step
|
||||
this feature does not have.
|
||||
"""
|
||||
|
||||
async def test_the_imported_source_resolves_to_the_import_run(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_activities: Any,
|
||||
import_mlflow_server: dict[str, Any],
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
) -> None:
|
||||
payload = import_bundle_factory()
|
||||
handle, _, _ = await _run_import(
|
||||
payload=payload,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
password='Factory1Model!Test',
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
result = await handle.result()
|
||||
model_name = import_names['model_name']
|
||||
|
||||
client = _mlflow_client(import_mlflow_server)
|
||||
version = next(iter(client.search_model_versions(f"name='{model_name}'")))
|
||||
assert version.source.endswith('/prediction_model'), version.source
|
||||
|
||||
# The production method, on the production repository object, against the real registry.
|
||||
resolved = import_activities.import_mlflow_repository.get_model_run_id(
|
||||
model_name, stage=version.current_stage
|
||||
)
|
||||
assert resolved == result['run_id'], (
|
||||
f'get_model_run_id resolved {resolved!r} from source {version.source!r}, '
|
||||
f'but the import created run {result["run_id"]!r}'
|
||||
)
|
||||
|
||||
artifact_roots = {artifact.path for artifact in client.list_artifacts(resolved)}
|
||||
assert REQUIRED_ARTIFACT_ROOTS <= artifact_roots, sorted(artifact_roots)
|
||||
@@ -1,384 +0,0 @@
|
||||
"""Each provisioning step fails as itself: its own code, its own sentence, nothing half-built.
|
||||
|
||||
The four steps after the gates — experiment creation, artifact upload, registration, the model
|
||||
document — each get one failure, and every one lands `IMPORT_ERROR` rather than
|
||||
`ORCHESTRATOR_VALIDATION_ERROR`. That split is not cosmetic: it is the table's existing distinction
|
||||
between "the person who submitted this can fix it" and "only support can", and a provisioning
|
||||
failure is always the second kind.
|
||||
|
||||
**How each failure is caused**, in design E7's order of preference:
|
||||
|
||||
| Step | Cause | Preference |
|
||||
|---|---|---|
|
||||
| `experiment_creation` | the MLflow server is genuinely stopped | broken dependency |
|
||||
| `artifact_upload` | targeted patch on `log_artifacts` | patch, justified in place |
|
||||
| `registration` | targeted patch on `create_model_version` | patch, justified in place |
|
||||
| `model_document` | a real unique index refuses the activity's own insert | broken dependency |
|
||||
|
||||
Two departures from what the tasks anticipated, both recorded in the test that found them:
|
||||
|
||||
- **Registration needs a patch too**, not just artifact upload. One process serves both tracking and
|
||||
the registry, so stopping it fails experiment creation three steps earlier; and MLflow validates
|
||||
nothing about a registered model name except that it is non-empty, which gate 6 already requires.
|
||||
Both patches are on the single client call the step *is*, never on the activity or the workflow, so
|
||||
what is under test — the workflow attributing this step's failure to this step — is untouched.
|
||||
- **The model document's unique-index violation is on `model_config.target`**, not on `id`. The
|
||||
duplicate-`id` state the task describes cannot be constructed at all; `TestModelDocument`'s
|
||||
docstring shows why, and the mechanism the task wanted — a real index refusing a real insert — is
|
||||
kept intact.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import tempfile
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from temporalio.client import WorkflowFailureError
|
||||
|
||||
from e2e.import_harness import IMPORT_BUCKET, IMPORT_PREFIX, SEEDED_MODEL_IDS
|
||||
from laborious.utils.bundle.steps import ImportStep
|
||||
from laborious.utils.models.import_errors import IMPORT_FAILURE_CATALOG
|
||||
|
||||
pytestmark = [pytest.mark.import_e2e, pytest.mark.asyncio]
|
||||
|
||||
PLATFORM_ERROR_STATUS = 'IMPORT_ERROR'
|
||||
FACTORY_PASSWORD = 'Factory1Model!Test'
|
||||
|
||||
|
||||
def _as_dict(value: Any) -> dict[str, Any]:
|
||||
return json.loads(value) if isinstance(value, str) else (value or {})
|
||||
|
||||
|
||||
def _assert_platform_failure(
|
||||
*,
|
||||
row: dict[str, Any] | None,
|
||||
step: ImportStep,
|
||||
workflow_id: str,
|
||||
object_key: str,
|
||||
) -> None:
|
||||
"""The shared claim of task 7.5, asserted by every case in this module.
|
||||
|
||||
`error_message` is the catalog sentence *exactly* — not "contains" — which is the strongest
|
||||
possible form of "carries no path, table name, object key or traceback": there is nothing in it
|
||||
that did not come from the catalog.
|
||||
"""
|
||||
assert row is not None, 'the failure lost the row it was recorded on'
|
||||
assert row['status'] == PLATFORM_ERROR_STATUS, (
|
||||
f"status is {row['status']!r}; a provisioning failure is not the user's to fix"
|
||||
)
|
||||
|
||||
code, reason, _responsibility = IMPORT_FAILURE_CATALOG[step]
|
||||
detail = _as_dict(row['orchestrator_response_data'])['import']
|
||||
assert detail['step'] == step.value, detail
|
||||
assert detail['code'] == code.value, detail
|
||||
assert detail['workflow_id'] == workflow_id, detail
|
||||
assert row['error_message'] == reason
|
||||
|
||||
message = row['error_message'] or ''
|
||||
for forbidden in (
|
||||
object_key,
|
||||
IMPORT_BUCKET,
|
||||
IMPORT_PREFIX,
|
||||
'experiment_run',
|
||||
'Traceback',
|
||||
'Error',
|
||||
'mlflow',
|
||||
'Mongo',
|
||||
tempfile.gettempdir(),
|
||||
):
|
||||
assert forbidden not in message, f'{forbidden!r} leaked into error_message: {message!r}'
|
||||
|
||||
|
||||
async def _upload_and_start(
|
||||
*,
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
**factory_kwargs: Any,
|
||||
) -> tuple[Any, int, str]:
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory(**factory_kwargs))
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
return handle, import_run_id, object_key
|
||||
|
||||
|
||||
class TestExperimentCreation:
|
||||
"""The tracking server is not there. No patch, no mock — the process is stopped."""
|
||||
|
||||
async def test_a_stopped_tracking_server_fails_at_experiment_creation(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_mlflow_stopped: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
with import_mlflow_stopped():
|
||||
handle, import_run_id, object_key = await _upload_and_start(
|
||||
import_bundle_factory=import_bundle_factory,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError):
|
||||
await handle.result()
|
||||
row = import_read_status_row(import_run_id)
|
||||
|
||||
_assert_platform_failure(
|
||||
row=row,
|
||||
step=ImportStep.EXPERIMENT_CREATION,
|
||||
workflow_id=handle.id,
|
||||
object_key=object_key,
|
||||
)
|
||||
|
||||
# Download, decryption, experiment creation and artifact upload are one fused activity
|
||||
# attempt (`model-import-fix-bugs`, bug 8): the workflow only learns the bundle's names once
|
||||
# that whole attempt returns, and an experiment-creation failure means it never does. The
|
||||
# placeholder names inserted with the row are therefore still there — narrowed from the old
|
||||
# "names before MLflow is touched" invariant, per `model-import-status-record`'s updated spec.
|
||||
assert row is not None
|
||||
assert row['experiment_name'] == 'import-e2e-placeholder'
|
||||
|
||||
# Nothing in Mongo: the document is the last step and never ran.
|
||||
documents = list(import_models_collection.find({}, {'_id': False}))
|
||||
assert {str(d['id']) for d in documents} == set(SEEDED_MODEL_IDS)
|
||||
import_assert_nothing_provisioned(
|
||||
experiment_name=import_names['experiment_name'],
|
||||
model_name=import_names['model_name'],
|
||||
object_key=object_key,
|
||||
)
|
||||
|
||||
|
||||
class TestArtifactUpload:
|
||||
"""The step design E7 names as the one a patch is needed for."""
|
||||
|
||||
async def test_a_failing_artifact_upload_is_recorded_as_its_own_step(
|
||||
self,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_activities: Any,
|
||||
import_mlflow_server: dict[str, Any],
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
"""Why a patch: with the server up and the bundle valid, nothing else fails only this step.
|
||||
|
||||
Hostile input cannot reach here — a bundle that gets past gate 7 has a readable tree, and
|
||||
one that does not never reaches artifact upload. A broken dependency cannot isolate it
|
||||
either: the only dependency this step has is the tracking server, and stopping that fails
|
||||
experiment creation one step earlier. So the patch is on the single client call the step
|
||||
makes, `log_artifacts`; everything around it — the run creation, the parameters, the
|
||||
workflow's attribution of the failure — is the real code.
|
||||
"""
|
||||
client = import_activities.import_mlflow_repository.client
|
||||
|
||||
def refuse(*_args: Any, **_kwargs: Any) -> None:
|
||||
raise OSError('artifact store is unreachable')
|
||||
|
||||
monkeypatch.setattr(client, 'log_artifacts', refuse)
|
||||
|
||||
handle, import_run_id, object_key = await _upload_and_start(
|
||||
import_bundle_factory=import_bundle_factory,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError):
|
||||
await handle.result()
|
||||
|
||||
_assert_platform_failure(
|
||||
row=import_read_status_row(import_run_id),
|
||||
step=ImportStep.ARTIFACT_UPLOAD,
|
||||
workflow_id=handle.id,
|
||||
object_key=object_key,
|
||||
)
|
||||
|
||||
# The experiment exists — it was the previous step and it succeeded. Nothing after does.
|
||||
mlflow_client = _fresh_client(import_mlflow_server)
|
||||
assert mlflow_client.get_experiment_by_name(import_names['experiment_name']) is not None
|
||||
assert not list(mlflow_client.search_model_versions(f"name='{import_names['model_name']}'"))
|
||||
assert import_models_collection.find_one({'name': import_names['model_name']}) is None
|
||||
|
||||
|
||||
class TestRegistration:
|
||||
"""The artifacts are stored and the model still cannot be registered."""
|
||||
|
||||
async def test_a_failing_registration_is_recorded_as_its_own_step(
|
||||
self,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_activities: Any,
|
||||
import_mlflow_server: dict[str, Any],
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
"""Why a patch here too, contrary to design E7's expectation that this step needs none.
|
||||
|
||||
Hostile input was the first thing tried: a registered model name MLflow refuses. It does not
|
||||
exist — `_validate_model_name` rejects only the empty string, which gate 6 already refuses,
|
||||
and the sqlite backend store does not enforce the `VARCHAR(256)` the schema declares, so
|
||||
even a 300-character name is accepted. A broken dependency was the second: one process
|
||||
serves both tracking and the registry, so stopping it fails experiment creation three steps
|
||||
earlier. The patch is on `create_model_version`, the one call this step is.
|
||||
"""
|
||||
client = import_activities.import_mlflow_repository.client
|
||||
|
||||
def refuse(*_args: Any, **_kwargs: Any) -> None:
|
||||
raise RuntimeError('registry write rejected')
|
||||
|
||||
monkeypatch.setattr(client, 'create_model_version', refuse)
|
||||
|
||||
handle, import_run_id, object_key = await _upload_and_start(
|
||||
import_bundle_factory=import_bundle_factory,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError):
|
||||
await handle.result()
|
||||
|
||||
_assert_platform_failure(
|
||||
row=import_read_status_row(import_run_id),
|
||||
step=ImportStep.REGISTRATION,
|
||||
workflow_id=handle.id,
|
||||
object_key=object_key,
|
||||
)
|
||||
|
||||
# The run and its artifacts survive — "the model's files were stored, but the model itself
|
||||
# could not be registered" is exactly what the catalog sentence for this step says.
|
||||
mlflow_client = _fresh_client(import_mlflow_server)
|
||||
experiment = mlflow_client.get_experiment_by_name(import_names['experiment_name'])
|
||||
assert experiment is not None
|
||||
runs = mlflow_client.search_runs([experiment.experiment_id])
|
||||
assert len(runs) == 1
|
||||
assert {a.path for a in mlflow_client.list_artifacts(runs[0].info.run_id)} >= {
|
||||
'prediction_model',
|
||||
'data_model',
|
||||
}
|
||||
assert import_models_collection.find_one({'name': import_names['model_name']}) is None
|
||||
|
||||
|
||||
class TestModelDocument:
|
||||
"""The last step, failed by a real unique index refusing the real insert.
|
||||
|
||||
**Why the index is on `model_config.target` and not on `id`.** Task 7.4 asks for "a document
|
||||
already holding the id the activity is about to compute". That state cannot be constructed, and
|
||||
why is worth more than the test would have been: `_next_model_document_id` computes
|
||||
`max(int(id)) + 1` over `distinct('id')` of the *same* collection, so every value that
|
||||
stringifies to the colliding id is counted by that maximum. Put a document at `'8'` and the next
|
||||
id becomes `'9'`. Make the id unparseable and it is ignored by the numbering *and* does not
|
||||
collide. There is no third case — the numbering cannot collide with itself, and the only real
|
||||
duplicate-key path on `id` is two imports racing, which this change lists as a non-goal. That is
|
||||
a good property of the implementation, not a gap.
|
||||
|
||||
So the mechanism the task asked for is kept — a real unique index refusing the activity's own
|
||||
insert — on a key that can actually collide. `model_config.target` is `'SE'` on the seeded
|
||||
documents and `'SE'` on the one this import writes, so a unique index there makes the insert a
|
||||
genuine `DuplicateKeyError` from the server. No patch, no mock: the same class of failure, from
|
||||
the same enforcement mechanism, at the same step.
|
||||
"""
|
||||
|
||||
async def test_a_unique_index_violation_fails_at_model_document(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_mlflow_server: dict[str, Any],
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
# All three seeded documents carry target 'SE', so the index has to be buildable first:
|
||||
# give two of them their own target and leave exactly one holding 'SE'. That one is what the
|
||||
# import's own document — also target 'SE' — collides with.
|
||||
for ordinal, seeded_id in enumerate(sorted(SEEDED_MODEL_IDS)[1:], start=1):
|
||||
import_models_collection.update_one(
|
||||
{'id': seeded_id}, {'$set': {'model_config.target': f'OTHER_{ordinal}'}}
|
||||
)
|
||||
import_models_collection.create_index('model_config.target', unique=True, name='e2e_target')
|
||||
try:
|
||||
handle, import_run_id, object_key = await _upload_and_start(
|
||||
import_bundle_factory=import_bundle_factory,
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
parameters={'target_variable': 'SE'},
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError):
|
||||
await handle.result()
|
||||
row = import_read_status_row(import_run_id)
|
||||
finally:
|
||||
# The collection is dropped and reseeded per test, but dropping the index here keeps the
|
||||
# failure local even if that ever changes.
|
||||
import_models_collection.drop_index('e2e_target')
|
||||
|
||||
_assert_platform_failure(
|
||||
row=row,
|
||||
step=ImportStep.MODEL_DOCUMENT,
|
||||
workflow_id=handle.id,
|
||||
object_key=object_key,
|
||||
)
|
||||
|
||||
# No document for this model, and the seeded ones are exactly as they were.
|
||||
assert import_models_collection.find_one({'name': import_names['model_name']}) is None
|
||||
assert {
|
||||
str(document['id']) for document in import_models_collection.find({}, {'_id': False})
|
||||
} == set(SEEDED_MODEL_IDS)
|
||||
|
||||
# Everything before it stands — registered, and unusable, which is the point of the split
|
||||
# between this step's sentence and the registration step's.
|
||||
mlflow_client = _fresh_client(import_mlflow_server)
|
||||
versions = list(mlflow_client.search_model_versions(f"name='{import_names['model_name']}'"))
|
||||
assert len(versions) == 1
|
||||
|
||||
|
||||
def _fresh_client(server: dict[str, Any]) -> Any:
|
||||
"""An MlflowClient bound to the test server, independent of the one the activities hold."""
|
||||
import mlflow
|
||||
|
||||
return mlflow.tracking.MlflowClient(tracking_uri=server['tracking_uri'])
|
||||
@@ -1,514 +0,0 @@
|
||||
"""What happens when the *record* cannot be written — the cases that gate this feature's release.
|
||||
|
||||
Everything else in this suite asks "did the import do the right thing". This file asks the harder
|
||||
question: the import did the right thing, and then could not say so. Design D20 answers it in two
|
||||
halves, and both are release-critical.
|
||||
|
||||
**A database that comes back.** The terminal write retries, so a status database restarted during a
|
||||
deploy still records the correct verdict. Without this, every rolling restart would leave imports
|
||||
looking unfinished forever.
|
||||
|
||||
**A database that does not come back.** The model is fully provisioned, and the workflow must not
|
||||
claim otherwise. It **completes**, and the verdict it could not write goes out through the failure
|
||||
channel — log, metric, notification — under `IMPORT_STATUS_NOT_RECORDED`. An import that worked must
|
||||
never be reported as an import that failed just because the audit row was unreachable: somebody
|
||||
would re-import a model that is already there.
|
||||
|
||||
That last point was worth discovering the hard way. The first version of this file asserted the run
|
||||
*failed*; it does not, and it should not. The spec says so in as many words — "the run is not
|
||||
reported as failed" — and the implementation agrees with the spec.
|
||||
|
||||
And the case that is the evidence for a migration in another repository: **a table without `V12`**.
|
||||
The claim succeeds (`V11` already gave the table `run_type`, and `PENDING` and the running status are
|
||||
pre-`V12` values), the model is fully provisioned, and only the *terminal* value is new — so the last
|
||||
write of the import is the first thing that fails. `ImportStatusRejectedError` is non-retryable, so
|
||||
it lands on the first attempt with no wait, and the row is left running forever. That is what
|
||||
"apply `V12` before turning `IMPORT_MODEL_ENABLED` on" buys, as a test rather than a warning in a
|
||||
document.
|
||||
|
||||
Everything here runs on a **real clock** (design E9): under time skipping the backoff sleeps are
|
||||
virtual, so a restarted container would never get wall-clock time to come back and the recovery test
|
||||
would race by construction. The retry intervals are shortened by `import_status_retry_shortening`;
|
||||
the production values stay asserted by unit test 12.1 of `import-model-workflow`, because what these
|
||||
tests are about is the behaviour, not the arithmetic.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
import mlflow
|
||||
import pytest
|
||||
from temporalio.client import WorkflowFailureError
|
||||
|
||||
from e2e.import_harness import failure_text, failure_types
|
||||
|
||||
pytestmark = [pytest.mark.import_e2e, pytest.mark.slow, pytest.mark.asyncio]
|
||||
|
||||
FACTORY_PASSWORD = 'Factory1Model!Test'
|
||||
RUNNING_STATUS = 'ORCHESTRATOR_WAITING_PROC'
|
||||
NOT_RECORDED_CODE = 'IMPORT_STATUS_NOT_RECORDED'
|
||||
STATUS_WRITE_MARKER = 'step=status_write'
|
||||
PROVISIONING_STEPS = ('experiment_creation', 'artifact_upload', 'registration', 'model_document')
|
||||
|
||||
|
||||
def _as_dict(value: Any) -> dict[str, Any]:
|
||||
return json.loads(value) if isinstance(value, str) else (value or {})
|
||||
|
||||
|
||||
def _client(server: dict[str, Any]) -> Any:
|
||||
return mlflow.tracking.MlflowClient(tracking_uri=server['tracking_uri'])
|
||||
|
||||
|
||||
def _channel_text(mock_logger: Any) -> str:
|
||||
"""Everything the failure channel said, as one string.
|
||||
|
||||
The channel is a log line, a metric and a notification, and the log line is the only sink that
|
||||
carries the whole classification — code, marker and the verdict it could not record. Reading the
|
||||
logger is therefore reading the channel, and it is synchronous, which the notification path is
|
||||
not.
|
||||
"""
|
||||
return '\n'.join(str(call) for call in mock_logger.mock_calls)
|
||||
|
||||
|
||||
def _assert_fully_provisioned(
|
||||
*, server: dict[str, Any], names: dict[str, str], collection: Any
|
||||
) -> str:
|
||||
"""Everything the import was supposed to create exists. Returns the run id."""
|
||||
client = _client(server)
|
||||
experiment = client.get_experiment_by_name(names['experiment_name'])
|
||||
assert experiment is not None, 'no experiment: the import never got as far as provisioning'
|
||||
runs = client.search_runs([experiment.experiment_id])
|
||||
assert len(runs) == 1, f'expected one run, found {len(runs)}'
|
||||
assert {a.path for a in client.list_artifacts(runs[0].info.run_id)} >= {
|
||||
'prediction_model',
|
||||
'data_model',
|
||||
}
|
||||
versions = list(client.search_model_versions(f"name='{names['model_name']}'"))
|
||||
assert len(versions) == 1, f'expected one registered version, found {len(versions)}'
|
||||
assert collection.find_one({'name': names['model_name']}) is not None, 'no model document'
|
||||
return runs[0].info.run_id
|
||||
|
||||
|
||||
def _assert_no_provisioning_step_is_blamed(text_of_failure: str) -> None:
|
||||
"""The import's own steps all succeeded, so none of them may appear as the cause."""
|
||||
for step in PROVISIONING_STEPS:
|
||||
assert step not in text_of_failure, f'{step!r} is blamed, but that step succeeded'
|
||||
|
||||
|
||||
async def _wait_for_claim(read_row: Any, import_run_id: int, timeout: float = 30.0) -> None:
|
||||
"""Block until the workflow has taken ownership of the row.
|
||||
|
||||
The earliest deterministic point after which every status write belongs to this import. Used to
|
||||
place an outage that must cover the whole rest of the run.
|
||||
"""
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
row = read_row(import_run_id)
|
||||
if row is not None and row['status'] == RUNNING_STATUS:
|
||||
return
|
||||
await asyncio.sleep(0.1)
|
||||
raise AssertionError('the row was never claimed, so the outage could not be placed')
|
||||
|
||||
|
||||
async def _wait_for_provisioning(
|
||||
*, server: dict[str, Any], names: dict[str, str], collection: Any, timeout: float = 90.0
|
||||
) -> None:
|
||||
"""Block until the model document exists — the observable end of provisioning.
|
||||
|
||||
Polling the two destinations rather than patching an activity: "the document exists" *is* the
|
||||
end of provisioning, so there is nothing to hook and nothing to fake. This is what lets the
|
||||
outage below be placed exactly on the terminal write.
|
||||
"""
|
||||
client = _client(server)
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
document = collection.find_one({'name': names['model_name']})
|
||||
if document is not None and list(
|
||||
client.search_model_versions(f"name='{names['model_name']}'")
|
||||
):
|
||||
return
|
||||
await asyncio.sleep(0.2)
|
||||
raise AssertionError('provisioning never finished, so the outage could not be placed')
|
||||
|
||||
|
||||
class TestTheDatabaseComesBack:
|
||||
"""Stopped and restarted inside the retry envelope: the verdict still lands."""
|
||||
|
||||
async def test_a_restart_inside_the_envelope_still_records_the_verdict(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_mlflow_server: dict[str, Any],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute_real_clock: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_bff_stopped: Any,
|
||||
) -> None:
|
||||
"""The outage lands on the terminal write, and the server is back before the retries run out.
|
||||
|
||||
The stop is placed after provisioning rather than at the start, so what is being tested is
|
||||
the *terminal* write's retry envelope specifically. A stop at the start would also exercise
|
||||
the claim's envelope, which is a different (and already covered) code path.
|
||||
"""
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
handle = await import_execute_real_clock(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
|
||||
await _wait_for_provisioning(
|
||||
server=import_mlflow_server,
|
||||
names=import_names,
|
||||
collection=import_models_collection,
|
||||
)
|
||||
# Down for two seconds — several retry attempts, well inside the shortened envelope.
|
||||
async with import_bff_stopped():
|
||||
await asyncio.sleep(2)
|
||||
|
||||
result = await handle.result()
|
||||
assert result['model_name'] == import_names['model_name']
|
||||
|
||||
run_id = _assert_fully_provisioned(
|
||||
server=import_mlflow_server,
|
||||
names=import_names,
|
||||
collection=import_models_collection,
|
||||
)
|
||||
assert result['run_id'] == run_id
|
||||
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None
|
||||
assert row['status'] == 'IMPORT_SUCCESS', (
|
||||
f'the verdict did not land after the restart: the row says {row["status"]!r}'
|
||||
)
|
||||
assert row['error_message'] is None
|
||||
assert _as_dict(row['orchestrator_response_data'])['import']['workflow_id'] == handle.id
|
||||
|
||||
async def test_a_broken_progress_update_does_not_delay_the_verdict(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_mlflow_server: dict[str, Any],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute_real_clock: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_bff_stopped: Any,
|
||||
) -> None:
|
||||
"""A mid-flight outage that only the progress hints see. The import does not care.
|
||||
|
||||
Progress writes are hints nobody reads — the terminal write rewrites status, code, reason and
|
||||
step from one classification — so they get their own short policy. If they shared the terminal
|
||||
write's envelope, a database the import does not need would stall a working import for
|
||||
minutes at every stage update.
|
||||
|
||||
The outage is placed early, so it lands on hints rather than on the terminal write, and the
|
||||
assertion is that the import still finished with the correct verdict.
|
||||
"""
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
handle = await import_execute_real_clock(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
|
||||
# Early and brief: the claim has landed, the gates are running, the terminal write is far off.
|
||||
await asyncio.sleep(1.0)
|
||||
async with import_bff_stopped():
|
||||
await asyncio.sleep(1.5)
|
||||
|
||||
started = time.monotonic()
|
||||
result = await handle.result()
|
||||
elapsed = time.monotonic() - started
|
||||
|
||||
assert result['model_name'] == import_names['model_name']
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None
|
||||
assert row['status'] == 'IMPORT_SUCCESS'
|
||||
assert elapsed < 60, (
|
||||
f'the import took {elapsed:.1f}s after the outage ended; a lost hint should cost '
|
||||
f'its own short policy and nothing more'
|
||||
)
|
||||
|
||||
|
||||
class TestTheDatabaseStaysDown:
|
||||
"""Past the envelope. The model is real, the record is not, and the run says which is which."""
|
||||
|
||||
@pytest.fixture
|
||||
def import_status_retry_shortening(self) -> dict[str, Any]:
|
||||
"""Shorter still, because these tests wait for the envelope to *expire*.
|
||||
|
||||
Six half-second attempts inside an eight-second envelope: long enough that retries really
|
||||
happen, short enough that giving up is something a test suite can wait for.
|
||||
"""
|
||||
return {
|
||||
'initial_interval': 0.5,
|
||||
'backoff_coefficient': 1.0,
|
||||
'maximum_attempts': 6,
|
||||
'hint_maximum_attempts': 2,
|
||||
'envelope_seconds': 8,
|
||||
'write_timeout_seconds': 3,
|
||||
}
|
||||
|
||||
async def test_a_successful_import_is_not_reported_as_failed(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_mlflow_server: dict[str, Any],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute_real_clock: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_bff_stopped: Any,
|
||||
mock_logger: Any,
|
||||
) -> None:
|
||||
"""The most important assertion in this file.
|
||||
|
||||
The model is registered, the document is written, the artifacts are stored — and the row
|
||||
cannot be updated. The workflow **completes**: reporting failure here would tell somebody to
|
||||
re-import a model that is already on the platform. The verdict goes out through the failure
|
||||
channel instead, under the not-recorded code, carrying what it could not write.
|
||||
|
||||
The outage is placed as soon as the row has been **claimed**, not once provisioning has
|
||||
finished. Waiting for provisioning looks more precise and is in fact a race: the terminal
|
||||
write follows the model document by milliseconds, while `docker stop` takes seconds, so the
|
||||
write lands before the database is actually gone and the test proves nothing. Stopping right
|
||||
after the claim covers everything that follows — the progress hints (which are allowed to
|
||||
fail) and the terminal write (which is the subject) — with no timing to get right.
|
||||
"""
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
handle = await import_execute_real_clock(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
|
||||
await _wait_for_claim(import_read_status_row, import_run_id)
|
||||
async with import_bff_stopped():
|
||||
result = await handle.result()
|
||||
|
||||
# The run completed. This is the assertion, not an incidental detail.
|
||||
assert result['model_name'] == import_names['model_name']
|
||||
assert result['import_run_id'] == import_run_id
|
||||
|
||||
_assert_fully_provisioned(
|
||||
server=import_mlflow_server,
|
||||
names=import_names,
|
||||
collection=import_models_collection,
|
||||
)
|
||||
|
||||
channel = _channel_text(mock_logger)
|
||||
assert NOT_RECORDED_CODE in channel, (
|
||||
'the unrecordable verdict never reached the failure channel'
|
||||
)
|
||||
assert STATUS_WRITE_MARKER in channel, (
|
||||
'the channel does not mark this as a status-write failure'
|
||||
)
|
||||
assert 'verdict_not_recorded=' in channel, (
|
||||
'the channel reports the failure but not the verdict it could not write'
|
||||
)
|
||||
_assert_no_provisioning_step_is_blamed(
|
||||
'\n'.join(
|
||||
line for line in channel.splitlines() if 'model import failed' in line
|
||||
).replace('step=model_document', '')
|
||||
)
|
||||
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None
|
||||
assert row['status'] == RUNNING_STATUS, (
|
||||
f'the row should still be running, not {row["status"]!r}'
|
||||
)
|
||||
|
||||
async def test_a_failed_import_keeps_its_original_cause(
|
||||
self,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
import_status_table_pre_v12: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_activities: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute_real_clock: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
mock_logger: Any,
|
||||
) -> None:
|
||||
"""An import that already failed, whose failure then cannot be written down either.
|
||||
|
||||
Two failures, and only one of them answers "why did this import not work". The status write
|
||||
is the second, and it must not displace the first: somebody told "the status could not be
|
||||
recorded" would go and look at the database, when what actually happened is that their model's
|
||||
files could not be stored.
|
||||
|
||||
**Why a pre-`V12` table instead of stopping the database.** The outage has to cover the
|
||||
terminal write and nothing earlier, and an import that fails at a gate reaches its terminal
|
||||
write within a second of starting — there is no window to aim at, and the first version of
|
||||
this test proved only that a stopped database *after* a completed write changes nothing.
|
||||
A pre-`V12` table has no timing at all: `IMPORT_ERROR` is a value it has never heard of, so
|
||||
the terminal write of a *failed* import is refused deterministically while every write before
|
||||
it succeeds. The failure being tested is the same one either way — the record cannot take the
|
||||
verdict — and this version cannot race.
|
||||
"""
|
||||
client = import_activities.import_mlflow_repository.client
|
||||
|
||||
def refuse(*_args: Any, **_kwargs: Any) -> None:
|
||||
raise OSError('artifact store is unreachable')
|
||||
|
||||
monkeypatch.setattr(client, 'log_artifacts', refuse)
|
||||
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
handle = await import_execute_real_clock(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError) as failure:
|
||||
await handle.result()
|
||||
|
||||
# The reported cause is the step that actually failed.
|
||||
text_of_failure = failure_text(failure.value)
|
||||
assert 'artifact store is unreachable' in text_of_failure, (
|
||||
f'the reported cause lost the original failure: {text_of_failure}'
|
||||
)
|
||||
assert 'ImportStatusRejectedError' not in failure_types(failure.value), (
|
||||
f'the status write displaced the original cause: {failure_types(failure.value)}'
|
||||
)
|
||||
assert 'the import log database refused' not in text_of_failure, (
|
||||
'the run is reported as a status-write failure rather than an artifact-upload failure'
|
||||
)
|
||||
|
||||
# The status write still reported *itself*, through the channel, as the secondary failure.
|
||||
channel = _channel_text(mock_logger)
|
||||
assert NOT_RECORDED_CODE in channel, 'the unrecordable verdict was never reported'
|
||||
assert STATUS_WRITE_MARKER in channel
|
||||
|
||||
# And the row is left running, because the verdict never landed.
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None
|
||||
assert row['status'] == RUNNING_STATUS
|
||||
|
||||
|
||||
class TestMissingV12:
|
||||
"""A table without `V12`. The evidence for why the migration gates enabling the worker."""
|
||||
|
||||
async def test_the_terminal_write_is_refused_on_its_first_attempt(
|
||||
self,
|
||||
import_status_table_pre_v12: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_mlflow_server: dict[str, Any],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute_real_clock: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_bff_engine: Any,
|
||||
mock_logger: Any,
|
||||
) -> None:
|
||||
"""Provisioned, then unrecordable — and unrecordable *immediately*.
|
||||
|
||||
The shape matters as much as the outcome. `ImportStatusRejectedError` is non-retryable: a
|
||||
`CHECK` violation is not going to start passing because it was tried again, and waiting the
|
||||
full 151-second envelope would turn a misconfiguration into a two-and-a-half-minute stall on
|
||||
every single import. The elapsed time is what proves it — the shortened policy still allows
|
||||
24 attempts, so a retried failure could not finish this fast.
|
||||
"""
|
||||
from sqlalchemy import text as sql_text
|
||||
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
|
||||
started = time.monotonic()
|
||||
handle = await import_execute_real_clock(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
result = await handle.result()
|
||||
elapsed = time.monotonic() - started
|
||||
|
||||
# The run completed: an import that provisioned a model is not a failed import, even when
|
||||
# the last thing it wanted to do was impossible.
|
||||
assert result['model_name'] == import_names['model_name']
|
||||
|
||||
# The claim and the names write both landed — the pre-V12 table was usable right up to the
|
||||
# last statement, which is exactly what makes this failure mode easy to miss.
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None
|
||||
assert row['status'] == RUNNING_STATUS, (
|
||||
f'the row is {row["status"]!r}; a pre-V12 table cannot hold a terminal import status, '
|
||||
f'so the row is left looking as though the import never finished'
|
||||
)
|
||||
assert row['experiment_name'] == import_names['experiment_name']
|
||||
assert row['run_name'] == f'import-{import_names["model_name"]}-v2'
|
||||
|
||||
# The model is fully provisioned. This is the whole problem: the work is done and invisible.
|
||||
_assert_fully_provisioned(
|
||||
server=import_mlflow_server,
|
||||
names=import_names,
|
||||
collection=import_models_collection,
|
||||
)
|
||||
|
||||
# The verdict reached the failure channel under the not-recorded code.
|
||||
channel = _channel_text(mock_logger)
|
||||
assert NOT_RECORDED_CODE in channel
|
||||
assert STATUS_WRITE_MARKER in channel
|
||||
assert 'verdict_not_recorded=' in channel
|
||||
|
||||
# Non-retryable: refused on the first attempt, with no retry wait.
|
||||
assert elapsed < 60, (
|
||||
f'the import took {elapsed:.1f}s; a non-retryable rejection must not be retried'
|
||||
)
|
||||
|
||||
# And the constraint that refused it is the one V12 adds — asserted directly, so the test
|
||||
# says *why* the write was impossible rather than only that it was.
|
||||
with import_bff_engine.begin() as connection:
|
||||
statuses = connection.execute(
|
||||
sql_text(
|
||||
'SELECT pg_get_constraintdef(oid) FROM pg_constraint '
|
||||
"WHERE conname = 'chk_experiment_run_status'"
|
||||
)
|
||||
).scalar_one()
|
||||
assert 'IMPORT_SUCCESS' not in statuses, (
|
||||
'the pre-V12 fixture already allows the import statuses, so this proved nothing'
|
||||
)
|
||||
@@ -1,272 +0,0 @@
|
||||
"""Importing the same model twice, and what the model document is allowed to say.
|
||||
|
||||
Two claims, and they pull in opposite directions, which is why they share a file.
|
||||
|
||||
**A second import is additive.** It creates a new run in the *same* experiment and a new registry
|
||||
version of the *same* registered model, and it promotes nothing. That is the whole of it: no
|
||||
overwrite, no stage change, no reconciliation.
|
||||
|
||||
**A second import touches the document not at all.** `active`, `target` and the flavors of a model
|
||||
already on the platform are operator-owned — somebody turned that model on, and re-importing it must
|
||||
not turn it off, retarget it or reset how it is loaded. So the document is written on the first
|
||||
import and left byte-identical on every one after it.
|
||||
|
||||
The document's content is asserted exactly rather than loosely, because it *is* the model listing:
|
||||
`model_config.transform_flavor`, `model_config.predict_flavor` and `model_config.target` are what
|
||||
the runtime reads to know how to load the model, and a missing or wrong one makes an imported model
|
||||
either invisible or unloadable.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import mlflow
|
||||
import pytest
|
||||
|
||||
from e2e.import_harness import SEEDED_MODEL_IDS
|
||||
|
||||
pytestmark = [pytest.mark.import_e2e, pytest.mark.asyncio]
|
||||
|
||||
FACTORY_PASSWORD = 'Factory1Model!Test'
|
||||
CONTRACTED_MODEL_CONFIG_KEYS = {
|
||||
'transform_flavor',
|
||||
'predict_flavor',
|
||||
'retention_minutes',
|
||||
'target',
|
||||
'skip_predict_timestamp',
|
||||
}
|
||||
|
||||
|
||||
def _client(server: dict[str, Any]) -> Any:
|
||||
return mlflow.tracking.MlflowClient(tracking_uri=server['tracking_uri'])
|
||||
|
||||
|
||||
async def _import_once(
|
||||
*,
|
||||
payload: bytes,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
) -> dict[str, Any]:
|
||||
"""One successful import. Raises if it did not succeed, so a caller never asserts on a failure."""
|
||||
object_key, digest = import_upload_bundle(payload)
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
result = await handle.result()
|
||||
return {'result': result, 'import_run_id': import_run_id, 'object_key': object_key}
|
||||
|
||||
|
||||
class TestSecondImport:
|
||||
"""The same model, imported twice in one test — the only way to observe the second as second."""
|
||||
|
||||
async def test_a_second_import_adds_a_version_and_a_run_and_nothing_else(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_mlflow_server: dict[str, Any],
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
"""Both bundles declare the same model and experiment, and different `model_version`s.
|
||||
|
||||
The second bundle declaring `model_version` '7' while the first declared '2' is the point of
|
||||
task 8.2: the destination version number is the registry's own count, and the bundle's
|
||||
number is a record of where the model came from, logged as `source.model_version` and used
|
||||
for nothing else. If the two were confused, the second version here would be '7'.
|
||||
"""
|
||||
first = await _import_once(
|
||||
payload=import_bundle_factory(metadata={'model_version': '2'}),
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
document_after_first = import_models_collection.find_one(
|
||||
{'name': import_names['model_name']}, {'_id': False}
|
||||
)
|
||||
assert document_after_first is not None, 'the first import wrote no document'
|
||||
|
||||
second = await _import_once(
|
||||
payload=import_bundle_factory(metadata={'model_version': '7'}),
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
|
||||
model_name = import_names['model_name']
|
||||
client = _client(import_mlflow_server)
|
||||
|
||||
# 8.1 — a new version of the same registered model.
|
||||
assert first['result']['version'] == '1'
|
||||
assert second['result']['version'] == '2'
|
||||
versions = list(client.search_model_versions(f"name='{model_name}'"))
|
||||
assert len(versions) == 2, f'expected two versions, found {len(versions)}'
|
||||
assert {v.version for v in versions} == {'1', '2'}
|
||||
|
||||
# 8.2 — destination numbering is independent of what the bundles declared.
|
||||
assert second['result']['version'] != '7'
|
||||
runs_by_id = {
|
||||
run.info.run_id: run
|
||||
for run in client.search_runs(
|
||||
[client.get_experiment_by_name(import_names['experiment_name']).experiment_id]
|
||||
)
|
||||
}
|
||||
assert runs_by_id[first['result']['run_id']].data.params['source.model_version'] == '2'
|
||||
assert runs_by_id[second['result']['run_id']].data.params['source.model_version'] == '7'
|
||||
|
||||
# 8.1 — a new run, in the *same* experiment, and both runs are still there.
|
||||
assert first['result']['run_id'] != second['result']['run_id']
|
||||
assert len(runs_by_id) == 2
|
||||
assert {run.info.run_name for run in runs_by_id.values()} == {
|
||||
f'import-{model_name}-v2',
|
||||
f'import-{model_name}-v7',
|
||||
}
|
||||
|
||||
# 8.1 — nothing promoted, by either import.
|
||||
assert all(version.current_stage in (None, 'None') for version in versions), [
|
||||
(v.version, v.current_stage) for v in versions
|
||||
]
|
||||
|
||||
# 8.3 — the document is byte-identical to what the first import wrote.
|
||||
document_after_second = import_models_collection.find_one(
|
||||
{'name': model_name}, {'_id': False}
|
||||
)
|
||||
assert document_after_second == document_after_first
|
||||
|
||||
# And both rows record their own success.
|
||||
for run in (first, second):
|
||||
row = import_read_status_row(run['import_run_id'])
|
||||
assert row is not None
|
||||
assert row['status'] == 'IMPORT_SUCCESS'
|
||||
|
||||
|
||||
class TestModelDocumentContent:
|
||||
"""What the written document says, asserted exactly rather than by presence."""
|
||||
|
||||
async def test_the_document_continues_the_collection_and_carries_the_contract(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
) -> None:
|
||||
"""`id` is `max(int(id)) + 1` across gapped ids — 1, 2, 7 give 8, not 4 and not "3".
|
||||
|
||||
The gaps are what make this worth asserting: a count would give 4, and a lexicographic
|
||||
maximum would give "8" here but "10" for a collection holding "9" and "30". The seeded ids
|
||||
are strings, as the baseline migration stores them, so `int()` is doing real work.
|
||||
"""
|
||||
await _import_once(
|
||||
payload=import_bundle_factory(parameters={'target_variable': 'PRODUCT_QUALITY'}),
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
|
||||
document = import_models_collection.find_one(
|
||||
{'name': import_names['model_name']}, {'_id': False}
|
||||
)
|
||||
assert document is not None
|
||||
|
||||
expected_id = str(max(int(seeded) for seeded in SEEDED_MODEL_IDS) + 1)
|
||||
assert document['id'] == expected_id, (
|
||||
f'{document["id"]!r} does not continue {SEEDED_MODEL_IDS}'
|
||||
)
|
||||
assert document['id'] == '8'
|
||||
|
||||
assert document['name'] == import_names['model_name']
|
||||
assert document['active'] is False
|
||||
|
||||
assert set(document['model_config']) == CONTRACTED_MODEL_CONFIG_KEYS, (
|
||||
f'model_config holds {sorted(document["model_config"])}, not exactly the '
|
||||
f'contracted keys'
|
||||
)
|
||||
assert document['model_config'] == {
|
||||
'transform_flavor': 'sklearn',
|
||||
'predict_flavor': 'joblib',
|
||||
'retention_minutes': 60,
|
||||
'target': 'PRODUCT_QUALITY',
|
||||
'skip_predict_timestamp': True,
|
||||
}
|
||||
|
||||
# The document has exactly these fields and no others — nothing was smuggled in alongside.
|
||||
assert set(document) == {'id', 'name', 'active', 'model_config'}
|
||||
|
||||
async def test_a_name_already_in_the_collection_is_left_untouched(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
"""An operator-owned document with this model's name survives the import unchanged.
|
||||
|
||||
`active: True` and a hand-set target are exactly what a re-import must not undo: somebody
|
||||
turned this model on and pointed it at something. The import notices the name and stops.
|
||||
"""
|
||||
existing = {
|
||||
'id': '42',
|
||||
'name': import_names['model_name'],
|
||||
'active': True,
|
||||
'model_config': {
|
||||
'transform_flavor': 'sklearn',
|
||||
'predict_flavor': 'pyfunc',
|
||||
'retention_minutes': 5,
|
||||
'target': 'OPERATOR_CHOSE_THIS',
|
||||
},
|
||||
}
|
||||
import_models_collection.insert_one(dict(existing))
|
||||
before = list(import_models_collection.find({}, {'_id': False}))
|
||||
|
||||
outcome = await _import_once(
|
||||
payload=import_bundle_factory(parameters={'target_variable': 'SOMETHING_ELSE'}),
|
||||
import_upload_bundle=import_upload_bundle,
|
||||
import_insert_status_row=import_insert_status_row,
|
||||
import_build_input=import_build_input,
|
||||
import_execute=import_execute,
|
||||
import_envelope=import_envelope,
|
||||
)
|
||||
|
||||
# The import still succeeded — a document that is already there is not a failure.
|
||||
row = import_read_status_row(outcome['import_run_id'])
|
||||
assert row is not None
|
||||
assert row['status'] == 'IMPORT_SUCCESS'
|
||||
|
||||
after = list(import_models_collection.find({}, {'_id': False}))
|
||||
assert after == before, 'the collection changed'
|
||||
|
||||
matching = [d for d in after if d['name'] == import_names['model_name']]
|
||||
assert len(matching) == 1, 'a duplicate was inserted'
|
||||
assert matching[0] == existing
|
||||
@@ -1,602 +0,0 @@
|
||||
"""The import log row: who owns it, when, and which columns the import is allowed to touch.
|
||||
|
||||
This is the file where the *real* `CHECK` constraints matter. Everything here would pass against a
|
||||
table without `V12` or without the length and pattern checks, and none of it would mean anything.
|
||||
|
||||
Four claims:
|
||||
|
||||
- **Ordering.** The claim is the workflow's first act, before any gate runs, and the names are on
|
||||
the row while the import is still in flight — because a person watching the import list has to see
|
||||
*which* model is being imported while it is being imported.
|
||||
- **Ownership.** The columns the inserting side owns — `username`, `request_data`, `bucket_name`,
|
||||
`file_name`, `run_type`, `created_at` — are never rewritten. The import updates status, the
|
||||
reason, the two names and the JSON detail, and nothing else.
|
||||
- **Refusal.** A row this workflow may not own is not written to at all. Not with an error, not with
|
||||
a status: the loudness goes to the Temporal error, the log and the metric, because a row that
|
||||
belongs to somebody else is not the place to be loud.
|
||||
- **Conformance.** Every status value an import writes is one the real constraint accepts, a
|
||||
training row's vocabulary still works beside it, and the feature added no column.
|
||||
|
||||
Two observations here are made from *inside* a running import, by wrapping the first call the step
|
||||
under observation makes. That is a patch, and it is an observer rather than a failure injector: it
|
||||
snapshots the row and delegates. Ordering cannot be asserted from a finished run — the terminal
|
||||
write is precisely what erases the intermediate state — so the assertion has to stand somewhere in
|
||||
the middle.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
from temporalio.client import WorkflowFailureError
|
||||
|
||||
from e2e.import_harness import IMPORT_PREFIX, failure_types
|
||||
from laborious.utils.bundle.steps import ImportStep
|
||||
from laborious.utils.models.import_errors import IMPORT_FAILURE_CATALOG
|
||||
|
||||
# `asyncio` is applied per class rather than to the whole module: `TestConstraintConformance`
|
||||
# talks to Postgres directly and its tests are synchronous, and a module-wide `asyncio` mark on
|
||||
# a sync test is a pytest warning, not a no-op.
|
||||
pytestmark = [pytest.mark.import_e2e]
|
||||
|
||||
FACTORY_PASSWORD = 'Factory1Model!Test'
|
||||
RUNNING_STATUS = 'ORCHESTRATOR_WAITING_PROC'
|
||||
VALIDATION_ERROR_STATUS = 'ORCHESTRATOR_VALIDATION_ERROR'
|
||||
|
||||
# The columns the inserting side owns. The import may never rewrite one of these.
|
||||
INSERTER_OWNED_COLUMNS = (
|
||||
'username',
|
||||
'request_data',
|
||||
'bucket_name',
|
||||
'file_name',
|
||||
'run_type',
|
||||
'created_at',
|
||||
)
|
||||
|
||||
# The 13 columns V1…V12 leave. D21 says the feature adds none; this is what checks it.
|
||||
EXPECTED_COLUMNS = {
|
||||
'id',
|
||||
'experiment_name',
|
||||
'run_name',
|
||||
'username',
|
||||
'status',
|
||||
'error_message',
|
||||
'created_at',
|
||||
'updated_at',
|
||||
'bucket_name',
|
||||
'file_name',
|
||||
'request_data',
|
||||
'orchestrator_response_data',
|
||||
'run_type',
|
||||
}
|
||||
|
||||
|
||||
def _as_dict(value: Any) -> dict[str, Any]:
|
||||
return json.loads(value) if isinstance(value, str) else (value or {})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestOrdering:
|
||||
"""When each write happens, observed from inside the run rather than after it."""
|
||||
|
||||
async def test_the_row_is_claimed_before_the_first_gate_runs(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_bundle_factory: Any,
|
||||
import_activities: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
"""Snapshot the row from inside `stat_object` — the first call gate 1 makes.
|
||||
|
||||
Standing here rather than after the run is not a convenience: the terminal write rewrites
|
||||
status, reason and detail, so a finished row cannot say what it looked like before the gates.
|
||||
`stat_object` is patched with a wrapper that records and then delegates — an observer, not a
|
||||
failure. Everything the import does still really happens.
|
||||
"""
|
||||
client = import_activities.import_minio_repository.minio_client
|
||||
real_stat = client.stat_object
|
||||
snapshots: list[dict[str, Any] | None] = []
|
||||
|
||||
def observing_stat(bucket: str, key: str, *args: Any, **kwargs: Any) -> Any:
|
||||
snapshots.append(import_read_status_row(observed_run_id[0]))
|
||||
return real_stat(bucket, key, *args, **kwargs)
|
||||
|
||||
observed_run_id: list[int] = [0]
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
observed_run_id[0] = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
client.stat_object = observing_stat
|
||||
try:
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=observed_run_id[0],
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
await handle.result()
|
||||
finally:
|
||||
client.stat_object = real_stat
|
||||
|
||||
assert snapshots, 'gate 1 never ran, so nothing was observed'
|
||||
claimed = snapshots[0]
|
||||
assert claimed is not None
|
||||
assert claimed['status'] == RUNNING_STATUS, (
|
||||
f'at the first gate the row was {claimed["status"]!r}, not the running status'
|
||||
)
|
||||
detail = _as_dict(claimed['orchestrator_response_data'])
|
||||
assert detail['import']['workflow_id'] == handle.id
|
||||
assert detail['import']['step'] in (
|
||||
ImportStep.RECEIVED.value,
|
||||
ImportStep.DOWNLOAD.value,
|
||||
), detail
|
||||
|
||||
async def test_the_names_are_on_the_row_while_the_import_is_still_running(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_activities: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
"""Observed at the first call registration makes, so the experiment and run already exist.
|
||||
|
||||
Download, decryption, experiment creation and artifact upload are fused into one activity
|
||||
attempt (`model-import-status-record`'s design, decision on the fused step): the names are
|
||||
not known to the workflow, and so cannot be recorded, until that whole activity returns —
|
||||
the guarantee is anchored to the two steps that follow it, registration and the model
|
||||
document, not to the experiment or run, which by construction already exist by then.
|
||||
|
||||
This is the stronger form of the same claim task 7.1 makes from a failure: here the import
|
||||
*succeeds*, and the names were still already written before registration ran — so their
|
||||
presence is not an artefact of the run having stopped.
|
||||
"""
|
||||
client = import_activities.import_mlflow_repository.client
|
||||
real_get_registered_model = client.get_registered_model
|
||||
snapshots: list[dict[str, Any] | None] = []
|
||||
observed_run_id: list[int] = [0]
|
||||
|
||||
def observing_get_registered_model(*args: Any, **kwargs: Any) -> Any:
|
||||
snapshots.append(import_read_status_row(observed_run_id[0]))
|
||||
return real_get_registered_model(*args, **kwargs)
|
||||
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
observed_run_id[0] = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
client.get_registered_model = observing_get_registered_model
|
||||
try:
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=observed_run_id[0],
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
await handle.result()
|
||||
finally:
|
||||
client.get_registered_model = real_get_registered_model
|
||||
|
||||
assert snapshots, 'registration never ran'
|
||||
mid_flight = snapshots[0]
|
||||
assert mid_flight is not None
|
||||
assert mid_flight['status'] == RUNNING_STATUS
|
||||
assert mid_flight['experiment_name'] == import_names['experiment_name']
|
||||
assert mid_flight['run_name'] == f'import-{import_names["model_name"]}-v2'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestOwnership:
|
||||
"""The inserting side's columns survive an import untouched, whatever the outcome."""
|
||||
|
||||
@pytest.mark.parametrize('outcome', ['success', 'failure'])
|
||||
async def test_the_inserted_columns_are_never_rewritten(
|
||||
self,
|
||||
outcome: str,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
payload = (
|
||||
import_bundle_factory()
|
||||
if outcome == 'success'
|
||||
else import_bundle_factory(include_metadata=False)
|
||||
)
|
||||
object_key, digest = import_upload_bundle(payload)
|
||||
import_run_id = import_insert_status_row(
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
username='owner.of.the.row@example.com',
|
||||
)
|
||||
|
||||
# Snapshotted before the run, not after: `import_execute` waits for the workflow, so a
|
||||
# snapshot taken from its return value would already be the post-import row and every
|
||||
# comparison below would be a row against itself.
|
||||
before = import_read_status_row(import_run_id)
|
||||
assert before is not None
|
||||
assert before['status'] == 'PENDING'
|
||||
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
|
||||
if outcome == 'success':
|
||||
await handle.result()
|
||||
else:
|
||||
with pytest.raises(WorkflowFailureError):
|
||||
await handle.result()
|
||||
|
||||
after = import_read_status_row(import_run_id)
|
||||
assert after is not None
|
||||
for column in INSERTER_OWNED_COLUMNS:
|
||||
assert after[column] == before[column], (
|
||||
f'{column} was rewritten: {before[column]!r} -> {after[column]!r}'
|
||||
)
|
||||
|
||||
# And the import *did* write what it owns, so this is not passing by doing nothing at all.
|
||||
assert after['status'] != before['status']
|
||||
assert after['updated_at'] >= before['updated_at']
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestClaimGuards:
|
||||
"""A row this workflow may not own is not written to. At all."""
|
||||
|
||||
async def test_an_unknown_id_fails_the_run_and_writes_nothing(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_bff_engine: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
# A real, claimable row exists — so "nothing was written" is a claim about *this* run
|
||||
# choosing the right row, not about the table being empty.
|
||||
bystander_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
before = _snapshot_table(import_bff_engine)
|
||||
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=bystander_id + 10_000,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError) as failure:
|
||||
await handle.result()
|
||||
|
||||
assert 'ImportLogRowNotClaimableError' in failure_types(failure.value)
|
||||
assert _snapshot_table(import_bff_engine) == before, 'the table changed'
|
||||
import_assert_nothing_provisioned(
|
||||
experiment_name=import_names['experiment_name'],
|
||||
model_name=import_names['model_name'],
|
||||
object_key=object_key,
|
||||
)
|
||||
|
||||
async def test_a_training_row_cannot_be_claimed_by_an_import(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_bff_engine: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
"""`run_type='TRAIN'` is a different feature's row, and the claim will not take it."""
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
training_id = import_insert_status_row(
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
run_type='TRAIN',
|
||||
experiment_name='a-training-experiment',
|
||||
)
|
||||
before = _snapshot_table(import_bff_engine)
|
||||
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=training_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError) as failure:
|
||||
await handle.result()
|
||||
|
||||
assert 'ImportLogRowNotClaimableError' in failure_types(failure.value)
|
||||
assert _snapshot_table(import_bff_engine) == before
|
||||
import_assert_nothing_provisioned(
|
||||
experiment_name=import_names['experiment_name'],
|
||||
model_name=import_names['model_name'],
|
||||
object_key=object_key,
|
||||
)
|
||||
|
||||
async def test_a_row_held_by_another_workflow_cannot_be_taken(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_bff_engine: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
"""Already running under somebody else's workflow id.
|
||||
|
||||
The claim is idempotent for *its own* workflow id and refuses every other, which is what
|
||||
makes a duplicated workflow start safe and a stolen row impossible.
|
||||
"""
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
held_id = import_insert_status_row(
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
status=RUNNING_STATUS,
|
||||
orchestrator_response_data={
|
||||
'import': {'workflow_id': 'import-held-by-somebody-else', 'step': 'received'}
|
||||
},
|
||||
)
|
||||
before = _snapshot_table(import_bff_engine)
|
||||
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=held_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError) as failure:
|
||||
await handle.result()
|
||||
|
||||
assert 'ImportLogRowNotClaimableError' in failure_types(failure.value)
|
||||
assert _snapshot_table(import_bff_engine) == before
|
||||
import_assert_nothing_provisioned(
|
||||
experiment_name=import_names['experiment_name'],
|
||||
model_name=import_names['model_name'],
|
||||
object_key=object_key,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestInputAgainstTheRow:
|
||||
"""A request that contradicts the row it claimed — recorded *on* that row, because it owns it."""
|
||||
|
||||
async def test_a_file_name_disagreeing_with_the_object_key_is_recorded_on_the_row(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_names: dict[str, str],
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
import_assert_nothing_provisioned: Any,
|
||||
) -> None:
|
||||
"""The difference from the claim guards above is where the error goes.
|
||||
|
||||
A row that could not be claimed is somebody else's, so nothing is written to it. This row
|
||||
*was* claimed — the contradiction is only found afterwards — so the failure belongs on it,
|
||||
where the person who submitted the import will see it.
|
||||
"""
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
import_run_id = import_insert_status_row(
|
||||
object_key=f'{IMPORT_PREFIX}a-different-object.sientia',
|
||||
expected_digest=digest,
|
||||
)
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
with pytest.raises(WorkflowFailureError):
|
||||
await handle.result()
|
||||
|
||||
row = import_read_status_row(import_run_id)
|
||||
assert row is not None
|
||||
code, reason, _responsibility = IMPORT_FAILURE_CATALOG[ImportStep.RECEIVED]
|
||||
assert row['status'] == VALIDATION_ERROR_STATUS
|
||||
assert row['error_message'] == reason
|
||||
detail = _as_dict(row['orchestrator_response_data'])['import']
|
||||
assert detail['step'] == ImportStep.RECEIVED.value
|
||||
assert detail['code'] == code.value
|
||||
# The record still names the object *it* was given, not the one the request named.
|
||||
assert row['file_name'] == f'{IMPORT_PREFIX}a-different-object.sientia'
|
||||
import_assert_nothing_provisioned(
|
||||
experiment_name=import_names['experiment_name'],
|
||||
model_name=import_names['model_name'],
|
||||
object_key=object_key,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestTrainingRowsAreUntouched:
|
||||
"""A training row sitting in the same table is never modified by an import."""
|
||||
|
||||
async def test_a_training_row_beside_an_import_is_unchanged(
|
||||
self,
|
||||
import_status_table: Any,
|
||||
import_models_collection: Any,
|
||||
import_bundle_factory: Any,
|
||||
import_upload_bundle: Any,
|
||||
import_insert_status_row: Any,
|
||||
import_build_input: Any,
|
||||
import_execute: Any,
|
||||
import_envelope: Any,
|
||||
import_read_status_row: Any,
|
||||
) -> None:
|
||||
object_key, digest = import_upload_bundle(import_bundle_factory())
|
||||
training_id = import_insert_status_row(
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
run_type='TRAIN',
|
||||
status='TRAINING_SUCCESS',
|
||||
experiment_name='someone-elses-training',
|
||||
username='trainer@example.com',
|
||||
)
|
||||
training_before = import_read_status_row(training_id)
|
||||
|
||||
import_run_id = import_insert_status_row(object_key=object_key, expected_digest=digest)
|
||||
handle = await import_execute(
|
||||
import_build_input(
|
||||
import_run_id=import_run_id,
|
||||
object_key=object_key,
|
||||
expected_digest=digest,
|
||||
password_envelope=import_envelope(FACTORY_PASSWORD),
|
||||
)
|
||||
)
|
||||
await handle.result()
|
||||
|
||||
assert import_read_status_row(import_run_id)['status'] == 'IMPORT_SUCCESS'
|
||||
assert import_read_status_row(training_id) == training_before
|
||||
|
||||
|
||||
class TestConstraintConformance:
|
||||
"""The real constraints, exercised directly — and the column list D21 promised not to change."""
|
||||
|
||||
def test_every_status_an_import_writes_is_accepted_on_an_import_row(
|
||||
self, import_status_table: Any, import_insert_status_row: Any, import_bff_engine: Any
|
||||
) -> None:
|
||||
"""`PENDING`, the running value, and both import terminals, on `run_type='IMPORT'`."""
|
||||
row_id = import_insert_status_row(
|
||||
object_key=f'{IMPORT_PREFIX}conformance.sientia', expected_digest='a' * 64
|
||||
)
|
||||
for status in ('PENDING', RUNNING_STATUS, VALIDATION_ERROR_STATUS, 'IMPORT_SUCCESS'):
|
||||
_set_status(import_bff_engine, row_id, status)
|
||||
_set_status(import_bff_engine, row_id, 'IMPORT_ERROR')
|
||||
|
||||
def test_the_terminal_values_do_not_cross_run_types(
|
||||
self, import_status_table: Any, import_insert_status_row: Any, import_bff_engine: Any
|
||||
) -> None:
|
||||
"""`chk_experiment_run_terminal_status_run_type`, in both directions.
|
||||
|
||||
This is the constraint `V12` adds, and it is the reason the migration gates enabling the
|
||||
worker: without it an import could write `TRAINING_SUCCESS` and a training run could write
|
||||
`IMPORT_ERROR`, and the import list would be unreadable.
|
||||
"""
|
||||
import_row = import_insert_status_row(
|
||||
object_key=f'{IMPORT_PREFIX}crossing.sientia', expected_digest='b' * 64
|
||||
)
|
||||
training_row = import_insert_status_row(
|
||||
object_key=f'{IMPORT_PREFIX}crossing.sientia',
|
||||
expected_digest='b' * 64,
|
||||
run_type='TRAIN',
|
||||
)
|
||||
|
||||
for row_id, forbidden in (
|
||||
(import_row, 'TRAINING_SUCCESS'),
|
||||
(import_row, 'TRAINING_ERROR'),
|
||||
(training_row, 'IMPORT_SUCCESS'),
|
||||
(training_row, 'IMPORT_ERROR'),
|
||||
):
|
||||
with pytest.raises(IntegrityError) as refusal:
|
||||
_set_status(import_bff_engine, row_id, forbidden)
|
||||
assert 'chk_experiment_run_terminal_status_run_type' in str(refusal.value)
|
||||
|
||||
# The shared stage values are deliberately outside the pairing: both run types pass through.
|
||||
for row_id in (import_row, training_row):
|
||||
for shared in ('PENDING', RUNNING_STATUS, VALIDATION_ERROR_STATUS):
|
||||
_set_status(import_bff_engine, row_id, shared)
|
||||
|
||||
def test_a_training_row_still_walks_its_own_vocabulary(
|
||||
self, import_status_table: Any, import_insert_status_row: Any, import_bff_engine: Any
|
||||
) -> None:
|
||||
training_row = import_insert_status_row(
|
||||
object_key=f'{IMPORT_PREFIX}training.sientia',
|
||||
expected_digest='c' * 64,
|
||||
run_type='TRAIN',
|
||||
)
|
||||
for status in ('PENDING', RUNNING_STATUS, 'TRAINING_SUCCESS'):
|
||||
_set_status(import_bff_engine, training_row, status)
|
||||
_set_status(import_bff_engine, training_row, 'TRAINING_ERROR')
|
||||
|
||||
def test_the_feature_added_no_column(
|
||||
self, import_status_table: Any, import_bff_engine: Any
|
||||
) -> None:
|
||||
"""Design D21: the error code and the workflow id live in the existing JSONB column.
|
||||
|
||||
A quiet `ALTER TABLE` would be the easiest way for this feature to break the BFF that owns
|
||||
the table, and the hardest to notice, so it is asserted rather than trusted.
|
||||
"""
|
||||
with import_bff_engine.begin() as connection:
|
||||
columns = {
|
||||
row[0]
|
||||
for row in connection.execute(
|
||||
text(
|
||||
'SELECT column_name FROM information_schema.columns '
|
||||
"WHERE table_schema = 'public' AND table_name = 'experiment_run'"
|
||||
)
|
||||
)
|
||||
}
|
||||
assert columns == EXPECTED_COLUMNS, (
|
||||
f'unexpected: {sorted(columns - EXPECTED_COLUMNS)}, '
|
||||
f'missing: {sorted(EXPECTED_COLUMNS - columns)}'
|
||||
)
|
||||
|
||||
|
||||
def _set_status(engine: Any, row_id: int, status: str) -> None:
|
||||
"""Write one status onto one row, so the real constraint gets to answer."""
|
||||
with engine.begin() as connection:
|
||||
connection.execute(
|
||||
text('UPDATE public.experiment_run SET status = :status WHERE id = :id'),
|
||||
{'status': status, 'id': row_id},
|
||||
)
|
||||
|
||||
|
||||
def _snapshot_table(engine: Any) -> list[dict[str, Any]]:
|
||||
"""Every row, ordered, as plain values — for "the table is byte-identical" assertions."""
|
||||
with engine.begin() as connection:
|
||||
return [
|
||||
dict(row)
|
||||
for row in connection.execute(
|
||||
text('SELECT * FROM public.experiment_run ORDER BY id')
|
||||
).mappings()
|
||||
]
|
||||
430
e2e/test_minimal_retrain.py
Normal file
430
e2e/test_minimal_retrain.py
Normal file
@@ -0,0 +1,430 @@
|
||||
"""
|
||||
End-to-end tests for the MinimalRetrain workflow.
|
||||
|
||||
The MLflow registry is fully stubbed because no real artifacts exist in a
|
||||
test container; we only validate that the workflow:
|
||||
|
||||
- Loads training data via ``load_query_with_minio_offload``.
|
||||
- Calls ``retrain_model`` with a payload pointing at MinIO.
|
||||
- Calls ``update_production_model`` only when retrain succeeds.
|
||||
- Persists ``sientia_data.log_retrain`` rows with all required columns;
|
||||
success rows carry the new ``version`` / ``mlflow_run_id`` /
|
||||
``mlflow_experiment_id`` while failure rows leave them ``NULL``.
|
||||
|
||||
The production DDL drops the legacy ``id`` / ``created_at`` columns and
|
||||
moves ``mlflow_experiment_id`` to ``int8`` and ``model_id`` to ``text``.
|
||||
The stubs used here therefore emit ``experiment_id`` as an integer to fit
|
||||
the new column type.
|
||||
"""
|
||||
|
||||
from contextlib import contextmanager
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import pandas as pd
|
||||
from sqlalchemy import text
|
||||
from temporalio.testing import WorkflowEnvironment
|
||||
from temporalio.worker import Worker
|
||||
|
||||
from e2e.helpers import (
|
||||
insert_target_data_for_drift,
|
||||
load_scenario_input,
|
||||
make_workflow_id,
|
||||
start_and_await_workflow,
|
||||
)
|
||||
from laborious.activities.activities import Activities
|
||||
from laborious.workflows.minimal_retrain import MinimalRetrain
|
||||
from sientia_model.wrappers.sientia_model import SientiaModel
|
||||
|
||||
# Columns defined by the production DDL for ``sientia_data.log_retrain``.
|
||||
# The legacy ``retrain_reports`` table had ``id`` and ``created_at``; the new
|
||||
# DDL drops both. ``mlflow_experiment_id`` is ``int8`` and ``model_id`` is
|
||||
# ``text``.
|
||||
EXPECTED_RETRAIN_REPORT_COLUMNS = [
|
||||
'mlflow_experiment_id',
|
||||
'mlflow_run_id',
|
||||
'model_id',
|
||||
'model_name',
|
||||
'status',
|
||||
'timestamp',
|
||||
'version',
|
||||
]
|
||||
|
||||
# Matches ``retrain_model`` return ``message`` when ``success`` is True (also written to ``log_retrain.status``).
|
||||
RETRAIN_ACTIVITY_SUCCESS_MESSAGE = 'Model retrained successfully.'
|
||||
|
||||
|
||||
class _FakeSientiaModelForMinimalRetrain(SientiaModel):
|
||||
"""
|
||||
Fake SientiaModel that uses the real SientiaModel lifecycle to surface
|
||||
index-alignment issues during ``retrain()``.
|
||||
|
||||
It intentionally performs strict alignment inside ``_retrain_model``:
|
||||
``y.loc[x.index]``.
|
||||
"""
|
||||
|
||||
def __init__(self, *, target: str = 'sensor_1'):
|
||||
super().__init__(
|
||||
model_type='FakeMinimalRetrain',
|
||||
model_version='0.0.0',
|
||||
model=object(),
|
||||
transformer=object(),
|
||||
)
|
||||
self.target = target
|
||||
self.model_is_fitted = True
|
||||
self.force_retrain_error = False
|
||||
|
||||
def store_model( # type: ignore[override]
|
||||
self,
|
||||
name: str,
|
||||
signature=None,
|
||||
pip_requirements=None,
|
||||
code_path=None,
|
||||
) -> None:
|
||||
# No-op: E2E tests validate workflow persistence, not real MLflow artifacts.
|
||||
return None
|
||||
|
||||
def _predict(self, data: pd.DataFrame):
|
||||
pred = pd.DataFrame({'prediction': [0.5] * len(data)}, index=data.index)
|
||||
return pred, {}
|
||||
|
||||
def _transform(self, data: pd.DataFrame):
|
||||
out = data.drop(columns=[self.target], errors='ignore').copy()
|
||||
out.index = data.index
|
||||
return out, {}
|
||||
|
||||
def _train_transformer(self, train_data: pd.DataFrame, val_data: pd.DataFrame) -> None:
|
||||
return None
|
||||
|
||||
def _train_model(
|
||||
self,
|
||||
x: pd.DataFrame,
|
||||
y: pd.DataFrame,
|
||||
x_val: pd.DataFrame | None = None,
|
||||
y_val: pd.DataFrame | None = None,
|
||||
) -> None:
|
||||
return None
|
||||
|
||||
def _retrain_transformer(self, data: pd.DataFrame) -> None:
|
||||
return None
|
||||
|
||||
def _retrain_model(self, x: pd.DataFrame, y: pd.DataFrame | None) -> None:
|
||||
if self.force_retrain_error:
|
||||
raise RuntimeError('training did not converge')
|
||||
if y is None:
|
||||
return
|
||||
# Strict alignment on purpose to reproduce the production failure mode.
|
||||
_ = y.loc[x.index]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mlflow_repository_stub():
|
||||
"""
|
||||
Override the shared E2E fixture: return a real fake ``SientiaModel`` wrapper
|
||||
instead of a MagicMock wrapper.
|
||||
"""
|
||||
repo = MagicMock()
|
||||
repo._client = MagicMock()
|
||||
|
||||
wrapper = _FakeSientiaModelForMinimalRetrain(target='sensor_1')
|
||||
repo.get_cached_model = MagicMock(return_value=wrapper)
|
||||
return repo
|
||||
|
||||
|
||||
def _retrain_input(model_id: int, **overrides) -> dict:
|
||||
"""Load and override the minimal-retrain base scenario."""
|
||||
payload = load_scenario_input('minimal_retrain_base.json', model_id=model_id)
|
||||
payload.update(overrides)
|
||||
return payload
|
||||
|
||||
|
||||
def _seed_retrain_training_rows(postgres_engine, model_id: int) -> None:
|
||||
"""
|
||||
Insert training rows in long format that pivot cleanly into
|
||||
``index=timestamp`` / ``columns={sensor_1, sensor_2}`` for ``retrain_model``.
|
||||
"""
|
||||
target_timestamps = [
|
||||
(datetime.now(timezone.utc).replace(second=0, microsecond=0) - timedelta(minutes=10 - i))
|
||||
.strftime('%Y-%m-%d %H:%M:%S%z')
|
||||
for i in range(5)
|
||||
]
|
||||
insert_target_data_for_drift(
|
||||
postgres_engine,
|
||||
model_id=model_id,
|
||||
timestamps=target_timestamps,
|
||||
variables_values={
|
||||
'sensor_1': [10.0, 11.0, 12.0, 13.0, 14.0],
|
||||
'sensor_2': [20.0, 21.0, 22.0, 23.0, 24.0],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _configure_retrain_happy_path(mlflow_repository_stub) -> None:
|
||||
"""
|
||||
Wire ``mlflow_repository_stub`` so retrain + update_production succeed.
|
||||
|
||||
Mocks (in order of consumption):
|
||||
|
||||
- ``_client.get_model_version_by_alias``: returns ``mv`` with a stable
|
||||
``run_id`` (used as ``source_run_id``).
|
||||
- ``start_run``: returns a context manager yielding a ``run_info`` with
|
||||
run/experiment ids.
|
||||
- ``log_params``: inert.
|
||||
- ``_client.search_model_versions``: returns one registry entry whose
|
||||
``version`` is promoted by ``update_production_model``.
|
||||
- ``promote_to_alias``: inert success.
|
||||
"""
|
||||
mv_src = MagicMock()
|
||||
mv_src.run_id = 'source-run-id'
|
||||
|
||||
new_version = MagicMock()
|
||||
new_version.version = '7'
|
||||
new_version.run_id = 'retrain-run-id'
|
||||
|
||||
mlflow_repository_stub._client.get_model_version_by_alias.return_value = mv_src
|
||||
|
||||
@contextmanager
|
||||
def fake_start_run(**kwargs):
|
||||
run_info = MagicMock()
|
||||
run_info.run_id = 'retrain-run-id'
|
||||
# ``mlflow_experiment_id`` is ``int8`` in the new DDL, so we feed an
|
||||
# integer-compatible id from the stubbed run info.
|
||||
run_info.experiment_id = 4242
|
||||
yield run_info
|
||||
|
||||
mlflow_repository_stub.start_run.side_effect = fake_start_run
|
||||
mlflow_repository_stub.log_params = MagicMock(return_value=None)
|
||||
mlflow_repository_stub._client.search_model_versions.return_value = [new_version]
|
||||
mlflow_repository_stub.promote_to_alias = MagicMock(return_value=None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_minimal_retrain_happy_path_writes_success_report(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker_minimal_retrain: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
mlflow_repository_stub,
|
||||
):
|
||||
"""
|
||||
Scenario MR.1.1: Retrain succeeds. ``sientia_data.log_retrain`` must
|
||||
contain a success row with version/mlflow_run_id/mlflow_experiment_id
|
||||
populated and the registry must have been told to promote the new version
|
||||
to the configured alias.
|
||||
"""
|
||||
client = temporal_test_env.client
|
||||
model_id = 711
|
||||
|
||||
_seed_retrain_training_rows(postgres_engine, model_id)
|
||||
_configure_retrain_happy_path(mlflow_repository_stub)
|
||||
|
||||
input_data = _retrain_input(model_id)
|
||||
|
||||
with patch('laborious.activities.mlflow.mlflow.log_artifact') as log_artifact_mock:
|
||||
await start_and_await_workflow(
|
||||
client,
|
||||
MinimalRetrain.run,
|
||||
input_data,
|
||||
make_workflow_id('test-retrain-happy'),
|
||||
)
|
||||
|
||||
with postgres_engine.connect() as conn:
|
||||
rows = (
|
||||
conn.execute(
|
||||
text(
|
||||
'SELECT * FROM sientia_data.log_retrain '
|
||||
'WHERE model_id = :m'
|
||||
),
|
||||
{'m': str(model_id)},
|
||||
)
|
||||
.mappings()
|
||||
.all()
|
||||
)
|
||||
|
||||
assert len(rows) == 1
|
||||
row = rows[0]
|
||||
for column in EXPECTED_RETRAIN_REPORT_COLUMNS:
|
||||
assert column in row, f'Missing log_retrain column: {column}'
|
||||
|
||||
assert row['status'] == RETRAIN_ACTIVITY_SUCCESS_MESSAGE, (
|
||||
"Expected retrain_model to return success (experiment_response['success'] is True). "
|
||||
'Persisted log_retrain.status is the activity message; when success is False the run '
|
||||
'never reaches mlflow.log_artifact — diagnose the retrain failure from status below, '
|
||||
'not from a skipped artifact upload. '
|
||||
f"Got status={row['status']!r}, version={row.get('version')!r}, "
|
||||
f"mlflow_run_id={row.get('mlflow_run_id')!r}."
|
||||
)
|
||||
|
||||
assert log_artifact_mock.called, (
|
||||
'After a successful retrain, retrain_model must call mlflow.log_artifact for the '
|
||||
'input CSV inside start_run.'
|
||||
)
|
||||
|
||||
# ``model_id`` is now ``text``; compare against the stringified id.
|
||||
assert row['model_id'] == str(model_id)
|
||||
assert row['model_name'] == 'test_model'
|
||||
assert row['version'] == '7'
|
||||
assert row['mlflow_run_id'] == 'retrain-run-id'
|
||||
# ``mlflow_experiment_id`` is now ``int8``; assert the integer value
|
||||
# provided by the stubbed run info.
|
||||
assert row['mlflow_experiment_id'] == 4242
|
||||
assert row['timestamp'] is not None
|
||||
|
||||
mlflow_repository_stub.promote_to_alias.assert_called_once()
|
||||
promote_kwargs = mlflow_repository_stub.promote_to_alias.call_args.kwargs
|
||||
assert promote_kwargs['model_name'] == 'test_model'
|
||||
assert promote_kwargs['version'] == '7'
|
||||
assert promote_kwargs['alias'] == 'production'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_minimal_retrain_failure_writes_report_without_version_columns(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker_minimal_retrain: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
mlflow_repository_stub,
|
||||
):
|
||||
"""
|
||||
Scenario MR.2.1: ``wrapper.retrain`` raises. The activity must catch the
|
||||
error, return ``success=False`` so ``update_production_model`` is skipped,
|
||||
and ``format_retrain_report`` must produce a row with the error message
|
||||
and NULL version columns.
|
||||
"""
|
||||
client = temporal_test_env.client
|
||||
model_id = 721
|
||||
|
||||
_seed_retrain_training_rows(postgres_engine, model_id)
|
||||
_configure_retrain_happy_path(mlflow_repository_stub)
|
||||
mlflow_repository_stub.get_cached_model.return_value.force_retrain_error = True
|
||||
|
||||
input_data = _retrain_input(model_id)
|
||||
|
||||
with patch('laborious.activities.mlflow.mlflow.log_artifact'):
|
||||
await start_and_await_workflow(
|
||||
client,
|
||||
MinimalRetrain.run,
|
||||
input_data,
|
||||
make_workflow_id('test-retrain-failure'),
|
||||
)
|
||||
|
||||
with postgres_engine.connect() as conn:
|
||||
rows = (
|
||||
conn.execute(
|
||||
text(
|
||||
'SELECT * FROM sientia_data.log_retrain '
|
||||
'WHERE model_id = :m'
|
||||
),
|
||||
{'m': str(model_id)},
|
||||
)
|
||||
.mappings()
|
||||
.all()
|
||||
)
|
||||
|
||||
assert len(rows) == 1
|
||||
row = rows[0]
|
||||
assert row['model_id'] == str(model_id)
|
||||
assert row['model_name'] == 'test_model'
|
||||
assert 'training did not converge' in row['status']
|
||||
assert row['version'] is None
|
||||
assert row['mlflow_run_id'] is None
|
||||
assert row['mlflow_experiment_id'] is None
|
||||
|
||||
mlflow_repository_stub.promote_to_alias.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_minimal_retrain_missing_target_writes_failure_report(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker_minimal_retrain: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
mlflow_repository_stub,
|
||||
):
|
||||
"""
|
||||
Scenario MR.2.2: ``model_config`` does not declare ``target``. The retrain
|
||||
activity must short-circuit before any MLflow call and the report row must
|
||||
carry the explicit guard message.
|
||||
"""
|
||||
client = temporal_test_env.client
|
||||
model_id = 722
|
||||
|
||||
_seed_retrain_training_rows(postgres_engine, model_id)
|
||||
_configure_retrain_happy_path(mlflow_repository_stub)
|
||||
|
||||
input_data = _retrain_input(model_id, model_config={})
|
||||
|
||||
with patch('laborious.activities.mlflow.mlflow.log_artifact'):
|
||||
await start_and_await_workflow(
|
||||
client,
|
||||
MinimalRetrain.run,
|
||||
input_data,
|
||||
make_workflow_id('test-retrain-missing-target'),
|
||||
)
|
||||
|
||||
with postgres_engine.connect() as conn:
|
||||
row = (
|
||||
conn.execute(
|
||||
text(
|
||||
'SELECT * FROM sientia_data.log_retrain '
|
||||
'WHERE model_id = :m'
|
||||
),
|
||||
{'m': str(model_id)},
|
||||
)
|
||||
.mappings()
|
||||
.first()
|
||||
)
|
||||
|
||||
assert row is not None
|
||||
assert 'target' in row['status'].lower(), (
|
||||
f"expected target-missing message, got status={row['status']!r}"
|
||||
)
|
||||
assert row['version'] is None
|
||||
mlflow_repository_stub.get_cached_model.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_minimal_retrain_no_training_data_does_not_persist_report(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker_minimal_retrain: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
mlflow_repository_stub,
|
||||
):
|
||||
"""
|
||||
Scenario MR.3.1: When the training query returns no rows the workflow must
|
||||
not persist any report row. The workflow currently raises plain
|
||||
``ValueError`` which Temporal treats as a workflow-task failure (causing
|
||||
indefinite retries until the test environment times out), so the assertion
|
||||
here is constrained to the persistence side-effect. See ``e2e/CODE_ISSUES.md``
|
||||
issue MR-1 for the recommended ``ApplicationError`` fix.
|
||||
"""
|
||||
client = temporal_test_env.client
|
||||
model_id = 731
|
||||
|
||||
with postgres_engine.begin() as conn:
|
||||
conn.execute(text(f'DELETE FROM sientia_data.laborious_data WHERE model_id = {model_id}'))
|
||||
|
||||
_configure_retrain_happy_path(mlflow_repository_stub)
|
||||
|
||||
input_data = _retrain_input(model_id)
|
||||
|
||||
with pytest.raises(Exception), patch('laborious.activities.mlflow.mlflow.log_artifact'):
|
||||
await start_and_await_workflow(
|
||||
client,
|
||||
MinimalRetrain.run,
|
||||
input_data,
|
||||
make_workflow_id('test-retrain-no-data'),
|
||||
)
|
||||
|
||||
with postgres_engine.connect() as conn:
|
||||
count = conn.execute(
|
||||
text('SELECT COUNT(*) FROM sientia_data.log_retrain WHERE model_id = :m'),
|
||||
{'m': str(model_id)},
|
||||
).scalar()
|
||||
assert count == 0
|
||||
@@ -9,7 +9,12 @@ from sqlalchemy import text
|
||||
from temporalio.testing import WorkflowEnvironment
|
||||
from temporalio.worker import Worker
|
||||
|
||||
from e2e.helpers import insert_sample_data, make_workflow_id, start_and_await_workflow
|
||||
from e2e.helpers import (
|
||||
insert_sample_data,
|
||||
load_scenario_input,
|
||||
make_workflow_id,
|
||||
start_and_await_workflow,
|
||||
)
|
||||
from laborious.activities.activities import Activities
|
||||
from laborious.utils.models import minio_dataframe_payload as mdp
|
||||
from laborious.workflows.predictions_batch import PredictionsBatch
|
||||
@@ -29,45 +34,23 @@ async def test_load_query_with_minio_offload_writes_object_to_bucket(
|
||||
"""
|
||||
model_id = 501
|
||||
with postgres_engine.begin() as conn:
|
||||
conn.execute(
|
||||
text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}')
|
||||
)
|
||||
conn.execute(text(f'DELETE FROM sientia_data.laborious_data WHERE model_id = {model_id}'))
|
||||
insert_sample_data(postgres_engine, model_id, [1.0, 2.0])
|
||||
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'schedule_name': 'test-schedule',
|
||||
'model_name': 'test_model',
|
||||
'model_id': model_id,
|
||||
'workflow_name': 'predictions_batch',
|
||||
}
|
||||
}
|
||||
scenario_input = load_scenario_input('minio_offload_load_query.json', model_id=model_id)
|
||||
metadata = {'metadata': scenario_input['metadata']}
|
||||
with patch.object(mdp, 'OFFLOAD_THRESHOLD_BYTES', 1):
|
||||
payload = await test_activities_real_minio.load_query_with_minio_offload(
|
||||
{
|
||||
**metadata,
|
||||
'query': (
|
||||
'SELECT timestamp, variable, value, created_at '
|
||||
f'FROM predictions_schema.laborious_data WHERE model_id = {model_id}'
|
||||
),
|
||||
'model_name': 'test_model',
|
||||
'datetime_columns': ['timestamp', 'created_at'],
|
||||
}
|
||||
)
|
||||
payload = test_activities_real_minio.load_query_with_minio_offload(scenario_input)
|
||||
assert payload.object_key, 'offloaded payload must reference a MinIO object'
|
||||
assert payload.data is None or payload.data == {}, (
|
||||
'large payloads should not inline tabular dict'
|
||||
)
|
||||
assert payload.data is None or payload.data == {}, 'large payloads should not inline tabular dict'
|
||||
|
||||
df = await payload.retrieve(test_activities_real_minio.minio_repository, metadata['metadata'])
|
||||
df = payload.retrieve(test_activities_real_minio.minio_repository, metadata['metadata'])
|
||||
assert len(df) >= 1
|
||||
|
||||
client = minio_container.get_client()
|
||||
listed = list(client.list_objects('test-bucket', recursive=True))
|
||||
names = [getattr(o, 'object_name', None) or getattr(o, '_object_name', '') for o in listed]
|
||||
assert any(n and 'prediction_datasets' in n for n in names), (
|
||||
f'unexpected object listing: {names!r}'
|
||||
)
|
||||
assert any(n and 'prediction_datasets' in n for n in names), f'unexpected object listing: {names!r}'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -83,43 +66,12 @@ async def test_predictions_batch_with_minio_offload_path(
|
||||
"""
|
||||
model_id = 502
|
||||
with postgres_engine.begin() as conn:
|
||||
conn.execute(
|
||||
text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}')
|
||||
)
|
||||
conn.execute(
|
||||
text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')
|
||||
)
|
||||
conn.execute(
|
||||
text(f'DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}')
|
||||
)
|
||||
conn.execute(text(f'DELETE FROM sientia_data.laborious_data WHERE model_id = {model_id}'))
|
||||
conn.execute(text(f'DELETE FROM sientia_data.predictions WHERE model_id = {model_id}'))
|
||||
conn.execute(text(f'DELETE FROM sientia_data.transformed_data WHERE model_id = {model_id}'))
|
||||
insert_sample_data(postgres_engine, model_id, [10.0, 20.0, 30.0])
|
||||
|
||||
input_data = {
|
||||
'schedule_name': 'test-schedule',
|
||||
'model_name': 'test_model',
|
||||
'model_id': model_id,
|
||||
'query': (
|
||||
'SELECT timestamp, variable, value, created_at '
|
||||
f'FROM predictions_schema.laborious_data WHERE model_id = {model_id}'
|
||||
),
|
||||
'schema': 'predictions_schema',
|
||||
'table_name': 'predictions',
|
||||
'transform_table_name': 'transformed_data',
|
||||
'input_filters': {'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}},
|
||||
'mlflow_transform_filters': {'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}},
|
||||
'mlflow_predict_filters': {'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}},
|
||||
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
|
||||
'opc_output_config': {},
|
||||
'pi_web_api_output_config': {},
|
||||
'save_transform': True,
|
||||
'prediction_store_policy': 'lts:1',
|
||||
'model_config': {
|
||||
'retention_minutes': 0,
|
||||
'transform_flavor': 'sklearn',
|
||||
'predict_flavor': 'sklearn',
|
||||
},
|
||||
'datetime_columns': ['timestamp', 'created_at'],
|
||||
}
|
||||
input_data = load_scenario_input('minio_offload_workflow.json', model_id=model_id)
|
||||
|
||||
with patch.object(mdp, 'OFFLOAD_THRESHOLD_BYTES', 1):
|
||||
await start_and_await_workflow(
|
||||
@@ -131,6 +83,26 @@ async def test_predictions_batch_with_minio_offload_path(
|
||||
|
||||
with postgres_engine.connect() as conn:
|
||||
count = conn.execute(
|
||||
text(f'SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = {model_id}')
|
||||
text(f'SELECT COUNT(*) FROM sientia_data.predictions WHERE model_id = {model_id}')
|
||||
).scalar()
|
||||
assert count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_load_query_with_inline_payload_when_below_threshold(
|
||||
postgres_engine,
|
||||
test_activities_real_minio: Activities,
|
||||
):
|
||||
"""Scenario 4.2.1: payload stays inline when threshold is high enough."""
|
||||
model_id = 503
|
||||
with postgres_engine.begin() as conn:
|
||||
conn.execute(text(f'DELETE FROM sientia_data.laborious_data WHERE model_id = {model_id}'))
|
||||
insert_sample_data(postgres_engine, model_id, [1.0, 2.0])
|
||||
|
||||
scenario_input = load_scenario_input('minio_offload_load_query.json', model_id=model_id)
|
||||
with patch.object(mdp, 'OFFLOAD_THRESHOLD_BYTES', 10**9):
|
||||
payload = test_activities_real_minio.load_query_with_minio_offload(scenario_input)
|
||||
|
||||
assert payload.object_key is None
|
||||
assert payload.data is not None
|
||||
|
||||
@@ -6,6 +6,7 @@ Mock-based OPC tests remain in test_predictions_batch_format_export.py.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
|
||||
import pytest
|
||||
from temporalio.testing import WorkflowEnvironment
|
||||
@@ -25,8 +26,8 @@ from laborious.utils.repository.opc_repository import OpcRepository
|
||||
from laborious.workflows.predictions_batch import PredictionsBatch
|
||||
|
||||
|
||||
async def _reconnect_under_held_lock(
|
||||
repo: OpcRepository, acquired: asyncio.Event, release: asyncio.Event
|
||||
def _reconnect_under_held_lock(
|
||||
repo: OpcRepository, acquired: threading.Event, release: threading.Event
|
||||
) -> None:
|
||||
"""
|
||||
Hold the connection lock until released, then reconnect.
|
||||
@@ -37,13 +38,13 @@ async def _reconnect_under_held_lock(
|
||||
|
||||
Args:
|
||||
repo (OpcRepository): Connected repository.
|
||||
acquired (asyncio.Event): Set once the lock is held, before any write can race.
|
||||
release (asyncio.Event): Awaited while holding the lock; set by the caller.
|
||||
acquired (threading.Event): Set once the lock is held, before any write can race.
|
||||
release (threading.Event): Waited on while holding the lock; set by the caller.
|
||||
"""
|
||||
async with repo._connection_lock:
|
||||
with repo._connection_lock:
|
||||
acquired.set()
|
||||
await release.wait()
|
||||
await repo._reconnect_locked()
|
||||
release.wait()
|
||||
repo._reconnect_locked()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -183,10 +184,15 @@ async def test_scenario_3_2_5_opc_write_blocked_during_reconnect_real_server(
|
||||
|
||||
repo = test_activities_real_opc.opc_repository['1']
|
||||
repo._session_ready.clear()
|
||||
acquired = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
reconnect_task = asyncio.create_task(_reconnect_under_held_lock(repo, acquired, release))
|
||||
await acquired.wait()
|
||||
acquired = threading.Event()
|
||||
release = threading.Event()
|
||||
reconnect_thread = threading.Thread(
|
||||
target=_reconnect_under_held_lock,
|
||||
args=(repo, acquired, release),
|
||||
daemon=True,
|
||||
)
|
||||
reconnect_thread.start()
|
||||
acquired.wait(timeout=5.0)
|
||||
|
||||
input_data = get_base_input_data(model_id)
|
||||
input_data['opc_output_config'] = build_opc_output_config(opc_e2e_server.node_ids)
|
||||
@@ -201,7 +207,7 @@ async def test_scenario_3_2_5_opc_write_blocked_during_reconnect_real_server(
|
||||
)
|
||||
finally:
|
||||
release.set()
|
||||
await reconnect_task
|
||||
reconnect_thread.join(timeout=5.0)
|
||||
|
||||
assert_prediction(
|
||||
postgres_engine,
|
||||
|
||||
@@ -4,7 +4,7 @@ End-to-end tests for PredictionsBatch workflow - Format and Export scenarios.
|
||||
|
||||
from decimal import Decimal
|
||||
from typing import Any, cast
|
||||
from unittest.mock import AsyncMock, call
|
||||
from unittest.mock import call
|
||||
|
||||
import pytest
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
@@ -15,51 +15,15 @@ from temporalio.worker import Worker
|
||||
from e2e.helpers import (
|
||||
assert_prediction,
|
||||
insert_sample_data,
|
||||
load_scenario_input,
|
||||
make_workflow_id,
|
||||
start_and_await_workflow,
|
||||
)
|
||||
from laborious.activities.activities import Activities
|
||||
from laborious.workflows.predictions_batch import PredictionsBatch
|
||||
|
||||
base_input_data = {
|
||||
'schedule_name': 'test-schedule',
|
||||
'model_name': 'test_model',
|
||||
'model_id': 301,
|
||||
'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 301',
|
||||
'schema': 'predictions_schema',
|
||||
'table_name': 'predictions',
|
||||
'transform_table_name': 'transformed_data',
|
||||
'input_filters': {
|
||||
'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}},
|
||||
},
|
||||
'mlflow_transform_filters': {
|
||||
'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
|
||||
},
|
||||
'mlflow_predict_filters': {
|
||||
'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
|
||||
},
|
||||
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
|
||||
'opc_output_config': {},
|
||||
'pi_web_api_output_config': {},
|
||||
'save_transform': True,
|
||||
'prediction_store_policy': 'lts:1',
|
||||
'model_config': {
|
||||
'retention_minutes': 0,
|
||||
'transform_flavor': 'sklearn',
|
||||
'predict_flavor': 'sklearn',
|
||||
},
|
||||
'datetime_columns': ['timestamp', 'created_at'],
|
||||
}
|
||||
|
||||
base_query = 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = {model_id}'
|
||||
|
||||
|
||||
def get_base_input_data(model_id):
|
||||
return {
|
||||
**base_input_data,
|
||||
'model_id': model_id,
|
||||
'query': base_query.format(model_id=model_id),
|
||||
}
|
||||
return load_scenario_input('format_export_base.json', model_id=model_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -85,12 +49,8 @@ async def test_scenario_3_1_1_default_prediction_export(
|
||||
model_id = 311
|
||||
|
||||
with postgres_engine.begin() as conn:
|
||||
conn.execute(
|
||||
text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')
|
||||
)
|
||||
conn.execute(
|
||||
text(f'DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}')
|
||||
)
|
||||
conn.execute(text(f"DELETE FROM sientia_data.predictions WHERE model_id = {model_id}"))
|
||||
conn.execute(text(f"DELETE FROM sientia_data.transformed_data WHERE model_id = {model_id}"))
|
||||
insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
|
||||
|
||||
input_data = get_base_input_data(model_id)
|
||||
@@ -186,9 +146,7 @@ async def test_scenario_3_1_1_default_prediction_export(
|
||||
|
||||
with postgres_engine.connect() as conn:
|
||||
tf_count = conn.execute(
|
||||
text(
|
||||
f'SELECT COUNT(*) FROM predictions_schema.transformed_data WHERE model_id = {model_id}'
|
||||
)
|
||||
text(f"SELECT COUNT(*) FROM sientia_data.transformed_data WHERE model_id = {model_id}")
|
||||
).scalar()
|
||||
assert tf_count == 0, 'transform export must be skipped when path_flag is set'
|
||||
|
||||
@@ -202,6 +160,9 @@ async def test_scenario_3_1_1_default_prediction_export(
|
||||
)
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_scenario_3_1_2_export_with_opc_only(
|
||||
@@ -212,17 +173,17 @@ async def test_scenario_3_1_2_export_with_opc_only(
|
||||
):
|
||||
"""
|
||||
Scenario 3.1.2: Export with OPC only
|
||||
|
||||
|
||||
Description:
|
||||
Export to PostgreSQL and OPC server only (no PI Web API).
|
||||
|
||||
|
||||
Expected Behavior:
|
||||
- Normal formatting
|
||||
- PostgreSQL export executed
|
||||
- OPC export executed
|
||||
- PI Web API activity skipped
|
||||
- Metrics written with OPC metrics
|
||||
|
||||
|
||||
Assertions:
|
||||
- PI Web API activity NOT called
|
||||
- OPC activity called
|
||||
@@ -289,6 +250,7 @@ async def test_scenario_3_1_2_export_with_opc_only(
|
||||
assert_prediction(postgres_engine, model_id)
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_scenario_3_1_3_export_with_pi_web_api_only(
|
||||
@@ -299,17 +261,17 @@ async def test_scenario_3_1_3_export_with_pi_web_api_only(
|
||||
):
|
||||
"""
|
||||
Scenario 3.1.3: Export with PI Web API only
|
||||
|
||||
|
||||
Description:
|
||||
Export to PostgreSQL and PI Web API only (no OPC).
|
||||
|
||||
|
||||
Expected Behavior:
|
||||
- Normal formatting
|
||||
- PostgreSQL export executed
|
||||
- PI Web API export executed
|
||||
- OPC activity skipped
|
||||
- Metrics written without OPC metrics
|
||||
|
||||
|
||||
Assertions:
|
||||
- OPC activity NOT called
|
||||
- PI Web API activity called
|
||||
@@ -372,6 +334,7 @@ async def test_scenario_3_1_3_export_with_pi_web_api_only(
|
||||
assert_prediction(postgres_engine, model_id)
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_scenario_3_1_4_export_without_optional_outputs(
|
||||
@@ -382,16 +345,16 @@ async def test_scenario_3_1_4_export_without_optional_outputs(
|
||||
):
|
||||
"""
|
||||
Scenario 3.1.4: Export Without Optional Outputs
|
||||
|
||||
|
||||
Description:
|
||||
Export only to PostgreSQL (no OPC or PI Web API).
|
||||
|
||||
|
||||
Expected Behavior:
|
||||
- Normal formatting
|
||||
- Only PostgreSQL export executed
|
||||
- OPC and PI Web API activities skipped
|
||||
- Metrics written without OPC metrics
|
||||
|
||||
|
||||
Assertions:
|
||||
- PI Web API activity NOT called
|
||||
- OPC activity NOT called
|
||||
@@ -419,6 +382,7 @@ async def test_scenario_3_1_4_export_without_optional_outputs(
|
||||
assert_prediction(postgres_engine, model_id)
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_scenario_3_1_5_export_without_transformed_data(
|
||||
@@ -429,15 +393,15 @@ async def test_scenario_3_1_5_export_without_transformed_data(
|
||||
):
|
||||
"""
|
||||
Scenario 3.1.5: Export Without Transformed Data
|
||||
|
||||
|
||||
Description:
|
||||
Only prediction exported, no transform table.
|
||||
|
||||
|
||||
Expected Behavior:
|
||||
- Only prediction formatted and exported
|
||||
- Transform export skipped
|
||||
- Single PostgreSQL write
|
||||
|
||||
|
||||
Assertions:
|
||||
- format_transformed_data NOT called
|
||||
- One PostgreSQL export
|
||||
@@ -449,9 +413,7 @@ async def test_scenario_3_1_5_export_without_transformed_data(
|
||||
|
||||
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
||||
with postgres_engine.begin() as conn:
|
||||
conn.execute(
|
||||
text(f'DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}')
|
||||
)
|
||||
conn.execute(text(f"DELETE FROM sientia_data.transformed_data WHERE model_id = {model_id}"))
|
||||
|
||||
input_data = get_base_input_data(model_id)
|
||||
input_data['save_transform'] = False # Don't save transformed data
|
||||
@@ -541,16 +503,15 @@ async def test_scenario_3_1_5_export_without_transformed_data(
|
||||
|
||||
with postgres_engine.connect() as conn:
|
||||
result_query = conn.execute(
|
||||
text(
|
||||
f'SELECT COUNT(*) FROM predictions_schema.transformed_data WHERE model_id = {model_id}'
|
||||
)
|
||||
text(f"SELECT COUNT(*) FROM sientia_data.transformed_data WHERE model_id = {model_id}")
|
||||
)
|
||||
count = result_query.scalar()
|
||||
assert count == 0, f'Expected transform table to be empty, but found {count} records'
|
||||
assert count == 0, f"Expected transform table to be empty, but found {count} records"
|
||||
|
||||
assert_prediction(postgres_engine, model_id)
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_scenario_3_2_1_pi_web_api_write_error(
|
||||
@@ -574,8 +535,7 @@ async def test_scenario_3_2_1_pi_web_api_write_error(
|
||||
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
||||
|
||||
test_activities.pi_web_api_client.write_value.side_effect = Exception(
|
||||
'PI Web API service unavailable'
|
||||
)
|
||||
"PI Web API service unavailable")
|
||||
|
||||
input_data = get_base_input_data(model_id)
|
||||
input_data['pi_web_api_output_config'] = {
|
||||
@@ -603,14 +563,14 @@ async def test_scenario_3_2_1_pi_web_api_write_error(
|
||||
)
|
||||
|
||||
assert_prediction(
|
||||
postgres_engine,
|
||||
model_id,
|
||||
postgres_engine, model_id,
|
||||
prediction_confidence=13,
|
||||
comments='PI Web API service unavailable',
|
||||
)
|
||||
assert notification_inserts.call_count >= 1
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_scenario_3_2_2_opc_write_error(
|
||||
@@ -632,17 +592,14 @@ async def test_scenario_3_2_2_opc_write_error(
|
||||
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
||||
|
||||
opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data)
|
||||
opc_write_data.return_value = (
|
||||
False,
|
||||
{
|
||||
'notification_id': 'OPC_WRITE_DATA_ERROR_1',
|
||||
'message': 'OPC server unavailable',
|
||||
'block': 'opc_repository',
|
||||
'level': NotificationLevel.ERROR,
|
||||
'attachment_content': 'OPC server unavailable',
|
||||
},
|
||||
)
|
||||
|
||||
opc_write_data.return_value = (False, {
|
||||
'notification_id': 'OPC_WRITE_DATA_ERROR_1',
|
||||
'message': 'OPC server unavailable',
|
||||
'block': 'opc_repository',
|
||||
'level': NotificationLevel.ERROR,
|
||||
'attachment_content': 'OPC server unavailable',
|
||||
})
|
||||
|
||||
input_data = get_base_input_data(model_id)
|
||||
input_data['opc_output_config'] = {
|
||||
'1': {
|
||||
@@ -669,8 +626,7 @@ async def test_scenario_3_2_2_opc_write_error(
|
||||
)
|
||||
|
||||
assert_prediction(
|
||||
postgres_engine,
|
||||
model_id,
|
||||
postgres_engine, model_id,
|
||||
prediction_confidence=12,
|
||||
comments='Some data could not be written to OPC servers',
|
||||
)
|
||||
@@ -678,19 +634,16 @@ async def test_scenario_3_2_2_opc_write_error(
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_scenario_3_2_4_opc_session_bad_error(
|
||||
async def test_scenario_3_2_4_opc_session_bad_mock(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
):
|
||||
"""
|
||||
Scenario 3.2.4: OPC session/channel Tier-1 Bad* (e.g. BadSessionIdInvalid).
|
||||
|
||||
PostgreSQL stores prediction_confidence 14 and a stable session error comment.
|
||||
Scenario 3.2.4 (mock): Tier-1 session error maps to confidence 14.
|
||||
"""
|
||||
client = temporal_test_env.client
|
||||
|
||||
model_id = 324
|
||||
|
||||
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
||||
@@ -699,42 +652,87 @@ async def test_scenario_3_2_4_opc_session_bad_error(
|
||||
opc_write_data.return_value = (
|
||||
False,
|
||||
{
|
||||
'opc_error_kind': 'session_bad',
|
||||
'opc_status': 'BadSessionIdInvalid',
|
||||
'notification_id': 'OPC_WRITE_DATA_ERROR_1',
|
||||
'message': 'BadSessionIdInvalid',
|
||||
'message': 'OPC session invalid',
|
||||
'block': 'opc_repository',
|
||||
'level': NotificationLevel.ERROR,
|
||||
'attachment_content': 'BadSessionIdInvalid',
|
||||
'opc_error_kind': 'session_bad',
|
||||
'opc_status': 'BadSessionIdInvalid',
|
||||
},
|
||||
)
|
||||
|
||||
input_data = get_base_input_data(model_id)
|
||||
input_data['opc_output_config'] = {
|
||||
'1': {
|
||||
'prediction_tags': {
|
||||
'addr_1': {
|
||||
'data_type': 'float',
|
||||
}
|
||||
},
|
||||
'confidence_tags': {
|
||||
'addr_2': {
|
||||
'data_type': 'float',
|
||||
}
|
||||
},
|
||||
'prediction_tags': {'addr_1': {'data_type': 'float'}},
|
||||
'confidence_tags': {'addr_2': {'data_type': 'float'}},
|
||||
}
|
||||
}
|
||||
input_data['pi_web_api_output_config'] = None
|
||||
|
||||
await start_and_await_workflow(
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-opc-session-bad')
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-opc-session-bad-mock')
|
||||
)
|
||||
|
||||
assert_prediction(
|
||||
postgres_engine,
|
||||
model_id,
|
||||
prediction_confidence=14,
|
||||
comments='OPC UA session/channel error: BadSessionIdInvalid',
|
||||
comments_contains='OPC UA session/channel error: BadSessionIdInvalid',
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_scenario_3_2_5_opc_reconnect_in_progress_mock(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
):
|
||||
"""
|
||||
Scenario 3.2.5 (mock): reconnect_in_progress maps to confidence 14.
|
||||
"""
|
||||
client = temporal_test_env.client
|
||||
model_id = 325
|
||||
|
||||
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
||||
|
||||
opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data)
|
||||
opc_write_data.return_value = (
|
||||
False,
|
||||
{
|
||||
'opc_error_kind': 'reconnect_in_progress',
|
||||
'notification_id': 'OPC_WRITE_DATA_ERROR_1',
|
||||
'message': 'OPC reconnect in progress',
|
||||
'block': 'opc_repository',
|
||||
'level': NotificationLevel.ERROR,
|
||||
'attachment_content': 'reconnect',
|
||||
},
|
||||
)
|
||||
|
||||
input_data = get_base_input_data(model_id)
|
||||
input_data['opc_output_config'] = {
|
||||
'1': {
|
||||
'prediction_tags': {'addr_1': {'data_type': 'float'}},
|
||||
'confidence_tags': {'addr_2': {'data_type': 'float'}},
|
||||
}
|
||||
}
|
||||
input_data['pi_web_api_output_config'] = None
|
||||
|
||||
await start_and_await_workflow(
|
||||
client,
|
||||
PredictionsBatch.run,
|
||||
input_data,
|
||||
make_workflow_id('test-opc-reconnect-mock'),
|
||||
)
|
||||
|
||||
assert_prediction(
|
||||
postgres_engine,
|
||||
model_id,
|
||||
prediction_confidence=14,
|
||||
comments_contains='OPC UA reconnect in progress',
|
||||
)
|
||||
|
||||
|
||||
@@ -758,8 +756,8 @@ async def test_scenario_3_2_3_pi_web_api_partial_write_error(
|
||||
|
||||
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
||||
|
||||
test_activities.pi_web_api_client.write_value = AsyncMock(
|
||||
side_effect=[
|
||||
test_activities.pi_web_api_client.set_side_effect(
|
||||
[
|
||||
# Prediction batch: two web_ids requested, only one acknowledged.
|
||||
[{'WebId': 'web_id_1', 'Errors': []}],
|
||||
# Confidence write succeeds.
|
||||
@@ -793,8 +791,46 @@ async def test_scenario_3_2_3_pi_web_api_partial_write_error(
|
||||
)
|
||||
|
||||
assert_prediction(
|
||||
postgres_engine,
|
||||
model_id,
|
||||
postgres_engine, model_id,
|
||||
prediction_confidence=13,
|
||||
comments="The number of written tags does not match the number of tag names: Expected ['tag_1', 'tag_3'] tags, but ['tag_1'] tags were written.",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_scenario_3_3_1_combined_pi_and_opc_outputs(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
):
|
||||
"""
|
||||
Scenario 3.3.1: PI and OPC enabled together.
|
||||
"""
|
||||
client = temporal_test_env.client
|
||||
model_id = 333
|
||||
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
||||
|
||||
input_data = get_base_input_data(model_id)
|
||||
input_data['pi_web_api_output_config'] = {
|
||||
'endpoint': 'test_endpoint',
|
||||
'prediction_tags': {'tag_1': 'web_id_1'},
|
||||
'confidence_tags': {'tag_2': 'web_id_2'},
|
||||
}
|
||||
input_data['opc_output_config'] = {
|
||||
'1': {
|
||||
'prediction_tags': {'addr_1': {'data_type': 'float'}},
|
||||
'confidence_tags': {'addr_2': {'data_type': 'float'}},
|
||||
}
|
||||
}
|
||||
|
||||
await start_and_await_workflow(
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-pi-opc-combined')
|
||||
)
|
||||
|
||||
assert test_activities.pi_web_api_client.write_value.call_count == 2
|
||||
opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data)
|
||||
assert opc_write_data.call_count == 2
|
||||
assert_prediction(postgres_engine, model_id)
|
||||
|
||||
|
||||
@@ -1,277 +0,0 @@
|
||||
"""
|
||||
End-to-end tests for the explicit `joblib` flavor on the MLflow model load paths.
|
||||
|
||||
`joblib.load` is NOT mocked here: the artifact fixture holds a real object serialized with
|
||||
`joblib.dump(..., compress=3)` whose class lives only inside the artifact's `code/` directory.
|
||||
A passing run therefore proves the flavor decompresses the pickle **and** resolves the module,
|
||||
without any `mlflow.*.load_model` being called for that model.
|
||||
"""
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pandas as pd
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
from temporalio.testing import WorkflowEnvironment
|
||||
from temporalio.worker import Worker
|
||||
|
||||
from e2e.conftest import JOBLIB_FLAVOR_MODEL_NAME
|
||||
from e2e.helpers import (
|
||||
assert_prediction,
|
||||
assert_stop,
|
||||
insert_sample_data,
|
||||
make_workflow_id,
|
||||
start_and_await_workflow,
|
||||
)
|
||||
from laborious.activities.activities import Activities
|
||||
from laborious.workflows.predictions_batch import PredictionsBatch
|
||||
|
||||
base_query = 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = {model_id}'
|
||||
|
||||
FLAVOR_LOADERS = ('sklearn', 'pyfunc')
|
||||
|
||||
|
||||
def build_input_data(
|
||||
model_id: int,
|
||||
model_name: str,
|
||||
transform_flavor: str = 'sklearn',
|
||||
predict_flavor: str = 'sklearn',
|
||||
) -> dict:
|
||||
"""Workflow input for a plain transform + predict run of a given model."""
|
||||
return {
|
||||
'schedule_name': 'test-schedule',
|
||||
'model_name': model_name,
|
||||
'model_id': model_id,
|
||||
'query': base_query.format(model_id=model_id),
|
||||
'schema': 'predictions_schema',
|
||||
'table_name': 'predictions',
|
||||
'transform_table_name': 'transformed_data',
|
||||
'input_filters': {
|
||||
'SPECIFIC_VARIABLES_NULL_VALUES': {
|
||||
'POLICY': 'CONTINUE',
|
||||
'CONFIG': {'variables': ['sensor_1']},
|
||||
},
|
||||
},
|
||||
'mlflow_transform_filters': {
|
||||
'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
|
||||
},
|
||||
'mlflow_predict_filters': {
|
||||
'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
|
||||
},
|
||||
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
|
||||
'opc_output_config': {},
|
||||
'pi_web_api_output_config': {},
|
||||
'save_transform': True,
|
||||
'prediction_store_policy': 'lts:1',
|
||||
'model_config': {
|
||||
'retention_minutes': 0,
|
||||
'transform_flavor': transform_flavor,
|
||||
'predict_flavor': predict_flavor,
|
||||
},
|
||||
'datetime_columns': ['timestamp', 'created_at'],
|
||||
}
|
||||
|
||||
|
||||
def expected_joblib_prediction(model) -> float:
|
||||
"""Prediction the artifact's model produces for the transform output of this suite."""
|
||||
features = pd.DataFrame(
|
||||
{
|
||||
'feature_1': [0.234],
|
||||
'feature_2': [0.783],
|
||||
'timestamp': ['2024-01-01 12:00:00'],
|
||||
}
|
||||
)
|
||||
return float(model.predict(features)[0])
|
||||
|
||||
|
||||
def assert_single_prediction_close(postgres_engine, model_id, expected, tolerance=1e-9):
|
||||
"""Assert exactly one Good prediction row whose value matches `expected`."""
|
||||
with postgres_engine.connect() as conn:
|
||||
rows = conn.execute(
|
||||
text(
|
||||
'SELECT prediction, prediction_status '
|
||||
'FROM predictions_schema.predictions '
|
||||
f'WHERE model_id = {model_id}'
|
||||
)
|
||||
).fetchall()
|
||||
|
||||
assert len(rows) == 1, f'Expected one prediction record, got {len(rows)}'
|
||||
assert abs(float(rows[0][0]) - expected) < tolerance, (
|
||||
f'Expected prediction close to {expected}, got {rows[0][0]}'
|
||||
)
|
||||
assert rows[0][1] == 'Good', f"Expected prediction_status='Good', got {rows[0][1]}"
|
||||
|
||||
|
||||
def logged_messages(mock_logger, level: str) -> str:
|
||||
"""Concatenate every message logged at a given custom level."""
|
||||
calls = getattr(mock_logger, f'custom_{level}').call_args_list
|
||||
return ' '.join(str(logged.args[0]) for logged in calls)
|
||||
|
||||
|
||||
def loader_uris(patch_mlflow) -> list[str]:
|
||||
"""Every model URI handed to an MLflow flavor loader during the run."""
|
||||
return [
|
||||
str(loaded.args[0])
|
||||
for flavor in FLAVOR_LOADERS
|
||||
for loaded in getattr(patch_mlflow, flavor).load_model.call_args_list
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_prediction_workflow_with_joblib_flavor(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
joblib_flavor_artifacts,
|
||||
mlflow_client,
|
||||
patch_mlflow,
|
||||
):
|
||||
"""`predict_flavor: 'joblib'` deserializes the real artifact, and the prediction lands."""
|
||||
client = temporal_test_env.client
|
||||
model_id = 271
|
||||
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
||||
input_data = build_input_data(model_id, JOBLIB_FLAVOR_MODEL_NAME, predict_flavor='joblib')
|
||||
|
||||
await start_and_await_workflow(
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-flavor')
|
||||
)
|
||||
|
||||
assert_single_prediction_close(
|
||||
postgres_engine,
|
||||
model_id,
|
||||
expected_joblib_prediction(joblib_flavor_artifacts['prediction_model']['model']),
|
||||
)
|
||||
|
||||
assert mlflow_client.download_artifacts.call_count == 1
|
||||
assert mlflow_client.download_artifacts.call_args.args[1] == 'prediction_model'
|
||||
|
||||
# The joblib flavor never builds a `models:/` URI, so no flavor loader saw this model.
|
||||
assert all(JOBLIB_FLAVOR_MODEL_NAME not in uri for uri in loader_uris(patch_mlflow))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_prediction_workflow_with_joblib_transform_flavor(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
joblib_flavor_artifacts,
|
||||
mlflow_client,
|
||||
patch_mlflow,
|
||||
mock_mlflow_models,
|
||||
):
|
||||
"""`transform_flavor: 'joblib'` downloads and deserializes `data_model` on its own."""
|
||||
client = temporal_test_env.client
|
||||
model_id = 274
|
||||
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
||||
input_data = build_input_data(model_id, JOBLIB_FLAVOR_MODEL_NAME, transform_flavor='joblib')
|
||||
|
||||
await start_and_await_workflow(
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-transform')
|
||||
)
|
||||
|
||||
assert_prediction(postgres_engine, model_id, prediction=0.5)
|
||||
mock_mlflow_models['predict_model'].predict.assert_called()
|
||||
|
||||
assert mlflow_client.download_artifacts.call_count == 1
|
||||
assert mlflow_client.download_artifacts.call_args.args[1] == 'data_model'
|
||||
|
||||
# The joblib flavor resolves no artifact URI, so `get_run` is never reached and no flavor
|
||||
# loader was given a `data_model` URI.
|
||||
patch_mlflow.get_run.assert_not_called()
|
||||
assert all(not uri.endswith('/data_model') for uri in loader_uris(patch_mlflow))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_prediction_workflow_with_sklearn_flavor_is_unchanged(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
mlflow_client,
|
||||
mock_mlflow_models,
|
||||
):
|
||||
"""A model on an MLflow flavor keeps the current result and downloads nothing."""
|
||||
client = temporal_test_env.client
|
||||
model_id = 272
|
||||
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
||||
input_data = build_input_data(model_id, 'test_model')
|
||||
|
||||
await start_and_await_workflow(
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-sklearn-flavor')
|
||||
)
|
||||
|
||||
assert_prediction(postgres_engine, model_id, prediction=0.5)
|
||||
mock_mlflow_models['predict_model'].predict.assert_called()
|
||||
mlflow_client.download_artifacts.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_prediction_workflow_downloads_only_the_artifact_code_directory(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
joblib_flavor_artifacts,
|
||||
mlflow_client,
|
||||
mock_mlflow_models,
|
||||
):
|
||||
"""A flavor model whose artifact carries `code` downloads that directory — and only it.
|
||||
|
||||
The flavor loaders read the model URI themselves, so nothing else is fetched through the
|
||||
client: the `code` directory is the whole of what this workflow downloads.
|
||||
"""
|
||||
mlflow_client.list_artifacts.side_effect = lambda _run_id, artifact_path=None: (
|
||||
[MagicMock(path=f'{artifact_path}/code')] if artifact_path else []
|
||||
)
|
||||
|
||||
client = temporal_test_env.client
|
||||
model_id = 275
|
||||
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
||||
input_data = build_input_data(model_id, 'test_model')
|
||||
|
||||
await start_and_await_workflow(
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-code-only')
|
||||
)
|
||||
|
||||
assert_prediction(postgres_engine, model_id, prediction=0.5)
|
||||
|
||||
downloaded = [loaded.args[1] for loaded in mlflow_client.download_artifacts.call_args_list]
|
||||
assert sorted(downloaded) == ['data_model/code', 'prediction_model/code']
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_prediction_workflow_joblib_flavor_failure_reports_single_cause(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
joblib_flavor_artifact_without_pickle,
|
||||
mlflow_client,
|
||||
mock_logger,
|
||||
):
|
||||
"""The joblib flavor failing stops the workflow with one cause, not a chain of two."""
|
||||
mlflow_client.download_artifacts.side_effect = lambda _run_id, _artifact_name: (
|
||||
joblib_flavor_artifact_without_pickle
|
||||
)
|
||||
|
||||
client = temporal_test_env.client
|
||||
model_id = 273
|
||||
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
||||
input_data = build_input_data(model_id, JOBLIB_FLAVOR_MODEL_NAME, predict_flavor='joblib')
|
||||
|
||||
await start_and_await_workflow(
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-flavor-failed')
|
||||
)
|
||||
|
||||
assert_stop(postgres_engine, model_id)
|
||||
|
||||
logged = logged_messages(mock_logger, 'debug')
|
||||
assert 'No .pkl file found in prediction_model artifact' in logged
|
||||
assert 'Original error' not in logged
|
||||
@@ -4,12 +4,13 @@ End-to-end tests for PredictionsBatch workflow - Main workflow scenarios.
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
from temporalio.testing import WorkflowEnvironment
|
||||
from temporalio.worker import Worker
|
||||
|
||||
from e2e.helpers import make_workflow_id, start_and_await_workflow
|
||||
import pytest
|
||||
|
||||
from e2e.helpers import load_scenario_input, make_workflow_id, start_and_await_workflow
|
||||
from laborious.activities.activities import Activities
|
||||
from laborious.workflows.predictions_batch import PredictionsBatch
|
||||
|
||||
@@ -26,9 +27,9 @@ async def test_scenario_1_1_1_happy_path_complete_success(
|
||||
client = temporal_test_env.client
|
||||
|
||||
with postgres_engine.begin() as conn:
|
||||
conn.execute(text('DELETE FROM predictions_schema.laborious_data WHERE model_id = 123'))
|
||||
conn.execute(text('DELETE FROM sientia_data.laborious_data WHERE model_id = 123'))
|
||||
insert_sql = """
|
||||
INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
|
||||
INSERT INTO sientia_data.laborious_data (model_id, variable, value, timestamp, created_at)
|
||||
VALUES
|
||||
(123, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
|
||||
(123, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'),
|
||||
@@ -36,43 +37,7 @@ async def test_scenario_1_1_1_happy_path_complete_success(
|
||||
"""
|
||||
conn.execute(text(insert_sql))
|
||||
|
||||
input_data = {
|
||||
'metadata': {
|
||||
'metadata': {
|
||||
'model_id': 123,
|
||||
'model_name': 'test_model',
|
||||
'schedule_name': 'test-schedule',
|
||||
'workflow_name': 'predictions_batch',
|
||||
}
|
||||
},
|
||||
'schedule_name': 'test-schedule',
|
||||
'model_name': 'test_model',
|
||||
'model_id': 123,
|
||||
'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 123',
|
||||
'schema': 'predictions_schema',
|
||||
'table_name': 'predictions',
|
||||
'transform_table_name': 'transformed_data',
|
||||
'input_filters': {
|
||||
'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}},
|
||||
},
|
||||
'mlflow_transform_filters': {
|
||||
'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
|
||||
},
|
||||
'mlflow_predict_filters': {
|
||||
'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
|
||||
},
|
||||
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
|
||||
'opc_output_config': {},
|
||||
'pi_web_api_output_config': {},
|
||||
'save_transform': True,
|
||||
'prediction_store_policy': 'lts:1',
|
||||
'model_config': {
|
||||
'retention_minutes': 0,
|
||||
'transform_flavor': 'sklearn',
|
||||
'predict_flavor': 'sklearn',
|
||||
},
|
||||
'datetime_columns': ['timestamp', 'created_at'],
|
||||
}
|
||||
input_data = load_scenario_input('main_happy_path.json', model_id=123)
|
||||
|
||||
await start_and_await_workflow(
|
||||
client,
|
||||
@@ -81,7 +46,7 @@ async def test_scenario_1_1_1_happy_path_complete_success(
|
||||
make_workflow_id('test-predictions-batch'),
|
||||
)
|
||||
|
||||
schema_name = 'predictions_schema'
|
||||
schema_name = 'sientia_data'
|
||||
with postgres_engine.connect() as conn:
|
||||
result_query = conn.execute(
|
||||
text(
|
||||
@@ -125,42 +90,7 @@ async def test_scenario_1_2_1_sql_query_execution_error(
|
||||
"""Invalid SQL: workflow may complete with early exit; no prediction rows."""
|
||||
client = temporal_test_env.client
|
||||
|
||||
input_data = {
|
||||
'metadata': {
|
||||
'metadata': {
|
||||
'model_id': 128,
|
||||
'model_name': 'test_model',
|
||||
'schedule_name': 'test-schedule',
|
||||
'workflow_name': 'predictions_batch',
|
||||
}
|
||||
},
|
||||
'schedule_name': 'test-schedule',
|
||||
'model_name': 'test_model',
|
||||
'model_id': 128,
|
||||
'query': 'SELECT * FROM nonexistent_table WHERE invalid_syntax =',
|
||||
'schema': 'predictions_schema',
|
||||
'table_name': 'predictions',
|
||||
'transform_table_name': 'transformed_data',
|
||||
'input_filters': {
|
||||
'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}},
|
||||
},
|
||||
'mlflow_transform_filters': {
|
||||
'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
|
||||
},
|
||||
'mlflow_predict_filters': {
|
||||
'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
|
||||
},
|
||||
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
|
||||
'opc_output_config': {},
|
||||
'pi_web_api_output_config': {},
|
||||
'save_transform': True,
|
||||
'prediction_store_policy': 'lts:1',
|
||||
'model_config': {
|
||||
'retention_minutes': 0,
|
||||
'transform_flavor': 'sklearn',
|
||||
'predict_flavor': 'sklearn',
|
||||
},
|
||||
}
|
||||
input_data = load_scenario_input('main_sql_error.json', model_id=128)
|
||||
|
||||
await start_and_await_workflow(
|
||||
client,
|
||||
@@ -171,7 +101,7 @@ async def test_scenario_1_2_1_sql_query_execution_error(
|
||||
|
||||
with postgres_engine.connect() as conn:
|
||||
count = conn.execute(
|
||||
text('SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 128')
|
||||
text('SELECT COUNT(*) FROM sientia_data.predictions WHERE model_id = 128')
|
||||
).scalar()
|
||||
assert count == 0
|
||||
|
||||
@@ -187,22 +117,7 @@ async def test_scenario_1_2_2_missing_required_parameters(
|
||||
"""Missing query: workflow does not produce predictions and is terminated explicitly."""
|
||||
client = temporal_test_env.client
|
||||
|
||||
input_data = {
|
||||
'metadata': {
|
||||
'metadata': {
|
||||
'model_id': 129,
|
||||
'model_name': 'test_model',
|
||||
'schedule_name': 'test-schedule',
|
||||
'workflow_name': 'predictions_batch',
|
||||
}
|
||||
},
|
||||
'schedule_name': 'test-schedule',
|
||||
'model_name': 'test_model',
|
||||
'model_id': 129,
|
||||
'schema': 'predictions_schema',
|
||||
'table_name': 'predictions',
|
||||
'transform_table_name': 'transformed_data',
|
||||
}
|
||||
input_data = load_scenario_input('main_missing_required.json', model_id=129)
|
||||
|
||||
handle = await client.start_workflow(
|
||||
PredictionsBatch.run,
|
||||
@@ -216,7 +131,7 @@ async def test_scenario_1_2_2_missing_required_parameters(
|
||||
|
||||
with postgres_engine.connect() as conn:
|
||||
count = conn.execute(
|
||||
text('SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 129')
|
||||
text('SELECT COUNT(*) FROM sientia_data.predictions WHERE model_id = 129')
|
||||
).scalar()
|
||||
assert count == 0
|
||||
|
||||
@@ -235,53 +150,17 @@ async def test_scenario_1_2_3_invalid_datetime_column_specification(
|
||||
client = temporal_test_env.client
|
||||
|
||||
with postgres_engine.begin() as conn:
|
||||
conn.execute(text('DELETE FROM predictions_schema.laborious_data WHERE model_id = 130'))
|
||||
conn.execute(text('DELETE FROM sientia_data.laborious_data WHERE model_id = 130'))
|
||||
conn.execute(
|
||||
text(
|
||||
"""
|
||||
INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
|
||||
INSERT INTO sientia_data.laborious_data (model_id, variable, value, timestamp, created_at)
|
||||
VALUES (130, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
|
||||
"""
|
||||
)
|
||||
)
|
||||
|
||||
input_data = {
|
||||
'metadata': {
|
||||
'metadata': {
|
||||
'model_id': 130,
|
||||
'model_name': 'test_model',
|
||||
'schedule_name': 'test-schedule',
|
||||
'workflow_name': 'predictions_batch',
|
||||
}
|
||||
},
|
||||
'schedule_name': 'test-schedule',
|
||||
'model_name': 'test_model',
|
||||
'model_id': 130,
|
||||
'query': 'SELECT timestamp, variable, value FROM predictions_schema.laborious_data WHERE model_id = 130',
|
||||
'schema': 'predictions_schema',
|
||||
'table_name': 'predictions',
|
||||
'transform_table_name': 'transformed_data',
|
||||
'input_filters': {
|
||||
'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}},
|
||||
},
|
||||
'mlflow_transform_filters': {
|
||||
'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
|
||||
},
|
||||
'mlflow_predict_filters': {
|
||||
'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
|
||||
},
|
||||
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
|
||||
'opc_output_config': {},
|
||||
'pi_web_api_output_config': {},
|
||||
'save_transform': True,
|
||||
'prediction_store_policy': 'lts:1',
|
||||
'model_config': {
|
||||
'retention_minutes': 0,
|
||||
'transform_flavor': 'sklearn',
|
||||
'predict_flavor': 'sklearn',
|
||||
},
|
||||
'datetime_columns': ['nonexistent_column'],
|
||||
}
|
||||
input_data = load_scenario_input('main_invalid_datetime.json', model_id=130)
|
||||
|
||||
handle = await client.start_workflow(
|
||||
PredictionsBatch.run,
|
||||
@@ -295,7 +174,7 @@ async def test_scenario_1_2_3_invalid_datetime_column_specification(
|
||||
|
||||
with postgres_engine.connect() as conn:
|
||||
count = conn.execute(
|
||||
text('SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 130')
|
||||
text('SELECT COUNT(*) FROM sientia_data.predictions WHERE model_id = 130')
|
||||
).scalar()
|
||||
assert count == 0
|
||||
|
||||
|
||||
@@ -3,103 +3,66 @@ End-to-end tests for PredictionsBatch workflow - Prediction Process scenarios.
|
||||
"""
|
||||
|
||||
from decimal import Decimal
|
||||
from unittest.mock import MagicMock
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
from pymongo import MongoClient
|
||||
from sqlalchemy import text
|
||||
from temporalio.testing import WorkflowEnvironment
|
||||
from temporalio.worker import Worker
|
||||
|
||||
from e2e.helpers import (
|
||||
assert_continue,
|
||||
assert_postgres_unique_violation_in_chain,
|
||||
assert_prediction,
|
||||
assert_prediction_row_count,
|
||||
assert_repeat,
|
||||
assert_stop,
|
||||
insert_sample_data,
|
||||
insert_sample_prediction,
|
||||
load_scenario_input,
|
||||
make_workflow_id,
|
||||
start_and_await_workflow,
|
||||
)
|
||||
from laborious.activities.activities import Activities
|
||||
from laborious.utils.models import minio_dataframe_payload as minio_payload_module
|
||||
from laborious.workflows.predictions_batch import PredictionsBatch
|
||||
|
||||
base_input_data = {
|
||||
'schedule_name': 'test-schedule',
|
||||
'model_name': 'test_model',
|
||||
'model_id': 201,
|
||||
'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 201',
|
||||
'schema': 'predictions_schema',
|
||||
'table_name': 'predictions',
|
||||
'transform_table_name': 'transformed_data',
|
||||
'input_filters': {
|
||||
'SPECIFIC_VARIABLES_NULL_VALUES': {
|
||||
'POLICY': 'CONTINUE',
|
||||
'CONFIG': {'variables': ['sensor_1']},
|
||||
},
|
||||
},
|
||||
'mlflow_transform_filters': {
|
||||
'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
|
||||
},
|
||||
'mlflow_predict_filters': {
|
||||
'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
|
||||
},
|
||||
'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
|
||||
'opc_output_config': {},
|
||||
'pi_web_api_output_config': {},
|
||||
'save_transform': True,
|
||||
'prediction_store_policy': 'lts:1',
|
||||
'model_config': {
|
||||
'retention_minutes': 0,
|
||||
'transform_flavor': 'sklearn',
|
||||
'predict_flavor': 'sklearn',
|
||||
},
|
||||
'datetime_columns': ['timestamp', 'created_at'],
|
||||
}
|
||||
|
||||
base_query = 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = {model_id}'
|
||||
DISTINCT_BATCH_TIMESTAMP = '2024-01-01 13:00:00+00:00'
|
||||
HISTORY_TIMESTAMP = '2024-01-01 12:00:00+00:00'
|
||||
|
||||
|
||||
def get_base_input_data(model_id):
|
||||
return {
|
||||
**base_input_data,
|
||||
'model_id': model_id,
|
||||
'query': base_query.format(model_id=model_id),
|
||||
}
|
||||
return load_scenario_input('prediction_process_base.json', model_id=model_id)
|
||||
|
||||
|
||||
def insert_sample_prediction(postgres_engine, model_id):
|
||||
with postgres_engine.begin() as conn:
|
||||
conn.execute(
|
||||
text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')
|
||||
@pytest.fixture
|
||||
def bad_data_model(mlflow_repository_stub):
|
||||
mlflow_repository_stub.stub_wrapper.transform = MagicMock(
|
||||
side_effect=Exception('Bad data model')
|
||||
)
|
||||
return mlflow_repository_stub.stub_wrapper
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def bad_predict_model(mlflow_repository_stub):
|
||||
wrapper = mlflow_repository_stub.stub_wrapper
|
||||
|
||||
def _good_transform(data):
|
||||
result = pd.DataFrame(
|
||||
{
|
||||
'feature_1': [0.234] * len(data),
|
||||
'feature_2': [0.783] * len(data),
|
||||
}
|
||||
)
|
||||
insert_sql = f"""
|
||||
INSERT INTO predictions_schema.predictions (model_id, timestamp, prediction, prediction_confidence, prediction_status, comments, response_time)
|
||||
VALUES
|
||||
({model_id}, '2024-01-01 12:00:00+00:00', 10, 0, 'Good', '', 0.1)
|
||||
"""
|
||||
conn.execute(text(insert_sql))
|
||||
return (model_id, Decimal(10), Decimal(0), 'Good')
|
||||
result.index = data.index
|
||||
return result, {}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def bad_data_model(patch_mlflow):
|
||||
model = MagicMock(predict=MagicMock(side_effect=Exception('Bad data model')))
|
||||
patch_mlflow.sklearn.load_model = MagicMock(return_value=model)
|
||||
return model
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def bad_predict_model(patch_mlflow, mock_mlflow_models):
|
||||
model = MagicMock(predict=MagicMock(side_effect=Exception('Bad predict model')))
|
||||
|
||||
def mock_sklearn_load_model(model_uri):
|
||||
if 'data_model' in model_uri or 'transform' in model_uri.lower():
|
||||
return mock_mlflow_models['transform_model']
|
||||
return model
|
||||
|
||||
patch_mlflow.sklearn = MagicMock()
|
||||
patch_mlflow.sklearn.load_model = MagicMock(side_effect=mock_sklearn_load_model)
|
||||
return model
|
||||
wrapper.transform.side_effect = _good_transform
|
||||
wrapper.predict = MagicMock(side_effect=Exception('Bad predict model'))
|
||||
return wrapper
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -109,7 +72,7 @@ async def test_scenario_2_1_1_input_gate_triggers_continue(
|
||||
temporal_worker: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
mock_mlflow_models,
|
||||
mlflow_repository_stub,
|
||||
):
|
||||
"""Input gate CONTINUE: export default prediction; MLflow transform/predict not used."""
|
||||
client = temporal_test_env.client
|
||||
@@ -120,8 +83,78 @@ async def test_scenario_2_1_1_input_gate_triggers_continue(
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-continue-policy')
|
||||
)
|
||||
assert_continue(postgres_engine, model_id)
|
||||
mock_mlflow_models['transform_model'].predict.assert_not_called()
|
||||
mock_mlflow_models['predict_model'].predict.assert_not_called()
|
||||
mlflow_repository_stub.stub_wrapper.transform.assert_not_called()
|
||||
mlflow_repository_stub.stub_wrapper.predict.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_e2e_metadata_id_reaches_postgres_when_gate_triggers(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
mongo_container,
|
||||
mlflow_repository_stub,
|
||||
):
|
||||
"""
|
||||
input_gate CONTINUE: the gate's document is written to Mongo, and the resulting
|
||||
metadata_id lands on the predictions row — independently verified on both sides.
|
||||
"""
|
||||
client = temporal_test_env.client
|
||||
model_id = 261
|
||||
# sensor_1 NULL, sensor_2 present — matches `input_filters.SPECIFIC_VARIABLES_NULL_VALUES`
|
||||
# (`variables: ["sensor_1"]`, POLICY CONTINUE) in prediction_process_base.json.
|
||||
insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
|
||||
input_data = get_base_input_data(model_id)
|
||||
await start_and_await_workflow(
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-metadata-id-continue')
|
||||
)
|
||||
assert_continue(postgres_engine, model_id)
|
||||
|
||||
with postgres_engine.connect() as conn:
|
||||
metadata_id = conn.execute(
|
||||
text(f'SELECT metadata_id FROM sientia_data.predictions WHERE model_id = {model_id}')
|
||||
).scalar()
|
||||
assert metadata_id is not None, 'Expected a metadata_id on the predictions row'
|
||||
|
||||
# Independent readback: a fresh client against the same testcontainer, not the
|
||||
# `mongo_repository`/`PredictionMetadata` instances the workflow itself wrote through.
|
||||
mongo_port = mongo_container.get_exposed_port(27017)
|
||||
verify_client = MongoClient(f'mongodb://localhost:{mongo_port}')
|
||||
try:
|
||||
document = verify_client['test_db']['prediction_metadata'].find_one({'_id': metadata_id})
|
||||
finally:
|
||||
verify_client.close()
|
||||
assert document is not None, f'Expected a prediction_metadata document for id {metadata_id}'
|
||||
assert document['filters'] == [
|
||||
{'name': 'SPECIFIC_VARIABLES_NULL_VALUES', 'policy': 'CONTINUE', 'counter': 1}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_e2e_metadata_id_is_null_when_nothing_triggers(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
):
|
||||
"""Clean prediction, no gate triggers: predictions.metadata_id stays NULL."""
|
||||
client = temporal_test_env.client
|
||||
model_id = 262
|
||||
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
||||
input_data = get_base_input_data(model_id)
|
||||
await start_and_await_workflow(
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-metadata-id-none')
|
||||
)
|
||||
assert_prediction(postgres_engine, model_id, prediction=0.5, prediction_confidence=0)
|
||||
|
||||
with postgres_engine.connect() as conn:
|
||||
metadata_id = conn.execute(
|
||||
text(f'SELECT metadata_id FROM sientia_data.predictions WHERE model_id = {model_id}')
|
||||
).scalar()
|
||||
assert metadata_id is None, f'Expected metadata_id to be NULL, got {metadata_id!r}'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -131,7 +164,7 @@ async def test_scenario_2_1_2_input_gate_triggers_stop(
|
||||
temporal_worker: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
mock_mlflow_models,
|
||||
mlflow_repository_stub,
|
||||
):
|
||||
"""Input gate STOP: no export, no MLflow."""
|
||||
client = temporal_test_env.client
|
||||
@@ -143,30 +176,60 @@ async def test_scenario_2_1_2_input_gate_triggers_stop(
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-input-stop')
|
||||
)
|
||||
assert_stop(postgres_engine, model_id)
|
||||
mock_mlflow_models['transform_model'].predict.assert_not_called()
|
||||
mlflow_repository_stub.stub_wrapper.transform.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_scenario_2_1_3_input_gate_triggers_repeat(
|
||||
async def test_scenario_2_1_3_input_gate_repeat_batch_timestamp_equals_history_fails(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
mock_mlflow_models,
|
||||
mlflow_repository_stub,
|
||||
):
|
||||
"""Input gate REPEAT with existing history."""
|
||||
"""
|
||||
REPEAT uses ``last_timestamp`` from the batch payload as the new row's ``timestamp``.
|
||||
When it equals the only historical prediction row, Postgres rejects the duplicate key.
|
||||
"""
|
||||
client = temporal_test_env.client
|
||||
model_id = 213
|
||||
insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
|
||||
data = insert_sample_prediction(postgres_engine, model_id)
|
||||
insert_sample_data(postgres_engine, model_id, ['NULL', 78.2], data_timestamp=HISTORY_TIMESTAMP)
|
||||
insert_sample_prediction(postgres_engine, model_id, prediction_timestamp=HISTORY_TIMESTAMP)
|
||||
input_data = get_base_input_data(model_id)
|
||||
input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'REPEAT'
|
||||
with pytest.raises(Exception) as excinfo:
|
||||
await start_and_await_workflow(
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-input-repeat-collision')
|
||||
)
|
||||
assert_postgres_unique_violation_in_chain(excinfo.value)
|
||||
assert_prediction_row_count(postgres_engine, model_id, 1)
|
||||
mlflow_repository_stub.stub_wrapper.transform.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_scenario_2_1_3_input_gate_repeat_distinct_batch_timestamp_inserts_second_row(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
mlflow_repository_stub,
|
||||
):
|
||||
"""REPEAT succeeds when batch ``last_timestamp`` differs from the historical prediction row."""
|
||||
client = temporal_test_env.client
|
||||
model_id = 2131
|
||||
insert_sample_data(
|
||||
postgres_engine, model_id, ['NULL', 78.2], data_timestamp=DISTINCT_BATCH_TIMESTAMP
|
||||
)
|
||||
data = insert_sample_prediction(postgres_engine, model_id, prediction_timestamp=HISTORY_TIMESTAMP)
|
||||
input_data = get_base_input_data(model_id)
|
||||
input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'REPEAT'
|
||||
await start_and_await_workflow(
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-input-repeat')
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-input-repeat-ok')
|
||||
)
|
||||
assert_repeat(postgres_engine, model_id, data)
|
||||
mock_mlflow_models['transform_model'].predict.assert_not_called()
|
||||
mlflow_repository_stub.stub_wrapper.transform.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -182,9 +245,7 @@ async def test_scenario_2_1_4_input_gate_repeat_without_prior_prediction(
|
||||
model_id = 214
|
||||
insert_sample_data(postgres_engine, model_id, ['NULL', 78.2])
|
||||
with postgres_engine.begin() as conn:
|
||||
conn.execute(
|
||||
text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')
|
||||
)
|
||||
conn.execute(text(f'DELETE FROM sientia_data.predictions WHERE model_id = {model_id}'))
|
||||
input_data = get_base_input_data(model_id)
|
||||
input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'REPEAT'
|
||||
await start_and_await_workflow(
|
||||
@@ -214,7 +275,7 @@ async def test_scenario_2_2_1_transform_gate_triggers_continue(
|
||||
postgres_engine=postgres_engine,
|
||||
model_id=model_id,
|
||||
prediction_confidence=Decimal(10),
|
||||
comments='Unknown MLFlow API error',
|
||||
comments='Bad data model',
|
||||
)
|
||||
|
||||
|
||||
@@ -226,7 +287,7 @@ async def test_scenario_2_2_2_transform_gate_triggers_stop(
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
bad_data_model,
|
||||
mock_mlflow_models,
|
||||
mlflow_repository_stub,
|
||||
):
|
||||
client = temporal_test_env.client
|
||||
model_id = 222
|
||||
@@ -237,12 +298,12 @@ async def test_scenario_2_2_2_transform_gate_triggers_stop(
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-stop')
|
||||
)
|
||||
assert_stop(postgres_engine, model_id)
|
||||
mock_mlflow_models['predict_model'].predict.assert_not_called()
|
||||
mlflow_repository_stub.stub_wrapper.predict.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_scenario_2_2_3_transform_gate_triggers_repeat(
|
||||
async def test_scenario_2_2_3_transform_gate_repeat_batch_timestamp_equals_history_fails(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker: Worker,
|
||||
test_activities: Activities,
|
||||
@@ -251,12 +312,37 @@ async def test_scenario_2_2_3_transform_gate_triggers_repeat(
|
||||
):
|
||||
client = temporal_test_env.client
|
||||
model_id = 223
|
||||
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
||||
data = insert_sample_prediction(postgres_engine, model_id)
|
||||
insert_sample_data(postgres_engine, model_id, [60.0, 78.2], data_timestamp=HISTORY_TIMESTAMP)
|
||||
insert_sample_prediction(postgres_engine, model_id, prediction_timestamp=HISTORY_TIMESTAMP)
|
||||
input_data = get_base_input_data(model_id)
|
||||
input_data['mlflow_transform_filters']['API_ERROR']['POLICY'] = 'REPEAT'
|
||||
with pytest.raises(Exception) as excinfo:
|
||||
await start_and_await_workflow(
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-repeat-collision')
|
||||
)
|
||||
assert_postgres_unique_violation_in_chain(excinfo.value)
|
||||
assert_prediction_row_count(postgres_engine, model_id, 1)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_scenario_2_2_3_transform_gate_repeat_distinct_batch_timestamp_inserts_second_row(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
bad_data_model,
|
||||
):
|
||||
client = temporal_test_env.client
|
||||
model_id = 2231
|
||||
insert_sample_data(
|
||||
postgres_engine, model_id, [60.0, 78.2], data_timestamp=DISTINCT_BATCH_TIMESTAMP
|
||||
)
|
||||
data = insert_sample_prediction(postgres_engine, model_id, prediction_timestamp=HISTORY_TIMESTAMP)
|
||||
input_data = get_base_input_data(model_id)
|
||||
input_data['mlflow_transform_filters']['API_ERROR']['POLICY'] = 'REPEAT'
|
||||
await start_and_await_workflow(
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-repeat')
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-repeat-ok')
|
||||
)
|
||||
assert_repeat(postgres_engine, model_id, data)
|
||||
|
||||
@@ -268,19 +354,20 @@ async def test_scenario_2_2_4_transform_content_gate_nan_values_stop(
|
||||
temporal_worker: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
mock_mlflow_models,
|
||||
mlflow_repository_stub,
|
||||
):
|
||||
"""mlflow_content_gate triggers STOP when transform output is all NaN (NAN_VALUES filter)."""
|
||||
client = temporal_test_env.client
|
||||
model_id = 224
|
||||
|
||||
def all_nan_transform(data):
|
||||
num_rows = max(len(data), 1) if hasattr(data, '__len__') else 1
|
||||
result = pd.DataFrame({'feature_1': [np.nan] * num_rows, 'feature_2': [np.nan] * num_rows})
|
||||
result = pd.DataFrame(
|
||||
{'feature_1': [np.nan] * len(data), 'feature_2': [np.nan] * len(data)}
|
||||
)
|
||||
result.index = data.index
|
||||
return result
|
||||
return result, {}
|
||||
|
||||
mock_mlflow_models['transform_model'].predict = MagicMock(side_effect=all_nan_transform)
|
||||
mlflow_repository_stub.stub_wrapper.transform = MagicMock(side_effect=all_nan_transform)
|
||||
|
||||
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
||||
input_data = get_base_input_data(model_id)
|
||||
@@ -292,7 +379,7 @@ async def test_scenario_2_2_4_transform_content_gate_nan_values_stop(
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-content-stop')
|
||||
)
|
||||
assert_stop(postgres_engine, model_id)
|
||||
mock_mlflow_models['predict_model'].predict.assert_not_called()
|
||||
mlflow_repository_stub.stub_wrapper.predict.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -316,7 +403,7 @@ async def test_scenario_2_3_1_predict_gate_triggers_continue(
|
||||
postgres_engine=postgres_engine,
|
||||
model_id=model_id,
|
||||
prediction_confidence=Decimal(10),
|
||||
comments='Unknown MLFlow API error',
|
||||
comments='Bad predict model',
|
||||
)
|
||||
|
||||
|
||||
@@ -342,7 +429,7 @@ async def test_scenario_2_3_2_predict_gate_triggers_stop(
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_scenario_2_3_3_predict_gate_triggers_repeat(
|
||||
async def test_scenario_2_3_3_predict_gate_repeat_batch_timestamp_equals_history_fails(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker: Worker,
|
||||
test_activities: Activities,
|
||||
@@ -351,13 +438,39 @@ async def test_scenario_2_3_3_predict_gate_triggers_repeat(
|
||||
):
|
||||
client = temporal_test_env.client
|
||||
model_id = 233
|
||||
insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
|
||||
data = insert_sample_prediction(postgres_engine, model_id)
|
||||
insert_sample_data(postgres_engine, model_id, [23.5, 78.2], data_timestamp=HISTORY_TIMESTAMP)
|
||||
insert_sample_prediction(postgres_engine, model_id, prediction_timestamp=HISTORY_TIMESTAMP)
|
||||
input_data = get_base_input_data(model_id)
|
||||
input_data['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'REPEAT'
|
||||
input_data['path_priority'] = ['REPEAT', 'STOP', 'CONTINUE']
|
||||
with pytest.raises(Exception) as excinfo:
|
||||
await start_and_await_workflow(
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-predict-repeat-collision')
|
||||
)
|
||||
assert_postgres_unique_violation_in_chain(excinfo.value)
|
||||
assert_prediction_row_count(postgres_engine, model_id, 1)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_scenario_2_3_3_predict_gate_repeat_distinct_batch_timestamp_inserts_second_row(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
bad_predict_model,
|
||||
):
|
||||
client = temporal_test_env.client
|
||||
model_id = 2331
|
||||
insert_sample_data(
|
||||
postgres_engine, model_id, [23.5, 78.2], data_timestamp=DISTINCT_BATCH_TIMESTAMP
|
||||
)
|
||||
data = insert_sample_prediction(postgres_engine, model_id, prediction_timestamp=HISTORY_TIMESTAMP)
|
||||
input_data = get_base_input_data(model_id)
|
||||
input_data['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'REPEAT'
|
||||
input_data['path_priority'] = ['REPEAT', 'STOP', 'CONTINUE']
|
||||
await start_and_await_workflow(
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-predict-repeat')
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-predict-repeat-ok')
|
||||
)
|
||||
assert_repeat(postgres_engine, model_id, data)
|
||||
|
||||
@@ -374,12 +487,73 @@ async def test_scenario_2_4_1_input_empty_data_stop(
|
||||
client = temporal_test_env.client
|
||||
model_id = 241
|
||||
with postgres_engine.begin() as conn:
|
||||
conn.execute(
|
||||
text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}')
|
||||
)
|
||||
conn.execute(text(f'DELETE FROM sientia_data.laborious_data WHERE model_id = {model_id}'))
|
||||
input_data = get_base_input_data(model_id)
|
||||
input_data['input_filters'] = {'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}}
|
||||
await start_and_await_workflow(
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-empty-data-stop')
|
||||
)
|
||||
assert_stop(postgres_engine, model_id)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_scenario_2_4_1_priority_conflict_resolution(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
bad_data_model,
|
||||
):
|
||||
"""Conflicting filter outputs must honor configured path_priority order."""
|
||||
client = temporal_test_env.client
|
||||
model_id = 242
|
||||
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
||||
input_data = get_base_input_data(model_id)
|
||||
input_data['mlflow_transform_filters'] = {'API_ERROR': {'POLICY': 'CONTINUE', 'CONFIG': {}}}
|
||||
input_data['mlflow_predict_filters'] = {'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}}
|
||||
input_data['path_priority'] = ['STOP', 'CONTINUE', 'REPEAT']
|
||||
|
||||
await start_and_await_workflow(
|
||||
client, PredictionsBatch.run, input_data, make_workflow_id('test-priority-conflict')
|
||||
)
|
||||
assert_continue(
|
||||
postgres_engine=postgres_engine,
|
||||
model_id=model_id,
|
||||
prediction_confidence=Decimal(10),
|
||||
comments='Bad data model',
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_e2e_request_predict_inline_minio_payload_with_datetimeindex(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
mlflow_repository_stub,
|
||||
):
|
||||
"""
|
||||
High offload threshold forces inline tabular dicts; ``DatetimeIndex`` must serialize as JSON
|
||||
(string index keys via ``MinioDataFramePayload.from_dataframe``) so ``request_predict`` completes.
|
||||
"""
|
||||
client = temporal_test_env.client
|
||||
model_id = 252
|
||||
with postgres_engine.begin() as conn:
|
||||
conn.execute(text(f'DELETE FROM sientia_data.laborious_data WHERE model_id = {model_id}'))
|
||||
conn.execute(text(f'DELETE FROM sientia_data.predictions WHERE model_id = {model_id}'))
|
||||
conn.execute(text(f'DELETE FROM sientia_data.transformed_data WHERE model_id = {model_id}'))
|
||||
insert_sample_data(postgres_engine, model_id, [60.0, 78.2])
|
||||
input_data = get_base_input_data(model_id)
|
||||
|
||||
with patch.object(minio_payload_module, 'OFFLOAD_THRESHOLD_BYTES', 10**9):
|
||||
await start_and_await_workflow(
|
||||
client,
|
||||
PredictionsBatch.run,
|
||||
input_data,
|
||||
make_workflow_id('test-predict-inline-json-datetimeindex'),
|
||||
)
|
||||
|
||||
assert_prediction(postgres_engine, model_id, prediction=0.5, prediction_confidence=0)
|
||||
mlflow_repository_stub.stub_wrapper.predict.assert_called()
|
||||
|
||||
@@ -1,128 +0,0 @@
|
||||
"""
|
||||
End-to-end test for retraining a model served through the explicit `joblib` flavor.
|
||||
|
||||
Retraining is not part of the `predictions_batch` workflow, so it gets its own module. The test
|
||||
is a write→read round trip in one run: `retrain_model` loads the production artifact through the
|
||||
joblib flavor, `log_model` writes the retrained model back as `model.pkl` + `code/`, and the
|
||||
written artifact is then loaded again through the same flavor. Neither `joblib.dump` nor
|
||||
`joblib.load` is mocked; only MLflow's own I/O is.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
import shutil
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
from e2e.conftest import JOBLIB_FLAVOR_CODE_MODULE, JOBLIB_FLAVOR_MODEL_NAME
|
||||
from laborious.activities.activities import Activities
|
||||
|
||||
RETRAIN_INDEX = [
|
||||
'2024-01-01 12:00:00+0000',
|
||||
'2024-01-01 12:01:00+0000',
|
||||
'2024-01-01 12:02:00+0000',
|
||||
]
|
||||
|
||||
METADATA = {
|
||||
'model_id': 'test_model',
|
||||
'model_name': JOBLIB_FLAVOR_MODEL_NAME,
|
||||
'workflow_name': 'test_workflow',
|
||||
'schema_name': 'test_schedule',
|
||||
}
|
||||
|
||||
|
||||
def build_retrain_data() -> pd.DataFrame:
|
||||
"""Training frame indexed by timestamp strings, carrying the target column."""
|
||||
return pd.DataFrame(
|
||||
{
|
||||
'sensor_1': [60.0, 61.0, 62.0],
|
||||
'sensor_2': [78.2, 78.4, 78.6],
|
||||
'target': [1.0, 2.0, 3.0],
|
||||
},
|
||||
index=RETRAIN_INDEX,
|
||||
)
|
||||
|
||||
|
||||
class FakeArtifactStore:
|
||||
"""Minimal stand-in for the run's artifact tree, so the written layout is inspectable."""
|
||||
|
||||
def __init__(self, root: Path):
|
||||
self.root = root
|
||||
|
||||
def log_artifact(self, local_path, artifact_path=None):
|
||||
if artifact_path is None:
|
||||
# Retraining also logs the CSVs at the run root; irrelevant to this test.
|
||||
return
|
||||
target = self.root / artifact_path
|
||||
target.mkdir(parents=True, exist_ok=True)
|
||||
shutil.copy(local_path, target / Path(local_path).name)
|
||||
|
||||
def log_artifacts(self, local_dir, artifact_path=None):
|
||||
target = self.root / artifact_path if artifact_path else self.root
|
||||
shutil.copytree(local_dir, target, dirs_exist_ok=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_retrain_joblib_model_writes_pickle_and_code_and_reloads(
|
||||
test_activities: Activities,
|
||||
joblib_flavor_artifacts,
|
||||
mlflow_client,
|
||||
patch_mlflow,
|
||||
tmp_path,
|
||||
):
|
||||
"""A joblib retrain writes `model.pkl` + `code/`, and the joblib flavor reads it back."""
|
||||
repository = test_activities.model_monitoring_repository
|
||||
|
||||
store = FakeArtifactStore(tmp_path / 'run-retrained')
|
||||
patch_mlflow.log_artifact.side_effect = store.log_artifact
|
||||
patch_mlflow.log_artifacts.side_effect = store.log_artifacts
|
||||
patch_mlflow.get_run.return_value.data.params = {}
|
||||
patch_mlflow.start_run.return_value.__enter__.return_value.info.run_id = 'run-retrained'
|
||||
|
||||
output = await repository.retrain_model(
|
||||
data=build_retrain_data(),
|
||||
model_name=JOBLIB_FLAVOR_MODEL_NAME,
|
||||
model_config={
|
||||
'target': 'target',
|
||||
'transform_flavor': 'joblib',
|
||||
'predict_flavor': 'joblib',
|
||||
},
|
||||
metadata=METADATA,
|
||||
)
|
||||
|
||||
assert output['success'] is True, output.get('traceback')
|
||||
|
||||
# Both models were loaded through the joblib flavor, not through a flavor loader.
|
||||
downloaded = [loaded.args[1] for loaded in mlflow_client.download_artifacts.call_args_list]
|
||||
assert sorted(downloaded) == ['data_model', 'prediction_model']
|
||||
assert patch_mlflow.sklearn.load_model.call_count == 0
|
||||
assert patch_mlflow.pyfunc.load_model.call_count == 0
|
||||
|
||||
written = store.root / 'prediction_model'
|
||||
assert (written / 'model.pkl').is_file()
|
||||
assert (written / 'code' / f'{JOBLIB_FLAVOR_CODE_MODULE}.py').is_file()
|
||||
assert not (written / 'MLmodel').exists()
|
||||
assert (store.root / 'data_model' / 'model.pkl').is_file()
|
||||
|
||||
# Drop the module the first load put in place: the reload must resolve it from the `code/`
|
||||
# the retrain wrote, not from what is already imported.
|
||||
for code_dir in (artifact['code_dir'] for artifact in joblib_flavor_artifacts.values()):
|
||||
if code_dir in sys.path:
|
||||
sys.path.remove(code_dir)
|
||||
sys.modules.pop(JOBLIB_FLAVOR_CODE_MODULE, None)
|
||||
importlib.invalidate_caches()
|
||||
|
||||
mlflow_client.download_artifacts.side_effect = lambda _run_id, _artifact_name: str(written)
|
||||
|
||||
reloaded, _local_dir = await repository.load_predict_model(
|
||||
JOBLIB_FLAVOR_MODEL_NAME, METADATA, 'joblib'
|
||||
)
|
||||
|
||||
features = pd.DataFrame({'feature_1': [0.234], 'feature_2': [0.783]})
|
||||
expected = joblib_flavor_artifacts['prediction_model']['model'].predict(features)
|
||||
|
||||
assert reloaded.predict(features) == pytest.approx(expected)
|
||||
assert str(written / 'code') in sys.path
|
||||
333
e2e/test_simple_metrics.py
Normal file
333
e2e/test_simple_metrics.py
Normal file
@@ -0,0 +1,333 @@
|
||||
"""
|
||||
End-to-end tests for the SimpleMetrics workflow.
|
||||
|
||||
Coverage focus:
|
||||
|
||||
- Happy path computes rmse/mse/mae/r2 from predictions joined against ``laborious_data``
|
||||
and persists rows to ``sientia_data.simple_metrics`` with all required columns.
|
||||
- Subset metric selection (only rmse) writes exactly the requested rows.
|
||||
- Zero-variance target produces ``r2=0`` per division-by-zero guard.
|
||||
- Empty join (no overlapping data) short-circuits without persisting anything.
|
||||
"""
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from decimal import Decimal
|
||||
|
||||
import math
|
||||
import pytest
|
||||
from sqlalchemy import text
|
||||
from temporalio.testing import WorkflowEnvironment
|
||||
from temporalio.worker import Worker
|
||||
|
||||
from e2e.helpers import load_scenario_input, make_workflow_id, start_and_await_workflow
|
||||
from laborious.activities.activities import Activities
|
||||
from laborious.workflows.simple_metrics import SimpleMetrics
|
||||
|
||||
# Columns defined by the production DDL for ``sientia_data.simple_metrics``.
|
||||
EXPECTED_SIMPLE_METRICS_COLUMNS = [
|
||||
'id',
|
||||
'model_id',
|
||||
'metric',
|
||||
'value',
|
||||
'timestamp',
|
||||
'data_size',
|
||||
'interval_minutes',
|
||||
'created_at',
|
||||
]
|
||||
|
||||
# ``timestamp`` is now nullable per the new DDL (production code may write it
|
||||
# null when the upstream data has no usable instant); skip the non-null check
|
||||
# for it while still validating presence.
|
||||
NULLABLE_SIMPLE_METRICS_COLUMNS = {'timestamp'}
|
||||
|
||||
|
||||
def _simple_metrics_input(model_id: int, **overrides) -> dict:
|
||||
"""Load and override the simple-metrics base scenario."""
|
||||
payload = load_scenario_input('simple_metrics_base.json', model_id=model_id)
|
||||
payload.update(overrides)
|
||||
return payload
|
||||
|
||||
|
||||
def _seed_predictions_and_targets(
|
||||
postgres_engine,
|
||||
model_id: int,
|
||||
pairs: list[tuple[float, float]],
|
||||
target_name: str = 'sensor_target',
|
||||
offset_minutes: int = 6,
|
||||
) -> list[str]:
|
||||
"""
|
||||
Insert matching prediction/target rows used by the SimpleMetrics SQL JOIN.
|
||||
|
||||
For each ``(prediction, target)`` pair we write a row in ``predictions`` and
|
||||
a matching row in ``laborious_data`` with ``variable=target_name`` so the
|
||||
inner join in the workflow query yields one row per pair.
|
||||
|
||||
Args:
|
||||
- postgres_engine: SQLAlchemy engine bound to the test container.
|
||||
- model_id: Model id stamped on every row.
|
||||
- pairs: ``(prediction, target)`` pairs, one per minute.
|
||||
- target_name: Variable name in ``laborious_data`` representing the target.
|
||||
- offset_minutes: Earliest row sits this many minutes ago so timestamps fall
|
||||
inside the workflow's recent-data window.
|
||||
|
||||
Return:
|
||||
List of timestamp strings written for the inserted rows.
|
||||
"""
|
||||
base = datetime.now(timezone.utc).replace(second=0, microsecond=0) - timedelta(
|
||||
minutes=offset_minutes
|
||||
)
|
||||
timestamps = [
|
||||
(base + timedelta(minutes=i)).strftime('%Y-%m-%d %H:%M:%S%z')
|
||||
for i in range(len(pairs))
|
||||
]
|
||||
|
||||
prediction_rows = []
|
||||
target_rows = []
|
||||
for index, (prediction, target_value) in enumerate(pairs):
|
||||
ts = timestamps[index]
|
||||
prediction_rows.append(
|
||||
f"({model_id}, {prediction}, 0, 0, 'Good', '{ts}', '{ts}')"
|
||||
)
|
||||
target_rows.append(
|
||||
f"({model_id}, '{target_name}', {target_value}, '{ts}', '{ts}')"
|
||||
)
|
||||
|
||||
with postgres_engine.begin() as conn:
|
||||
conn.execute(text(f'DELETE FROM sientia_data.predictions WHERE model_id = {model_id}'))
|
||||
conn.execute(text(f'DELETE FROM sientia_data.laborious_data WHERE model_id = {model_id}'))
|
||||
# The SimpleMetrics SQL JOIN only filters ``predictions.model_id``; it does
|
||||
# NOT filter ``laborious_data.model_id`` (see ``e2e/CODE_ISSUES.md`` issue
|
||||
# SM-1). Without this cross-model cleanup, a previous test's target rows
|
||||
# under the same variable name would join into this test's predictions
|
||||
# whenever timestamps happened to overlap.
|
||||
conn.execute(
|
||||
text(
|
||||
"DELETE FROM sientia_data.laborious_data "
|
||||
"WHERE variable IN (:sensor_default, :target_name) "
|
||||
"AND timestamp >= NOW() - INTERVAL '120 minutes'"
|
||||
),
|
||||
{'sensor_default': 'sensor_target', 'target_name': target_name},
|
||||
)
|
||||
if prediction_rows:
|
||||
conn.execute(
|
||||
text(
|
||||
'INSERT INTO sientia_data.predictions '
|
||||
'(model_id, prediction, prediction_confidence, response_time, '
|
||||
'prediction_status, "timestamp", created_at) VALUES '
|
||||
+ ', '.join(prediction_rows)
|
||||
)
|
||||
)
|
||||
conn.execute(
|
||||
text(
|
||||
'INSERT INTO sientia_data.laborious_data '
|
||||
'(model_id, variable, value, "timestamp", created_at) VALUES '
|
||||
+ ', '.join(target_rows)
|
||||
)
|
||||
)
|
||||
return timestamps
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_simple_metrics_happy_path_persists_all_metrics_and_columns(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker_simple_metrics: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
):
|
||||
"""
|
||||
Scenario S.1.1: rmse/mse/mae/r2 are calculated from a deterministic
|
||||
prediction/target pair set and written one row per metric. Every column
|
||||
expected by ``sientia_data.simple_metrics`` must be populated (except the
|
||||
nullable ``timestamp`` column) and the numerical values must match
|
||||
closed-form expectations.
|
||||
"""
|
||||
client = temporal_test_env.client
|
||||
model_id = 511
|
||||
|
||||
pairs = [
|
||||
(1.0, 2.0),
|
||||
(2.0, 4.0),
|
||||
(3.0, 5.0),
|
||||
(4.0, 9.0),
|
||||
(5.0, 12.0),
|
||||
]
|
||||
diffs = [target - prediction for prediction, target in pairs]
|
||||
n = len(diffs)
|
||||
expected_rmse = math.sqrt(sum(d * d for d in diffs) / n)
|
||||
expected_mse = sum(d * d for d in diffs) / n
|
||||
expected_mae = sum(abs(d) for d in diffs) / n
|
||||
target_mean = sum(t for _, t in pairs) / n
|
||||
ss_res = sum((target - prediction) ** 2 for prediction, target in pairs)
|
||||
ss_tot = sum((t - target_mean) ** 2 for _, t in pairs)
|
||||
expected_r2 = 1.0 - (ss_res / ss_tot)
|
||||
|
||||
_seed_predictions_and_targets(postgres_engine, model_id=model_id, pairs=pairs)
|
||||
|
||||
input_data = _simple_metrics_input(model_id)
|
||||
await start_and_await_workflow(
|
||||
client,
|
||||
SimpleMetrics.run,
|
||||
input_data,
|
||||
make_workflow_id('test-simple-metrics-happy'),
|
||||
)
|
||||
|
||||
with postgres_engine.connect() as conn:
|
||||
rows = (
|
||||
conn.execute(
|
||||
text(
|
||||
'SELECT * FROM sientia_data.simple_metrics '
|
||||
'WHERE model_id = :m ORDER BY metric'
|
||||
),
|
||||
{'m': str(model_id)},
|
||||
)
|
||||
.mappings()
|
||||
.all()
|
||||
)
|
||||
|
||||
assert len(rows) == 4, f'Expected 4 metric rows, got {len(rows)}'
|
||||
for column in EXPECTED_SIMPLE_METRICS_COLUMNS:
|
||||
assert column in rows[0], f'Missing simple_metrics column: {column}'
|
||||
for row in rows:
|
||||
for column in EXPECTED_SIMPLE_METRICS_COLUMNS:
|
||||
if column in NULLABLE_SIMPLE_METRICS_COLUMNS:
|
||||
continue
|
||||
assert row[column] is not None, f"Column '{column}' is NULL in {dict(row)}"
|
||||
|
||||
by_metric = {row['metric']: row for row in rows}
|
||||
assert set(by_metric) == {'rmse', 'mse', 'mae', 'r2'}
|
||||
|
||||
def _decimal_close(actual, expected, places: int = 6) -> bool:
|
||||
return abs(float(actual) - expected) < 10 ** (-places)
|
||||
|
||||
assert _decimal_close(by_metric['rmse']['value'], expected_rmse)
|
||||
assert _decimal_close(by_metric['mse']['value'], expected_mse)
|
||||
assert _decimal_close(by_metric['mae']['value'], expected_mae)
|
||||
assert _decimal_close(by_metric['r2']['value'], expected_r2)
|
||||
|
||||
assert all(row['data_size'] == n for row in rows), 'data_size must equal target row count'
|
||||
assert all(row['interval_minutes'] == 60 for row in rows)
|
||||
# ``model_id`` is now ``text`` in the new DDL, so we compare with the
|
||||
# stringified test id rather than the numeric value.
|
||||
assert all(row['model_id'] == str(model_id) for row in rows)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_simple_metrics_subset_metrics_writes_only_requested_rows(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker_simple_metrics: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
):
|
||||
"""
|
||||
Scenario S.1.2: Requesting ``metrics=['rmse']`` must persist exactly one row
|
||||
with metric ``rmse`` and skip mse/mae/r2.
|
||||
"""
|
||||
client = temporal_test_env.client
|
||||
model_id = 512
|
||||
|
||||
pairs = [(1.0, 2.0), (2.0, 4.0), (3.0, 6.0)]
|
||||
_seed_predictions_and_targets(postgres_engine, model_id=model_id, pairs=pairs)
|
||||
|
||||
input_data = _simple_metrics_input(model_id, metrics=['rmse'])
|
||||
await start_and_await_workflow(
|
||||
client,
|
||||
SimpleMetrics.run,
|
||||
input_data,
|
||||
make_workflow_id('test-simple-metrics-subset'),
|
||||
)
|
||||
|
||||
with postgres_engine.connect() as conn:
|
||||
metrics = [
|
||||
r[0]
|
||||
for r in conn.execute(
|
||||
text(
|
||||
'SELECT metric FROM sientia_data.simple_metrics '
|
||||
'WHERE model_id = :m'
|
||||
),
|
||||
{'m': str(model_id)},
|
||||
).all()
|
||||
]
|
||||
assert metrics == ['rmse']
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_simple_metrics_zero_variance_target_returns_zero_r2(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker_simple_metrics: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
):
|
||||
"""
|
||||
Scenario S.2.1: When the target column has zero variance the activity must
|
||||
return ``r2 = 0`` (division-by-zero guard) and still persist all four metrics.
|
||||
"""
|
||||
client = temporal_test_env.client
|
||||
model_id = 521
|
||||
|
||||
pairs = [(0.0, 5.0), (1.0, 5.0), (2.0, 5.0), (3.0, 5.0)]
|
||||
_seed_predictions_and_targets(postgres_engine, model_id=model_id, pairs=pairs)
|
||||
|
||||
input_data = _simple_metrics_input(model_id)
|
||||
await start_and_await_workflow(
|
||||
client,
|
||||
SimpleMetrics.run,
|
||||
input_data,
|
||||
make_workflow_id('test-simple-metrics-zero-variance'),
|
||||
)
|
||||
|
||||
with postgres_engine.connect() as conn:
|
||||
r2_value = conn.execute(
|
||||
text(
|
||||
"SELECT value FROM sientia_data.simple_metrics "
|
||||
"WHERE model_id = :m AND metric = 'r2'"
|
||||
),
|
||||
{'m': str(model_id)},
|
||||
).scalar()
|
||||
assert r2_value is not None
|
||||
assert Decimal(str(r2_value)) == Decimal('0'), f'expected r2=0, got {r2_value!r}'
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.integration
|
||||
async def test_simple_metrics_no_overlapping_data_short_circuits(
|
||||
temporal_test_env: WorkflowEnvironment,
|
||||
temporal_worker_simple_metrics: Worker,
|
||||
test_activities: Activities,
|
||||
postgres_engine,
|
||||
):
|
||||
"""
|
||||
Scenario S.3.1: When the join produces no rows (no matching laborious_data
|
||||
row for the configured ``target``), the workflow returns early without
|
||||
invoking ``calculate_simple_metrics`` and writes nothing.
|
||||
"""
|
||||
client = temporal_test_env.client
|
||||
model_id = 531
|
||||
|
||||
# Insert predictions but no matching target rows for the configured variable.
|
||||
_seed_predictions_and_targets(
|
||||
postgres_engine,
|
||||
model_id=model_id,
|
||||
pairs=[(1.0, 1.0)],
|
||||
target_name='wrong_variable_name',
|
||||
)
|
||||
|
||||
input_data = _simple_metrics_input(model_id)
|
||||
await start_and_await_workflow(
|
||||
client,
|
||||
SimpleMetrics.run,
|
||||
input_data,
|
||||
make_workflow_id('test-simple-metrics-empty-join'),
|
||||
)
|
||||
|
||||
with postgres_engine.connect() as conn:
|
||||
count = conn.execute(
|
||||
text(
|
||||
'SELECT COUNT(*) FROM sientia_data.simple_metrics '
|
||||
'WHERE model_id = :m'
|
||||
),
|
||||
{'m': str(model_id)},
|
||||
).scalar()
|
||||
assert count == 0, 'Empty target data must short-circuit and skip persistence'
|
||||
@@ -1,2 +1,2 @@
|
||||
git+ssh://git@github.com/Aignosi/sientia-dataops-library.git:sientia-do
|
||||
git+ssh://git@github.com/Aignosi/sientia-mlops-library.git:sientia
|
||||
git+ssh://git@github.com/Aignosi/sientia-dataops-library.git:sientia_do
|
||||
git+ssh://git@github.com/Aignosi/sientia-model-library.git:sientia_model
|
||||
143
input_sample.json
Normal file
143
input_sample.json
Normal file
@@ -0,0 +1,143 @@
|
||||
{
|
||||
"models": [
|
||||
{
|
||||
"id": "1001",
|
||||
"name": "test-runtime",
|
||||
"active": false,
|
||||
"model_config": {
|
||||
"alias": "production",
|
||||
"retention_minutes": 60,
|
||||
"target": "Square"
|
||||
}
|
||||
}
|
||||
],
|
||||
"pipelines": [
|
||||
{
|
||||
"schedule_name": "laborious-test-runtime",
|
||||
"model_id": "1001",
|
||||
"workflow_type": "predictions_batch",
|
||||
"frequency": "60s",
|
||||
"max_retry_policy": 1,
|
||||
"query": "select * from sientia_data.laborious_data where model_id = 1 and \"timestamp\" > NOW() - INTERVAL '5 minutes' order by \"timestamp\" desc limit 30;",
|
||||
"retention_time": 60,
|
||||
"write_tags": [],
|
||||
"input_filters": [
|
||||
{
|
||||
"filter_name": "EMPTY_DATA",
|
||||
"policy": "STOP"
|
||||
},
|
||||
{
|
||||
"filter_name": "SPECIFIC_VARIABLES_NULL_VALUES",
|
||||
"policy": "CONTINUE",
|
||||
"config": {
|
||||
"variables": [
|
||||
"Counter"
|
||||
]
|
||||
}
|
||||
}
|
||||
],
|
||||
"mlflow_transform_filters": [
|
||||
{
|
||||
"filter_name": "API_ERROR",
|
||||
"policy": "REPEAT"
|
||||
},
|
||||
{
|
||||
"filter_name": "NAN_VALUES",
|
||||
"policy": "STOP"
|
||||
}
|
||||
],
|
||||
"mlflow_predict_filters": [
|
||||
{
|
||||
"filter_name": "API_ERROR",
|
||||
"policy": "CONTINUE"
|
||||
}
|
||||
],
|
||||
"path_priority": [
|
||||
"STOP",
|
||||
"CONTINUE",
|
||||
"REPEAT"
|
||||
],
|
||||
"active": true,
|
||||
"updated_at": {
|
||||
"$date": "2026-05-07T23:35:01.600Z"
|
||||
},
|
||||
"save_transform": false,
|
||||
"pi_web_api_output_config": {},
|
||||
"datetime_columns": [
|
||||
"timestamp",
|
||||
"created_at"
|
||||
]
|
||||
},
|
||||
{
|
||||
"schedule_name": "minimal-retrain-test-runtime",
|
||||
"model_id": "1001",
|
||||
"model_name": "test-runtime",
|
||||
"workflow_type": "minimal_retrain",
|
||||
"frequency": "1h",
|
||||
"max_retry_policy": 1,
|
||||
"query": "select * from sientia_data.laborious_data where model_id = 1 and \"timestamp\" > NOW() - INTERVAL '60 minutes' order by \"timestamp\" desc;",
|
||||
"schema": "sientia_data",
|
||||
"table_name": "log_retrain",
|
||||
"datetime_columns": ["timestamp", "created_at"],
|
||||
"model_config": {
|
||||
"target": "Square"
|
||||
},
|
||||
"active": true,
|
||||
"updated_at": {
|
||||
"$date": "2026-05-07T23:35:01.600Z"
|
||||
}
|
||||
},
|
||||
{
|
||||
"schedule_name": "drift-test-runtime",
|
||||
"model_id": "1001",
|
||||
"model_name": "test-runtime",
|
||||
"workflow_type": "drift",
|
||||
"frequency": "5m",
|
||||
"offset": "2m",
|
||||
"max_retry_policy": 1,
|
||||
"execution_timeout_seconds": 300,
|
||||
"task_timeout_seconds": 300,
|
||||
"interval": 5,
|
||||
"drift_metrics": [
|
||||
"kolmogorov_smirnov",
|
||||
"jensen_shannon",
|
||||
"wasserstein"
|
||||
],
|
||||
"chunk_period": "min",
|
||||
"schema": "sientia_data",
|
||||
"source_table_name": "laborious_data",
|
||||
"target_table_name": "drift_metrics",
|
||||
"model_config": {
|
||||
"target": "Square"
|
||||
},
|
||||
"active": true,
|
||||
"updated_at": {
|
||||
"$date": "2026-05-07T23:35:01.600Z"
|
||||
}
|
||||
},
|
||||
{
|
||||
"schedule_name": "simple-metrics-test-runtime",
|
||||
"model_id": "1001",
|
||||
"model_name": "test-runtime",
|
||||
"workflow_type": "simple_metrics",
|
||||
"frequency": "5m",
|
||||
"offset": "2m",
|
||||
"max_retry_policy": 1,
|
||||
"execution_timeout_seconds": 300,
|
||||
"task_timeout_seconds": 300,
|
||||
"interval_minutes": 5,
|
||||
"metrics": ["rmse", "mse", "mae", "r2"],
|
||||
"schema": "sientia_data",
|
||||
"predictions_table_name": "predictions",
|
||||
"data_table_name": "laborious_data",
|
||||
"target_table_name": "simple_metrics",
|
||||
"model_config": {
|
||||
"target": "Square"
|
||||
},
|
||||
"active": true,
|
||||
"updated_at": {
|
||||
"$date": "2026-05-18T23:35:01.600Z"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
25
inter_arrival.py
Normal file
25
inter_arrival.py
Normal file
@@ -0,0 +1,25 @@
|
||||
# %%
|
||||
|
||||
# Load logs.txt
|
||||
with open('logs.txt', 'r') as file:
|
||||
lines = file.readlines()
|
||||
|
||||
# %%
|
||||
import re
|
||||
# Grep "inter-arrival_s=number" with regex
|
||||
intervals = []
|
||||
for line in lines:
|
||||
match = re.search(r'inter-arrival_s=([0-9.]+)', line)
|
||||
if match:
|
||||
intervals.append(float(match.group(1)))
|
||||
# %%
|
||||
|
||||
print(intervals)
|
||||
# %%
|
||||
import matplotlib.pyplot as plt
|
||||
plt.plot(intervals)
|
||||
plt.ylabel('Inter-arrival time (s)')
|
||||
plt.xlabel('Sample')
|
||||
plt.title('Inter-arrival time distribution')
|
||||
plt.show()
|
||||
# %%
|
||||
305
laborious-temporal-plugin-store-migration-plan.md
Normal file
305
laborious-temporal-plugin-store-migration-plan.md
Normal file
@@ -0,0 +1,305 @@
|
||||
---
|
||||
tags:
|
||||
- engineering
|
||||
- sientia
|
||||
- runtime-system
|
||||
- laborious-temporal
|
||||
- plugin-store
|
||||
- migration-plan
|
||||
created: 2026-03-02
|
||||
modified: 2026-03-02
|
||||
created_by: Vitor Pimentel
|
||||
modified_by: Vitor Pimentel
|
||||
status: draft
|
||||
---
|
||||
|
||||
# Sientia Laborious Temporal — PluginStore & Wrapper Migration Plan
|
||||
|
||||
> Migration plan for evolving `sientia-dataops-laborious_temporal` from direct MLflow model loading to a runtime-aware architecture that uses Sientia model wrappers (`SientiaModel`) via their public methods, aligned with the runtime strategy.
|
||||
|
||||
## Summary
|
||||
|
||||
1. [[#Objectives and Scope|Objectives and Scope]] — What this migration must achieve
|
||||
2. [[#Existing State Overview (laborious_temporal)|Existing State Overview]] — Current responsibilities and coupling points
|
||||
3. [[#Requirements Mapping|Requirements Mapping]] — Functional and non-functional requirements
|
||||
4. [[#Target Architecture|Target Architecture]] — Desired runtime and model interaction architecture
|
||||
5. [[#Implementation Plan|Implementation Plan]] — Phased, detailed changes to apply
|
||||
6. [[#Testing Strategy|Testing Strategy]] — How to validate the new behavior
|
||||
7. [[#Rollout and Migration Strategy|Rollout and Migration Strategy]] — How to safely roll out the changes
|
||||
8. [[#Related Documents|Related Documents]] — Cross-links to supporting documents
|
||||
|
||||
---
|
||||
|
||||
## Objectives and Scope
|
||||
|
||||
This migration focuses on the `sientia-dataops-laborious_temporal` application and aims to:
|
||||
|
||||
- Keep the **runtime-aware deployment model** consistent with the rest of the runtime system (Helm + `RUNTIME` env var, runtime installation via PluginStore).
|
||||
- Ensure that **all interactions with models use the public methods of the Sientia wrapper** (`SientiaModel`):
|
||||
- Use `SientiaModel.train(...)` and `retrain(...)` for training and retraining flows.
|
||||
- Use `SientiaModel.predict(...)` and `SientiaModel.transform(...)` for inference and preprocessing.
|
||||
- **Use the shared MLflow repository** (`SientiaMLflowRepository`) for all MLflow operations (load, runs, artifacts, promotion, production lookup, metadata logging); do not implement these in Laborious.
|
||||
|
||||
Out of scope:
|
||||
|
||||
- Replacing MLflow as the tracking and registry backend.
|
||||
- Redesigning Temporal workflows (queues, retry policies) beyond what is required for the new model interaction style.
|
||||
|
||||
---
|
||||
|
||||
## Existing State Overview (laborious_temporal)
|
||||
|
||||
Key components in `sientia-dataops-laborious_temporal`:
|
||||
|
||||
- **MLflow activities** (`laborious/activities/mlflow.py`)
|
||||
- `MLFlow` class exposes Temporal activities for:
|
||||
- `request_transform` — loads transformation models from MLflow and applies them to input data.
|
||||
- `request_predict` — loads predictive models from MLflow and generates predictions.
|
||||
- `retrain_model` — orchestrates retraining using historical data stored in MinIO and MLflow registry.
|
||||
- `update_production_model` — promotes new versions to production.
|
||||
- `get_reference_data` — fetches evaluation/reference datasets from model artifacts.
|
||||
- These activities delegate ML-specific work to `MLFlowRepository`.
|
||||
|
||||
- **MLflow repository** (`laborious/utils/repository/model_repository.py`)
|
||||
- `MLFlowRepository` encapsulates the interaction with MLflow:
|
||||
- Model discovery and run resolution (`get_model_run_id`, `get_model_uri`, `get_experiment`, etc.).
|
||||
- Artifact download and loading for both transformer and prediction models.
|
||||
- Model caching and retention (`get_model`, `get_cached_operation`).
|
||||
- Transformation and prediction entry points:
|
||||
- `transform(...)` wraps `get_cached_operation(..., operation='transform')`.
|
||||
- `predict(...)` wraps `get_cached_operation(..., operation='predict')`.
|
||||
- Retraining orchestration (`fit_models`, `create_new_experiment`, `retrain_model`, `update_production_model`).
|
||||
- Today:
|
||||
- Models are loaded via MLflow flavors: sklearn, pyfunc, pytorch.
|
||||
- When `flavor == 'pyfunc'` and `load_wrapper=True`, the repository loads a wrapper via:
|
||||
- `raw_model = mlflow.pyfunc.load_model(artifact_path)`
|
||||
- `model = raw_model._model_impl.python_model`
|
||||
- Production models are resolved using **stages** in the Model Registry (for example, selecting the latest version in stage `Production`); **aliases such as `@production` are not used yet**, and models are registered explicitly as part of the current retrain/promotion flows.
|
||||
- The wrapper’s `_model_impl` class does not extend `SientiaModel`.
|
||||
- All this MLflow-specific logic is local to `laborious_temporal` and partially duplicated in `sientia-dataops-model-manager`, which motivates the extraction of a shared MLflow repository in `sientia-dataops-library` (see `mlflow-shared-repository-migration-plan`).
|
||||
|
||||
**MLflow:** All MLflow ops → [[mlflow-shared-repository-migration-plan|shared repository]]. Laborious uses the interface; `SientiaModel` lifecycle is in sientia-model-library.
|
||||
|
||||
### Current vs Target — High-level Flow
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
subgraph current [Current State — Laborious Temporal]
|
||||
direction TB
|
||||
TemporalWorker["Temporal Worker"]
|
||||
MlflowActivities["MLFlow Activities\nrequest_transform / request_predict / retrain_model"]
|
||||
MLFlowRepositoryNode["MLFlowRepository"]
|
||||
MLflowRegistry["MLflow Tracking + Registry"]
|
||||
RawModel["Loaded Model\n(sklearn / pyfunc / pytorch)"]
|
||||
|
||||
TemporalWorker -->|"start workflow\n(Temporal)"| MlflowActivities
|
||||
MlflowActivities -->|"call transform()/predict()/retrain_model()"| MLFlowRepositoryNode
|
||||
MLFlowRepositoryNode -->|"search_model_versions()\ncurrent_stage == 'Production'"| MLflowRegistry
|
||||
MLFlowRepositoryNode -->|"mlflow.*.load_model(model_uri)"| RawModel
|
||||
MLFlowRepositoryNode -->|"raw_model.predict(data)\nraw_model.fit(data)"| RawModel
|
||||
end
|
||||
|
||||
subgraph target [Target State — Laborious Temporal]
|
||||
direction TB
|
||||
TemporalWorker2["Temporal Worker"]
|
||||
MlflowActivities2["MLFlow Activities\n(same APIs)"]
|
||||
MLFlowRepositoryNode2["MLFlowRepository\n(wrapper-aware)"]
|
||||
MLflowRegistry2["MLflow Tracking + Registry\n(aliases enabled)"]
|
||||
SientiaWrapperNode["Wrapper Instance\n(extends SientiaModel)"]
|
||||
|
||||
TemporalWorker2 -->|"start workflow\n(Temporal)"| MlflowActivities2
|
||||
MlflowActivities2 -->|"call transform()/predict()/retrain_model()"| MLFlowRepositoryNode2
|
||||
MLFlowRepositoryNode2 -->|"get_model_version_by_alias('production')\n& models:/name@production"| MLflowRegistry2
|
||||
MLFlowRepositoryNode2 -->|"mlflow.pyfunc.load_model(...)"| SientiaWrapperNode
|
||||
MLFlowRepositoryNode2 -->|"wrapper.transform(...)\nwrapper.predict(...)\nwrapper.train()/retrain()"| SientiaWrapperNode
|
||||
SientiaWrapperNode -->|"store_model(...)\n(auto-register + update alias)"| MLflowRegistry2
|
||||
end
|
||||
```
|
||||
|
||||
### Current vs Target — Retrain Hot Path (Code Sketch)
|
||||
|
||||
Current retrain flow inside `MLFlowRepository.fit_models` / `retrain_model` (simplified):
|
||||
|
||||
```python
|
||||
data_model, _ = await self.download_model(
|
||||
model_name=model_name,
|
||||
metadata=metadata,
|
||||
model_type="transform",
|
||||
flavor=transform_flavor,
|
||||
load_wrapper=(transform_flavor == "pyfunc"),
|
||||
)
|
||||
|
||||
prediction_model, _ = await self.download_model(
|
||||
model_name=model_name,
|
||||
metadata=metadata,
|
||||
model_type="predict",
|
||||
flavor=predict_flavor,
|
||||
load_wrapper=(predict_flavor == "pyfunc"),
|
||||
)
|
||||
|
||||
treated_data_candidate = data_model.fit(data) # or data_model.predict(data)
|
||||
...
|
||||
prediction_model.fit(retrain_dataset) # direct fit on underlying model
|
||||
```
|
||||
|
||||
Target retrain flow when using wrappers that extend `SientiaModel`:
|
||||
|
||||
```python
|
||||
wrapper, _ = await self.download_model(
|
||||
model_name=model_name,
|
||||
metadata=metadata,
|
||||
# model_type="predict", # wrapper owns both transformer + model stack
|
||||
# flavor="pyfunc", flavor will always be pyfunc, this parameter will be removed
|
||||
# load_wrapper=True, wrapper will always be loaded from _model_impl
|
||||
)
|
||||
|
||||
# First-time training or full retrain using public API
|
||||
wrapper.train(
|
||||
train_data=train_df, # features + target
|
||||
val_data=val_df,
|
||||
target=target_name,
|
||||
)
|
||||
|
||||
# Incremental retrain (when appropriate)
|
||||
wrapper.retrain(full_retrain_df)
|
||||
|
||||
transformed_df, trans_meta = wrapper.transform(raw_df)
|
||||
pred_df, pred_meta = wrapper.predict({}, transformed_df, params={})
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Requirements Mapping
|
||||
|
||||
### Functional Requirements
|
||||
|
||||
| ID | Requirement | Description | Impacted Areas |
|
||||
| ----- | -------------------------------------------------- | ------------------------------------------------------------------------------------------------ | ------------------------------------------------------ |
|
||||
| FR-01 | Runtime detection and installation | Align worker startup with `RUNTIME` and runtime installation via PluginStore | Worker |
|
||||
| FR-02 | Wrapper‑based training using public API | Retraining must call the public `retrain(...)` method of `SientiaModel` | `fit_models`, `retrain_model` paths |
|
||||
| FR-03 | Wrapper‑based inference using public API | Inference must call `predict(...)` and `transform(...)` on the wrapper; obtain wrappers via `SientiaMLflowRepository` | Activities, shared repository |
|
||||
| FR-04 | Use shared MLflow repository | Delegate all MLflow operations (load, runs, artifacts, promotion, production lookup) to `SientiaMLflowRepository`; do not implement in Laborious | [[mlflow-shared-repository-migration-plan]] |
|
||||
| FR-05 | Model configuration for training/retraining | `model_config` must define `target` and `retention_minutes`; all models are pyfunc + wrapper (no flavor selection) | `model_config` structures |
|
||||
|
||||
### Non-Functional Requirements
|
||||
|
||||
| ID | Requirement | Description | Impacted Areas |
|
||||
| ------ | -------------------------------- | ------------------------------------------------------------------------------------------------------------------------ | --------------------------------------- |
|
||||
| NFR-01 | Consistent public API usage | All wrapper interactions must go through public `SientiaModel` methods; metadata logging is handled by the shared repository | Activities |
|
||||
| NFR-02 | Fail fast when wrappers unavailable | Where wrappers are not yet available, raise an explicit error requesting model update | Shared repository, config |
|
||||
| NFR-03 | Testability | Enable unit tests to validate wrapper-based flows and integration with shared repository | `tests/laborious` |
|
||||
| NFR-04 | Operational safety | Retraining and promotion behavior must remain auditable and robust | Retrain & promotion flows |
|
||||
|
||||
---
|
||||
|
||||
## Target Architecture
|
||||
|
||||
### Wrapper-centric model interactions
|
||||
|
||||
The target state for `laborious_temporal` is:
|
||||
|
||||
- All training and retraining logic goes through the wrapper’s public methods: `train(...)` and `retrain(...)`.
|
||||
- Inference and transformation use `transform(df)` and `predict(context, df, params)`.
|
||||
- All MLflow operations (load, runs, artifacts, promotion, production lookup) go through `SientiaMLflowRepository` (see [[mlflow-shared-repository-migration-plan]]).
|
||||
|
||||
### Use of Shared MLflow Repository
|
||||
|
||||
Laborious obtains wrappers and performs all MLflow operations via `SientiaMLflowRepository`. The shared repository (see [[mlflow-shared-repository-migration-plan]]) owns: alias-based URIs, pyfunc loading, wrapper extraction, promotion, runs, artifacts, and metadata logging. Laborious activities call `repo.load_wrapper(...)` and then the wrapper’s public methods (`transform`, `predict`, `train`, `retrain`); they do not implement MLflow logic.
|
||||
|
||||
---
|
||||
|
||||
## Implementation Plan
|
||||
|
||||
### Phase 0 — Design and Configuration Alignment
|
||||
|
||||
- **P0-01**: Catalog model types and flavors used by `laborious_temporal`:
|
||||
- For each active model:
|
||||
- Flavor will be always pyfunc
|
||||
- Whether a `_model_impl` wrapper is already present and extends `SientiaModel`.
|
||||
- **P0-02**: Define configuration fields in `model_config` for wrapper usage:
|
||||
- Example:
|
||||
- `target` field for training/retraining (already partially present).
|
||||
- `retention_minutes` field for model retention in minutes (unchanged).
|
||||
|
||||
### Phase 1 — Inference via SientiaModel Public API (using shared repository)
|
||||
|
||||
- **P1-01**: Replace `download_model` with `SientiaMLflowRepository.load_wrapper(...)`; remove local MLflow loading logic.
|
||||
- **P1-02**: Update `get_cached_operation` to call `wrapper.transform(data)` and `wrapper.predict({}, data)`; unpack returned metadata; metadata logging is handled by the shared repository.
|
||||
- **P1-03**: Ensure activities (`request_transform`, `request_predict`) remain unchanged externally (inputs/outputs unchanged).
|
||||
|
||||
### Phase 2 — Retraining via SientiaModel.retrain
|
||||
|
||||
- **P2-01**: Refactor `fit_models` to call `wrapper.retrain(data)` and `wrapper.train(...)`; use `SientiaMLflowRepository` for runs, metrics, artifacts, and promotion (no local MLflow logic).
|
||||
|
||||
### Phase 3 — MLflow Logging and Promotion (via shared repository)
|
||||
|
||||
- **P3-01**: In `create_new_experiment`, use the wrapper’s `store_model(...)` for model artifacts; use `SientiaMLflowRepository` for runs, metrics, and any additional MLflow operations.
|
||||
- **P3-02**: Use `SientiaMLflowRepository.promote_to_alias(...)` for production promotion; model registration is auto-handled by wrappers.
|
||||
|
||||
### Phase 4 — Runtime Alignment
|
||||
|
||||
- **P4-01**:`laborious_temporal` is also deployed via the runtime-aware Helm chart:
|
||||
- Read `RUNTIME` env var.
|
||||
- Install runtime via PluginStore before starting Temporal workers.
|
||||
- **P4-02**: Standardize worker queue name as `{project_name}-{runtime}-queue`.
|
||||
- **P4-03**: Fix quality pipelines to a single runtime (to be decided).
|
||||
|
||||
---
|
||||
|
||||
## Testing Strategy
|
||||
|
||||
- **T1 — Unit tests**
|
||||
- Add tests in `tests/laborious/utils/repository/test_model_repository.py` to cover:
|
||||
- Wrapper-based `get_cached_operation` for both `transform` and `predict` (using shared repository).
|
||||
- Wrapper-based `fit_models` and `retrain_model` paths calling `train` and `retrain` respectively.
|
||||
- Code coverage must be 100%.
|
||||
- **T2 — Integration tests**
|
||||
- Use existing end-to-end tests under `e2e/`:
|
||||
- Configure a model with `SientiaModel` wrapper (loaded via shared repository).
|
||||
- Run full prediction and retrain workflows; compare predictions, retrain outcomes, and MLflow artifacts.
|
||||
|
||||
---
|
||||
|
||||
## Rollout and Migration Strategy
|
||||
|
||||
- **R1 — Create models with the new architecture**
|
||||
- Update or create models in the modeling pipeline so they:
|
||||
- Use wrappers that extend `SientiaModel`.
|
||||
- Correctly implement `train`, `retrain`, `predict`, `transform`, and `store_model`.
|
||||
- Produce structured metadata in `transform_meta` and `pred_meta`.
|
||||
- Publish these models to a test store (or dedicated branch/experiment) for initial validation.
|
||||
|
||||
- **R2 — Provision runtimes for the new models**
|
||||
- Configure and install dedicated runtimes for the new models:
|
||||
- Ensure runtime dependencies (Python and system libraries) are available via PluginStore/runtime installer.
|
||||
- Validate that each runtime can:
|
||||
- Load the wrapper through MLflow.
|
||||
- Execute `transform` and `predict` end-to-end on sample data.
|
||||
|
||||
- **R3 — Run new models in real workflows**
|
||||
- Integrate the new wrapper-based models into real Laborious workflows, initially in non-critical environments:
|
||||
- Route only a subset of flows or entities to the new models.
|
||||
- Monitor logs (including metadata), business metrics, and retraining/promotion behavior.
|
||||
- Promote these models to production using aliases (`@production`) in MLflow 3+.
|
||||
|
||||
- **R4 — Migrate remaining models progressively**
|
||||
- Define migration waves by model family:
|
||||
- For each existing model:
|
||||
- Create or adapt a `SientiaModel` wrapper.
|
||||
- Provision the corresponding runtime.
|
||||
- Execute the test cycle (T1–T2) from the previous section.
|
||||
- Update aliases so production traffic uses the new wrapper.
|
||||
- After all models are migrated:
|
||||
- Remove legacy stage-based paths (`Production`) and non-wrapper models.
|
||||
- Simplify the codebase to assume wrappers + aliases only.
|
||||
|
||||
---
|
||||
|
||||
## Related Documents
|
||||
|
||||
- [[mlflow-shared-repository-migration-plan|MLflow Shared Repository Migration Plan]] — Concepts implemented in the common interface (production lookup, wrapper loading, promotion, metadata logging, etc.)
|
||||
- [[model-manager-plugin-store-migration-plan|Model Manager PluginStore Migration Plan]]
|
||||
- [[analytics-implementation-plan|Runtime Analytics Helm Implementation Plan]]
|
||||
- [[analytics|Runtime Analytics Architecture and Analysis]]
|
||||
- [[../model-plugin-system/06-end-to-end-flow|Model Plugin System — End-to-End Flow]]
|
||||
|
||||
@@ -6,90 +6,107 @@ with workflow.unsafe.imports_passed_through():
|
||||
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||
from sientia_do.observability.logger import Logger
|
||||
from sientia_do.repository.minio_repository import MinioRepository
|
||||
from sientia_do.repository.mongodb_repository import MongoDBRepository
|
||||
from sientia_do.repository.minio_repository_sync import MinioRepository
|
||||
from sientia_do.repository.mongodb_repository_sync import MongoDBRepository
|
||||
from sientia_do.utils.connectors_config import build_mongodb_config
|
||||
from sientia_model.model_repository.mlflow_repository import SientiaMLflowRepository
|
||||
from sientia_model.model_repository.plugin_store import PluginStore
|
||||
|
||||
from laborious.activities.api import API
|
||||
from laborious.activities.gates import Gates
|
||||
from laborious.activities.mlflow import MLFlow
|
||||
from laborious.activities.model_import import ModelImport
|
||||
from laborious.activities.model_metrics import ModelMetrics
|
||||
from laborious.activities.opc import OPC
|
||||
from laborious.activities.prediction_metadata import PredictionMetadata
|
||||
from laborious.activities.storage import Storage
|
||||
from laborious.utils.connectors_config import (
|
||||
build_import_config,
|
||||
build_import_status_config,
|
||||
)
|
||||
from laborious.utils.repository.model_repository import MLFlowRepository
|
||||
from laborious.utils.connectors_config import build_mlflow_config
|
||||
|
||||
|
||||
class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API, ModelImport):
|
||||
class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API, PredictionMetadata):
|
||||
"""
|
||||
Main activities orchestrator for the Laborious system.
|
||||
Central orchestrator for all Temporal activities used by Laborious workflows.
|
||||
|
||||
This class combines functionality from multiple activity classes to provide
|
||||
a unified interface for all workflow operations. It manages database connections,
|
||||
MLFlow model interactions, data quality validation, and OPC server communications.
|
||||
Composes Storage (Postgres + MinIO offload), MLFlow (wrapper-based inference and retrain
|
||||
via ``SientiaMLflowRepository``), Gates (data quality and ML response filters), OPC exports,
|
||||
drift/simple metrics, PI Web API writes, and PredictionMetadata (the single write point for
|
||||
gate metadata — see `laborious.activities.prediction_metadata`). The worker constructs one
|
||||
``Activities`` instance per process and registers its callables on multiple workers bound to
|
||||
different task queues.
|
||||
|
||||
The class implements multiple inheritance to combine specialized functionality:
|
||||
- Storage: Database operations and data persistence
|
||||
- MLFlow: Model inference and transformation operations
|
||||
- Gates: Data quality validation and filtering mechanisms
|
||||
- OPC: Real-time data export to OPC servers
|
||||
- ModelMetrics: Model performance metrics and drift detection
|
||||
- API: PI Web API export operations for industrial systems
|
||||
MLflow connectivity: unless ``mlflow_repository`` is injected (tests only), this class builds
|
||||
``SientiaMLflowRepository`` from ``build_mlflow_config()`` so tracking credentials and URL
|
||||
stay aligned with the rest of Laborious env-based configuration.
|
||||
|
||||
Attributes:
|
||||
postgres_config (dict): PostgreSQL connection configuration
|
||||
mlflow_config (dict): MLFlow server configuration
|
||||
opc_config (dict): OPC server configuration
|
||||
pi_web_api_config (dict): PI Web API server configuration
|
||||
logger (Logger): Logging and observability instance
|
||||
notification_handler (NotificationHandler): Notification management instance
|
||||
Inherits and exposes behaviour from mixins; the MLFlow mixin holds ``mlflow_repository``
|
||||
and ``plugin_store`` after ``__init__``.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
postgres_config: dict[str, Any],
|
||||
mlflow_config: dict[str, Any],
|
||||
plugin_store: PluginStore,
|
||||
minio_config: dict[str, Any],
|
||||
opc_config: dict[str, Any],
|
||||
pi_web_api_config: dict[str, Any],
|
||||
logger: Logger,
|
||||
notification_handler: NotificationHandler,
|
||||
mongo_config: dict[str, Any] | None = None,
|
||||
import_config: dict[str, Any] | None = None,
|
||||
import_status_config: dict[str, Any] | None = None,
|
||||
metrics_controller: MetricsController | None = None,
|
||||
mlflow_repository: SientiaMLflowRepository | None = None,
|
||||
mongo_repository: MongoDBRepository | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize the Activities orchestrator with all required configurations.
|
||||
Wire Postgres, MinIO, MLflow, OPC, gates, metrics, and PI Web API into a single object.
|
||||
|
||||
This constructor initializes all parent classes with their respective
|
||||
configurations and sets up the foundation for all activity operations.
|
||||
A single ``MetricsController`` instance is created (or reused) and passed to MinIO,
|
||||
MLflow repository, and all mixins so Prometheus and SDK metrics stay consistent.
|
||||
|
||||
Args:
|
||||
postgres_config: PostgreSQL connection configuration dictionary
|
||||
Required keys: host, port, user, password, dbname, min_connections, max_connections
|
||||
mlflow_config: MLFlow server configuration dictionary
|
||||
Required keys: host, port, username, password
|
||||
opc_config: OPC server configuration dictionary
|
||||
Can contain multiple server configurations
|
||||
pi_web_api_config: PI Web API server configuration dictionary
|
||||
Required keys: base_url, auth_type, auth_token
|
||||
logger: Logger instance for observability and debugging
|
||||
notification_handler: Notification handler for alerts and monitoring
|
||||
mongo_config: MongoDB configuration for the model listing the import writes. When
|
||||
absent, the import activities that need it refuse rather than guess — the
|
||||
notification handler keeps its own client either way.
|
||||
import_config: `build_import_config()`; read from the environment when absent
|
||||
import_status_config: `build_import_status_config()`; read from the environment when
|
||||
absent. This is the BFF database holding the import log: another database on the
|
||||
`POSTGRES_*` server, and a second connection that is never the one Storage holds.
|
||||
- postgres_config: Host, port, credentials, db name, and pool bounds for Storage.
|
||||
- plugin_store: ``PluginStore`` instance; the worker must call ``install_runtime`` before
|
||||
activities run so wrapper code is importable.
|
||||
- minio_config: Endpoint, keys, bucket, retention, and TLS flag for object storage payloads.
|
||||
- opc_config: Map of OPC server id to connection settings for ``OPC`` mixin.
|
||||
- pi_web_api_config: Base URL and auth for ``API`` mixin.
|
||||
- logger: Structured logger used across all activities.
|
||||
- notification_handler: Handler for alerts and persisted notifications.
|
||||
- metrics_controller: Optional shared controller; if ``None``, a new one is created.
|
||||
- mlflow_repository: Optional ``SientiaMLflowRepository`` for unit/e2e tests; in production
|
||||
leave unset so the repository is built from environment via ``build_mlflow_config()``.
|
||||
- mongo_repository: Optional ``MongoDBRepository`` (prediction metadata, best-effort) for
|
||||
unit/e2e tests; in production leave unset so the repository is built from environment
|
||||
via ``build_mongodb_config()``, same pattern as ``mlflow_repository`` above.
|
||||
|
||||
Raises:
|
||||
Exception: If any parent class initialization fails
|
||||
Exception: If any parent ``__init__`` fails (e.g. invalid config keys).
|
||||
|
||||
Return:
|
||||
None
|
||||
"""
|
||||
metrics_controller = MetricsController(logger=logger)
|
||||
|
||||
mc = metrics_controller or MetricsController(logger=logger)
|
||||
|
||||
# Production path: one shared MLflow client for all model registry / tracking calls.
|
||||
if mlflow_repository is None:
|
||||
mlflow_cfg = build_mlflow_config()
|
||||
mlflow_repository = SientiaMLflowRepository(
|
||||
host=mlflow_cfg['url'],
|
||||
username=mlflow_cfg['username'],
|
||||
password=mlflow_cfg['password'],
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=mc,
|
||||
)
|
||||
|
||||
if mongo_repository is None:
|
||||
mongo_cfg = build_mongodb_config()
|
||||
mongo_repository = MongoDBRepository(
|
||||
connection_string=mongo_cfg['connection_string'],
|
||||
database_name=mongo_cfg['database_name'],
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=mc,
|
||||
)
|
||||
|
||||
minio_repository = MinioRepository(
|
||||
endpoint=minio_config['endpoint_url'],
|
||||
@@ -98,20 +115,10 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API, ModelImport):
|
||||
bucket=minio_config['default_bucket'],
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
metrics_controller=mc,
|
||||
secure=minio_config['secure'],
|
||||
)
|
||||
|
||||
mlflow_repository = MLFlowRepository(
|
||||
f'{mlflow_config["host"]}:{mlflow_config["port"]}',
|
||||
mlflow_config['username'],
|
||||
mlflow_config['password'],
|
||||
logger,
|
||||
notification_handler,
|
||||
metrics_controller,
|
||||
)
|
||||
|
||||
# Initialize parent classes
|
||||
Storage.__init__(
|
||||
self,
|
||||
host=postgres_config['host'],
|
||||
@@ -125,16 +132,17 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API, ModelImport):
|
||||
minio_repository=minio_repository,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
metrics_controller=mc,
|
||||
)
|
||||
|
||||
MLFlow.__init__(
|
||||
self,
|
||||
mlflow_repository=mlflow_repository,
|
||||
plugin_store=plugin_store,
|
||||
minio_repository=minio_repository,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
metrics_controller=mc,
|
||||
)
|
||||
|
||||
Gates.__init__(
|
||||
@@ -142,7 +150,15 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API, ModelImport):
|
||||
minio_repository=minio_repository,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
metrics_controller=mc,
|
||||
)
|
||||
|
||||
PredictionMetadata.__init__(
|
||||
self,
|
||||
mongo_repository=mongo_repository,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=mc,
|
||||
)
|
||||
|
||||
OPC.__init__(
|
||||
@@ -150,14 +166,14 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API, ModelImport):
|
||||
opc_servers=opc_config,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
metrics_controller=mc,
|
||||
)
|
||||
|
||||
ModelMetrics.__init__(
|
||||
self,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
metrics_controller=mc,
|
||||
)
|
||||
|
||||
API.__init__(
|
||||
@@ -167,51 +183,23 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API, ModelImport):
|
||||
auth_token=pi_web_api_config['auth_token'],
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
metrics_controller=mc,
|
||||
)
|
||||
|
||||
mongo_repository = (
|
||||
MongoDBRepository(
|
||||
connection_string=mongo_config['connection_string'],
|
||||
database_name=mongo_config['database_name'],
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
if mongo_config
|
||||
else None
|
||||
)
|
||||
|
||||
ModelImport.__init__(
|
||||
self,
|
||||
status_config=import_status_config or build_import_status_config(),
|
||||
import_config=import_config or build_import_config(),
|
||||
minio_repository=minio_repository,
|
||||
mlflow_repository=mlflow_repository,
|
||||
mongo_repository=mongo_repository,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
async def shutdown(self):
|
||||
def shutdown(self) -> None:
|
||||
"""
|
||||
Gracefully shutdown all activities and clean up resources.
|
||||
Close database pools, sync clients, and OPC sessions in a defined order.
|
||||
|
||||
This method ensures proper cleanup of all resources including:
|
||||
- PostgreSQL connection pools
|
||||
- OPC server connections
|
||||
- PI Web API client connections
|
||||
- MLFlow model repositories
|
||||
- Any other resources that need explicit cleanup
|
||||
Should be invoked on worker exit so connection pools and OPC sessions are released
|
||||
cleanly before process termination.
|
||||
|
||||
The method should be called before the application terminates to ensure
|
||||
proper resource cleanup and prevent resource leaks.
|
||||
Return:
|
||||
None
|
||||
"""
|
||||
Storage.close(self)
|
||||
MLFlow.close(self)
|
||||
Gates.close(self)
|
||||
await OPC.aclose(self)
|
||||
OPC.close(self)
|
||||
ModelMetrics.close(self)
|
||||
API.close(self)
|
||||
ModelImport.close(self)
|
||||
PredictionMetadata.close(self)
|
||||
|
||||
@@ -11,7 +11,7 @@ with workflow.unsafe.imports_passed_through():
|
||||
from sientia_do.observability.logger import Logger
|
||||
from sientia_do.observability.metrics_controller import MetricsController
|
||||
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
|
||||
from sientia_do.repository.pi_web_api_client import PIWebAPIClient
|
||||
from sientia_do.repository.pi_web_api_client_sync import PIWebAPIClient
|
||||
|
||||
from laborious import metrics
|
||||
|
||||
@@ -105,7 +105,7 @@ class API(SientiaMonitoring):
|
||||
self.pi_web_api_client.close()
|
||||
SientiaMonitoring.shutdown(self)
|
||||
|
||||
async def process_pi_web_api_response(
|
||||
def process_pi_web_api_response(
|
||||
self,
|
||||
response_data: list[dict[str, Any]],
|
||||
tags: dict[str, str],
|
||||
@@ -156,7 +156,7 @@ class API(SientiaMonitoring):
|
||||
self.error(
|
||||
f'Error writing tag {tag_name}:{web_id} to PI Web API: {errors}', metadata
|
||||
)
|
||||
await self.emit_metric(
|
||||
self.emit_metric_sync(
|
||||
metric_object=metrics.PI_WEB_API_PREDICTION_WRITTEN_ERROR_COUNT,
|
||||
tags={
|
||||
**core_labels,
|
||||
@@ -165,7 +165,7 @@ class API(SientiaMonitoring):
|
||||
)
|
||||
confidence = PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
|
||||
else:
|
||||
await self.emit_metric(
|
||||
self.emit_metric_sync(
|
||||
metric_object=metrics.PI_WEB_API_PREDICTION_WRITTEN_COUNT,
|
||||
tags={
|
||||
**core_labels,
|
||||
@@ -182,7 +182,7 @@ class API(SientiaMonitoring):
|
||||
metadata,
|
||||
)
|
||||
|
||||
await self.send_notification_async(
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='WRITE_PI_WEB_API_PREDICTION_ERROR',
|
||||
message=f'The number of written tags does not match the number of tag names: Expected {tag_names} tags, but {written_tags} tags were written.\nResponse:\n {json.dumps(response_data, indent=4)}\nTags:\n {json.dumps(tags, indent=4)}',
|
||||
@@ -194,7 +194,7 @@ class API(SientiaMonitoring):
|
||||
return confidence, message
|
||||
|
||||
@activity.defn(name='write_pi_web_api_data')
|
||||
async def write_pi_web_api_data(self, input_data: dict[str, Any]) -> dict[Any, Any]:
|
||||
def write_pi_web_api_data(self, input_data: dict[str, Any]) -> dict[Any, Any]:
|
||||
"""
|
||||
Write prediction and confidence data to PI Web API.
|
||||
|
||||
@@ -232,7 +232,7 @@ class API(SientiaMonitoring):
|
||||
confidence_value = data.head(1)['prediction_confidence'].values[0]
|
||||
|
||||
try:
|
||||
prediction_response = await self.pi_web_api_client.write_value(
|
||||
prediction_response = self.pi_web_api_client.write_value(
|
||||
web_ids=prediction_tags,
|
||||
value={
|
||||
'Timestamp': data.head(1)['timestamp'].values[0],
|
||||
@@ -241,7 +241,7 @@ class API(SientiaMonitoring):
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
confidence, message = await self.process_pi_web_api_response(
|
||||
confidence, message = self.process_pi_web_api_response(
|
||||
response_data=prediction_response,
|
||||
tags=raw_prediction_tags,
|
||||
core_labels=core_labels,
|
||||
@@ -258,7 +258,7 @@ class API(SientiaMonitoring):
|
||||
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
await self.send_notification_async(
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='WRITE_PI_WEB_API_PREDICTION_ERROR',
|
||||
message=f'Error writing prediction data to PI Web API: {e}\n Tags: {raw_prediction_tags}',
|
||||
@@ -275,7 +275,7 @@ class API(SientiaMonitoring):
|
||||
return data.to_dict()
|
||||
|
||||
try:
|
||||
confidence_response = await self.pi_web_api_client.write_value(
|
||||
confidence_response = self.pi_web_api_client.write_value(
|
||||
web_ids=confidence_tags,
|
||||
value={
|
||||
'Timestamp': data.head(1)['timestamp'].values[0],
|
||||
@@ -284,7 +284,7 @@ class API(SientiaMonitoring):
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
await self.process_pi_web_api_response(
|
||||
self.process_pi_web_api_response(
|
||||
response_data=confidence_response,
|
||||
tags=raw_confidence_tags,
|
||||
core_labels=core_labels,
|
||||
@@ -293,7 +293,7 @@ class API(SientiaMonitoring):
|
||||
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
await self.send_notification_async(
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='WRITE_PI_WEB_API_CONFIDENCE_ERROR',
|
||||
message=f'Error writing confidence data to PI Web API: {e}\n Tags: {raw_confidence_tags}',
|
||||
|
||||
@@ -1,8 +1,5 @@
|
||||
from sientia_do.repository.minio_repository import MinioRepository
|
||||
from temporalio import activity, workflow
|
||||
|
||||
from laborious.utils.repository.minio_manager import MinioManager
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
import traceback
|
||||
from collections.abc import Callable, Mapping
|
||||
@@ -13,6 +10,8 @@ with workflow.unsafe.imports_passed_through():
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
from sientia_do.observability.logger import Logger
|
||||
from sientia_do.observability.metrics_controller import MetricsController
|
||||
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
|
||||
from sientia_do.repository.minio_repository_sync import MinioRepository
|
||||
from sientia_do.utils.formatters import create_sample_dict
|
||||
|
||||
from laborious import metrics
|
||||
@@ -28,6 +27,7 @@ with workflow.unsafe.imports_passed_through():
|
||||
InputFilterFunc = Callable[[DataFrame, dict[str, Any]], bool]
|
||||
ResponseFilterFunc = Callable[[dict[str, Any], dict[str, Any]], bool]
|
||||
ContentFilterFunc = Callable[[DataFrame, dict[str, Any]], bool]
|
||||
InputCounterFunc = Callable[[DataFrame, dict[str, Any]], int]
|
||||
|
||||
# Input filter function mappings
|
||||
input_filter_functions: dict[str, InputFilterFunc] = {
|
||||
@@ -35,6 +35,42 @@ input_filter_functions: dict[str, InputFilterFunc] = {
|
||||
'EMPTY_DATA': filter_empty_data,
|
||||
}
|
||||
|
||||
|
||||
def _count_specific_variables_null_values(data: DataFrame, config: dict[str, Any]) -> int:
|
||||
"""
|
||||
Count rows with a null value in one of the configured variables.
|
||||
|
||||
Mirrors `filter_specific_variables_null_values`'s own trigger condition, so the count
|
||||
reflects exactly what made that filter fire.
|
||||
|
||||
Return:
|
||||
int: Number of rows with a null value in `config['variables']`.
|
||||
"""
|
||||
if data.empty:
|
||||
return 0
|
||||
return int((data['variable'].isin(config['variables']) & data['value'].isna()).sum())
|
||||
|
||||
|
||||
def _count_empty_data(_data: DataFrame, _config: dict[str, Any]) -> int:
|
||||
"""
|
||||
Count rows for the EMPTY_DATA filter.
|
||||
|
||||
`EMPTY_DATA` only triggers when there are no rows at all, so the count is always 0.
|
||||
|
||||
Return:
|
||||
int: Always 0.
|
||||
"""
|
||||
return 0
|
||||
|
||||
|
||||
# Per-filter row counts for the `counter` field in prediction metadata (gate-local: mirrors each
|
||||
# filter's own trigger condition instead of changing `conditional_filters.py`'s bool contract,
|
||||
# which `mlflow_content_gate` also depends on for EMPTY_DATA).
|
||||
input_filter_counters: dict[str, InputCounterFunc] = {
|
||||
'SPECIFIC_VARIABLES_NULL_VALUES': _count_specific_variables_null_values,
|
||||
'EMPTY_DATA': _count_empty_data,
|
||||
}
|
||||
|
||||
# Confidence mappings kept separate from function maps to avoid Union types
|
||||
input_path_confidence: Mapping[str, int] = {
|
||||
'STOP': -1,
|
||||
@@ -42,6 +78,7 @@ input_path_confidence: Mapping[str, int] = {
|
||||
'REPEAT': -1,
|
||||
}
|
||||
|
||||
|
||||
# MLFlow response filter function mappings
|
||||
mlflow_response_filter_functions: dict[str, ResponseFilterFunc] = {
|
||||
'API_ERROR': api_error_filter,
|
||||
@@ -53,12 +90,34 @@ mlflow_response_path_confidence: Mapping[str, int] = {
|
||||
'REPEAT': -1,
|
||||
}
|
||||
|
||||
|
||||
def _count_nan_values(data: DataFrame, _config: dict[str, Any]) -> int:
|
||||
"""
|
||||
Count rows for the NAN_VALUES filter.
|
||||
|
||||
`NAN_VALUES` only triggers when every value in the DataFrame is NaN, so every row in it
|
||||
matched the condition — same convention as `_count_specific_variables_null_values`:
|
||||
counting rows that satisfied the filter's own trigger condition.
|
||||
|
||||
Return:
|
||||
int: Number of rows in the evaluated DataFrame.
|
||||
"""
|
||||
return len(data)
|
||||
|
||||
|
||||
# MLFlow content filter function mappings
|
||||
mlflow_content_filter_functions: dict[str, ContentFilterFunc] = {
|
||||
'NAN_VALUES': nan_values_filter,
|
||||
'EMPTY_DATA': filter_empty_data,
|
||||
}
|
||||
|
||||
# `EMPTY_DATA` reuses `_count_empty_data` — same function (`filter_empty_data`) and same
|
||||
# trigger condition (absence of rows) as `input_gate`'s own EMPTY_DATA.
|
||||
mlflow_content_filter_counters: dict[str, InputCounterFunc] = {
|
||||
'NAN_VALUES': _count_nan_values,
|
||||
'EMPTY_DATA': _count_empty_data,
|
||||
}
|
||||
|
||||
mlflow_content_path_confidence: Mapping[str, int] = {
|
||||
'STOP': -1,
|
||||
'CONTINUE': 18,
|
||||
@@ -66,7 +125,7 @@ mlflow_content_path_confidence: Mapping[str, int] = {
|
||||
}
|
||||
|
||||
|
||||
class Gates(MinioManager):
|
||||
class Gates(SientiaMonitoring):
|
||||
"""
|
||||
Data quality gates and filtering activities for the Laborious system.
|
||||
|
||||
@@ -106,8 +165,12 @@ class Gates(MinioManager):
|
||||
Raises:
|
||||
Exception: If BaseActivity initialization fails
|
||||
"""
|
||||
MinioManager.__init__(
|
||||
self, minio_repository, logger, notification_handler, metrics_controller
|
||||
self.minio_repository = minio_repository
|
||||
SientiaMonitoring.__init__(
|
||||
self,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
def close(self) -> None:
|
||||
@@ -115,7 +178,12 @@ class Gates(MinioManager):
|
||||
Close the gates activity and clean up resources.
|
||||
"""
|
||||
|
||||
MinioManager.close(self)
|
||||
if self.minio_repository is not None:
|
||||
try:
|
||||
self.minio_repository.close()
|
||||
finally:
|
||||
self.minio_repository = None
|
||||
SientiaMonitoring.shutdown(self)
|
||||
|
||||
def __del__(self):
|
||||
self.close()
|
||||
@@ -155,7 +223,9 @@ class Gates(MinioManager):
|
||||
return policy, filter_config
|
||||
|
||||
@activity.defn(name='input_gate')
|
||||
async def input_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
|
||||
def input_gate(
|
||||
self, input_data: dict[str, Any]
|
||||
) -> tuple[str | None, int, str, dict[str, Any] | None]:
|
||||
"""
|
||||
Apply input data quality filters and validation.
|
||||
|
||||
@@ -170,6 +240,9 @@ class Gates(MinioManager):
|
||||
3. Determines appropriate path decisions (STOP, CONTINUE, REPEAT)
|
||||
4. Provides confidence scores and detailed comments
|
||||
5. Handles errors gracefully with notification integration
|
||||
6. Builds a metadata document describing which filter(s) triggered, with their policy
|
||||
and a counter, whenever one does — this activity only computes it, it never writes it
|
||||
(see `PredictionMetadata.write_prediction_metadata` for the actual write)
|
||||
|
||||
Args:
|
||||
input_data: Configuration and data for input validation
|
||||
@@ -180,10 +253,14 @@ class Gates(MinioManager):
|
||||
- path_priority (list[str]): Priority order for path decisions
|
||||
|
||||
Returns:
|
||||
tuple: (path_flag, confidence, comment)
|
||||
tuple: (path_flag, confidence, comment, document)
|
||||
- path_flag (str | None): Decision path (STOP, CONTINUE, REPEAT, or None)
|
||||
- confidence (int): Confidence score for the decision
|
||||
- comment (str): Detailed explanation of the decision
|
||||
- document (dict | None): Metadata describing the filter(s) that produced
|
||||
`path_flag`, shaped `{'filters': [{'name', 'policy', 'counter'}, ...]}`;
|
||||
`None` when nothing was filtered. Not written anywhere by this activity —
|
||||
the caller decides whether and how to persist it.
|
||||
|
||||
Raises:
|
||||
Exception: If filter execution fails or configuration is invalid
|
||||
@@ -194,10 +271,11 @@ class Gates(MinioManager):
|
||||
|
||||
filters = input_data['filters']
|
||||
payload = MinioDataFramePayload.from_dict(input_data['data'])
|
||||
data = await payload.retrieve(self.minio_repository, metadata)
|
||||
data = payload.retrieve(self.minio_repository, metadata)
|
||||
path_priority = input_data['path_priority']
|
||||
|
||||
filter_output = []
|
||||
# (filter_name, policy, filter_config) per filter that triggered, in configured order
|
||||
triggered_filters: list[tuple[str, str, dict[str, Any]]] = []
|
||||
|
||||
self._debug_dataframe('Input data:', data, metadata)
|
||||
self.debug(f'Filters: {filters}', metadata)
|
||||
@@ -211,10 +289,10 @@ class Gates(MinioManager):
|
||||
try:
|
||||
if input_filter_functions[fil](data, filter_config):
|
||||
self.debug(f'Data not passed the input filter {fil}:{config}', metadata)
|
||||
filter_output.append(policy)
|
||||
triggered_filters.append((fil, policy, filter_config))
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
await self.send_notification_async(
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id=f'INTPUT_GATE_ERROR__{fil}',
|
||||
message=f'Error in filter {fil}:{config}: \n {e}',
|
||||
@@ -223,19 +301,51 @@ class Gates(MinioManager):
|
||||
attachment_content=trace,
|
||||
)
|
||||
|
||||
triggered_policies = [policy for _, policy, _ in triggered_filters]
|
||||
|
||||
for path_flag in path_priority:
|
||||
if path_flag in filter_output:
|
||||
if path_flag in triggered_policies:
|
||||
self.info(f'Input gate result: {path_flag}', metadata)
|
||||
return path_flag, input_path_confidence[path_flag], 'Input data with bad quality'
|
||||
|
||||
matched = [
|
||||
(fil, filter_config)
|
||||
for fil, policy, filter_config in triggered_filters
|
||||
if policy == path_flag
|
||||
]
|
||||
document = {
|
||||
'filters': [
|
||||
{
|
||||
'name': fil,
|
||||
'policy': path_flag,
|
||||
# `.get(..., lambda *_: 0)`, not `[fil]`: a filter added to
|
||||
# `input_filter_functions` without a matching counter here degrades to
|
||||
# counter=0 instead of a raw KeyError crashing this prediction (review
|
||||
# finding, 2026-09-30 — see test enforcing key parity between the two).
|
||||
'counter': input_filter_counters.get(fil, lambda *_: 0)(
|
||||
data, filter_config
|
||||
),
|
||||
}
|
||||
for fil, filter_config in matched
|
||||
],
|
||||
}
|
||||
|
||||
return (
|
||||
path_flag,
|
||||
input_path_confidence[path_flag],
|
||||
'Input data with bad quality',
|
||||
document,
|
||||
)
|
||||
|
||||
self.info('Nothing was filtered by the input gate', metadata)
|
||||
|
||||
del data
|
||||
|
||||
return None, 0, ''
|
||||
return None, 0, '', None
|
||||
|
||||
@activity.defn(name='mlflow_response_gate')
|
||||
async def mlflow_response_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
|
||||
def mlflow_response_gate(
|
||||
self, input_data: dict[str, Any]
|
||||
) -> tuple[str | None, int, str, dict[str, Any] | None]:
|
||||
"""
|
||||
Validate MLFlow API response quality and integrity.
|
||||
|
||||
@@ -249,6 +359,9 @@ class Gates(MinioManager):
|
||||
3. Determines path decisions based on response validation results
|
||||
4. Provides confidence scores and detailed validation comments
|
||||
5. Handles API errors and response validation failures
|
||||
6. Builds a metadata document describing which filter(s) triggered, with their policy
|
||||
and a counter, whenever one does — this activity only computes it, it never writes it
|
||||
(see `PredictionMetadata.write_prediction_metadata` for the actual write)
|
||||
|
||||
Args:
|
||||
input_data: Configuration and data for response validation
|
||||
@@ -260,10 +373,14 @@ class Gates(MinioManager):
|
||||
- path_priority (list[str]): Priority order for path decisions
|
||||
|
||||
Returns:
|
||||
tuple: (path_flag, confidence, comment)
|
||||
tuple: (path_flag, confidence, comment, document)
|
||||
- path_flag (str | None): Decision path (STOP, CONTINUE, REPEAT, or None)
|
||||
- confidence (int): Confidence score for the decision
|
||||
- comment (str): Detailed explanation of the decision
|
||||
- document (dict | None): Metadata describing the filter(s) that produced
|
||||
`path_flag`, shaped `{'filters': [{'name', 'policy', 'counter'}, ...]}`;
|
||||
`None` when nothing was filtered. Not written anywhere by this activity —
|
||||
the caller decides whether and how to persist it.
|
||||
|
||||
Raises:
|
||||
Exception: If response validation fails or configuration is invalid
|
||||
@@ -279,15 +396,16 @@ class Gates(MinioManager):
|
||||
self.debug(f'Filters: {filters}', metadata)
|
||||
|
||||
payload = MinioDataFramePayload.from_dict(raw_data)
|
||||
data = await payload.retrieve(self.minio_repository, metadata)
|
||||
data = payload.retrieve(self.minio_repository, metadata)
|
||||
|
||||
gate_type = input_data['type']
|
||||
path_priority = input_data['path_priority']
|
||||
|
||||
filter_output = []
|
||||
|
||||
comments = []
|
||||
|
||||
# (filter_name, policy, filter_config) per filter that triggered, in configured order
|
||||
triggered_filters: list[tuple[str, str, dict[str, Any]]] = []
|
||||
|
||||
status = payload.status or {}
|
||||
|
||||
for fil, config in filters.items():
|
||||
@@ -296,19 +414,25 @@ class Gates(MinioManager):
|
||||
policy, filter_config = self._read_filter_entry(config)
|
||||
try:
|
||||
if mlflow_response_filter_functions[fil](status, filter_config):
|
||||
filter_output.append(policy)
|
||||
comments.append(status.get('message', 'Unknown MLFlow API error'))
|
||||
await self.send_notification_async(
|
||||
triggered_filters.append((fil, policy, filter_config))
|
||||
# The real error lives one level in, under `content` — `request_transform`/
|
||||
# `request_predict` nest `message`/`traceback` there on failure (mlflow.py),
|
||||
# reusing the same `content` key the success path uses for the DataFrame
|
||||
# (review finding, 2026-10-08: reading `status['message']` directly always
|
||||
# fell back to the generic default, in every real production occurrence).
|
||||
error_content = status.get('content') or {}
|
||||
comments.append(error_content.get('message', 'Unknown MLFlow API error'))
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id=f'{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}',
|
||||
message=status.get('message', 'Unknown MLFlow API error'),
|
||||
message=error_content.get('message', 'Unknown MLFlow API error'),
|
||||
block='mlflow_gate',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=status.get('traceback'),
|
||||
attachment_content=error_content.get('traceback'),
|
||||
)
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
await self.send_notification_async(
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id=f'MLFLOW_GATE_RESPONSE_FILTER__{fil}',
|
||||
message=f'Error in filter {fil}:{config}: \n {e}',
|
||||
@@ -317,19 +441,47 @@ class Gates(MinioManager):
|
||||
attachment_content=trace,
|
||||
)
|
||||
|
||||
triggered_policies = [policy for _, policy, _ in triggered_filters]
|
||||
|
||||
for path_flag in path_priority:
|
||||
if path_flag in filter_output:
|
||||
if path_flag in triggered_policies:
|
||||
self.info(f'Mlflow response gate result: {path_flag}', metadata)
|
||||
return path_flag, mlflow_response_path_confidence[path_flag], ', '.join(comments)
|
||||
|
||||
matched = [fil for fil, policy, _ in triggered_filters if policy == path_flag]
|
||||
document = {
|
||||
'filters': [
|
||||
{
|
||||
'name': fil,
|
||||
'policy': path_flag,
|
||||
# No per-row/per-item unit to count here — `api_error_filter`
|
||||
# evaluates a single API response as a whole, not tabular data —
|
||||
# so, unlike `input_gate`'s per-filter counters, this is a fixed
|
||||
# value: 1 (one API error event). Inlined rather than a counter
|
||||
# function/dict, since there's only one response filter today and
|
||||
# this isn't an override of anything (review finding).
|
||||
'counter': 1,
|
||||
}
|
||||
for fil in matched
|
||||
],
|
||||
}
|
||||
|
||||
return (
|
||||
path_flag,
|
||||
mlflow_response_path_confidence[path_flag],
|
||||
', '.join(comments),
|
||||
document,
|
||||
)
|
||||
|
||||
self.info('Nothing was filtered by the mlflow response gate', metadata)
|
||||
|
||||
del data
|
||||
|
||||
return None, 0, ''
|
||||
return None, 0, '', None
|
||||
|
||||
@activity.defn(name='mlflow_content_gate')
|
||||
async def mlflow_content_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
|
||||
def mlflow_content_gate(
|
||||
self, input_data: dict[str, Any]
|
||||
) -> tuple[str | None, int, str, dict[str, Any] | None]:
|
||||
"""
|
||||
Validate MLFlow prediction content quality and integrity.
|
||||
|
||||
@@ -343,6 +495,9 @@ class Gates(MinioManager):
|
||||
3. Determines path decisions based on content validation results
|
||||
4. Provides confidence scores and detailed validation comments
|
||||
5. Handles content validation failures and quality issues
|
||||
6. Builds a metadata document describing which filter(s) triggered, with their policy
|
||||
and a counter, whenever one does — this activity only computes it, it never writes it
|
||||
(see `PredictionMetadata.write_prediction_metadata` for the actual write)
|
||||
|
||||
Args:
|
||||
input_data: Configuration and data for content validation
|
||||
@@ -354,10 +509,14 @@ class Gates(MinioManager):
|
||||
- path_priority (list[str]): Priority order for path decisions
|
||||
|
||||
Returns:
|
||||
tuple: (path_flag, confidence, comment)
|
||||
tuple: (path_flag, confidence, comment, document)
|
||||
- path_flag (str | None): Decision path (STOP, CONTINUE, REPEAT, or None)
|
||||
- confidence (int): Confidence score for the decision
|
||||
- comment (str): Detailed explanation of the decision
|
||||
- document (dict | None): Metadata describing the filter(s) that produced
|
||||
`path_flag`, shaped `{'filters': [{'name', 'policy', 'counter'}, ...]}`;
|
||||
`None` when nothing was filtered. Not written anywhere by this activity —
|
||||
the caller decides whether and how to persist it.
|
||||
|
||||
Raises:
|
||||
Exception: If content validation fails or configuration is invalid
|
||||
@@ -368,12 +527,13 @@ class Gates(MinioManager):
|
||||
filters = input_data['filters']
|
||||
|
||||
payload = MinioDataFramePayload.from_dict(input_data['data'])
|
||||
data = await payload.retrieve(self.minio_repository, metadata)
|
||||
data = payload.retrieve(self.minio_repository, metadata)
|
||||
|
||||
gate_type = input_data['type']
|
||||
path_priority = input_data['path_priority']
|
||||
|
||||
filter_output = []
|
||||
# (filter_name, policy, filter_config) per filter that triggered, in configured order
|
||||
triggered_filters: list[tuple[str, str, dict[str, Any]]] = []
|
||||
|
||||
self._debug_dataframe('Input data:', data, metadata)
|
||||
self.debug(f'Filters: \n {filters}', metadata)
|
||||
@@ -384,8 +544,8 @@ class Gates(MinioManager):
|
||||
policy, filter_config = self._read_filter_entry(config)
|
||||
try:
|
||||
if mlflow_content_filter_functions[fil](data, filter_config):
|
||||
filter_output.append(policy)
|
||||
await self.send_notification_async(
|
||||
triggered_filters.append((fil, policy, filter_config))
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id=f'{gate_type.upper()}_GATE_CONTENT_FILTER__{fil}',
|
||||
message=f'Data not passed the content filter {fil}:{config}',
|
||||
@@ -395,7 +555,7 @@ class Gates(MinioManager):
|
||||
)
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
await self.send_notification_async(
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id=f'MLFLOW_GATE_CONTENT_FILTER__{fil}',
|
||||
message=f'Error in filter {fil}:{config}: \n {e}',
|
||||
@@ -404,20 +564,46 @@ class Gates(MinioManager):
|
||||
attachment_content=trace,
|
||||
)
|
||||
|
||||
triggered_policies = [policy for _, policy, _ in triggered_filters]
|
||||
|
||||
for path_flag in path_priority:
|
||||
if path_flag in filter_output:
|
||||
if path_flag in triggered_policies:
|
||||
self.info(f'Mlflow content gate result: {path_flag}', metadata)
|
||||
|
||||
matched = [
|
||||
(fil, filter_config)
|
||||
for fil, policy, filter_config in triggered_filters
|
||||
if policy == path_flag
|
||||
]
|
||||
document = {
|
||||
'filters': [
|
||||
{
|
||||
'name': fil,
|
||||
'policy': path_flag,
|
||||
# `.get(..., lambda *_: 0)`, not `[fil]`: same defensive fallback as
|
||||
# `input_gate` (review finding, 2026-09-30) — a filter added to
|
||||
# `mlflow_content_filter_functions` without a matching counter here
|
||||
# degrades to counter=0 instead of a raw KeyError.
|
||||
'counter': mlflow_content_filter_counters.get(fil, lambda *_: 0)(
|
||||
data, filter_config
|
||||
),
|
||||
}
|
||||
for fil, filter_config in matched
|
||||
],
|
||||
}
|
||||
|
||||
return (
|
||||
path_flag,
|
||||
mlflow_content_path_confidence[path_flag],
|
||||
'Transformed data not passed the content filter',
|
||||
document,
|
||||
)
|
||||
|
||||
self.info('Nothing was filtered by the mlflow content gate', metadata)
|
||||
|
||||
del data
|
||||
|
||||
return None, 0, ''
|
||||
return None, 0, '', None
|
||||
|
||||
def get_prediction_store_policy(
|
||||
self, prediction_store_policy: str, metadata: dict[str, Any]
|
||||
@@ -471,7 +657,7 @@ class Gates(MinioManager):
|
||||
return policy_type, int(policy_value)
|
||||
|
||||
@activity.defn(name='format_transformed_data')
|
||||
async def format_transformed_data(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
|
||||
def format_transformed_data(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
|
||||
"""
|
||||
Format transformed data for storage and export operations.
|
||||
|
||||
@@ -507,7 +693,7 @@ class Gates(MinioManager):
|
||||
self.info('Formatting transformed data...', metadata)
|
||||
|
||||
payload = MinioDataFramePayload.from_dict(input_data['data'])
|
||||
data = await payload.retrieve(self.minio_repository, metadata)
|
||||
data = payload.retrieve(self.minio_repository, metadata)
|
||||
|
||||
data['timestamp'] = data.index
|
||||
data = data.reset_index(drop=True)
|
||||
@@ -515,7 +701,7 @@ class Gates(MinioManager):
|
||||
data = data.melt(id_vars='timestamp', var_name='variable', value_name='value')
|
||||
data['model_id'] = model_id
|
||||
|
||||
return await MinioDataFramePayload.from_dataframe(
|
||||
return MinioDataFramePayload.from_dataframe(
|
||||
dataframe=data,
|
||||
minio_repo=self.minio_repository,
|
||||
model_name=input_data['model_name'],
|
||||
@@ -526,7 +712,7 @@ class Gates(MinioManager):
|
||||
)
|
||||
|
||||
@activity.defn(name='format_prediction')
|
||||
async def format_prediction(self, input_data: dict[str, Any]) -> dict:
|
||||
def format_prediction(self, input_data: dict[str, Any]) -> dict:
|
||||
"""
|
||||
Format prediction data according to configured storage policies.
|
||||
|
||||
@@ -558,7 +744,7 @@ class Gates(MinioManager):
|
||||
self.info('Formatting prediction...', metadata)
|
||||
|
||||
payload = MinioDataFramePayload.from_dict(input_data['data'])
|
||||
data = await payload.retrieve(self.minio_repository, metadata)
|
||||
data = payload.retrieve(self.minio_repository, metadata)
|
||||
|
||||
# Create timestamp column from index and reset index
|
||||
data['timestamp'] = data.index
|
||||
@@ -607,7 +793,7 @@ class Gates(MinioManager):
|
||||
return data.to_dict()
|
||||
|
||||
@activity.defn(name='format_default_prediction')
|
||||
async def format_default_prediction(self, input_data: dict[str, Any]) -> dict:
|
||||
def format_default_prediction(self, input_data: dict[str, Any]) -> dict:
|
||||
"""
|
||||
Create and format default prediction data for error conditions.
|
||||
|
||||
@@ -628,6 +814,9 @@ class Gates(MinioManager):
|
||||
- model_id (str): Unique identifier for the ML model
|
||||
- prediction_confidence (float): Confidence score (typically low for errors)
|
||||
- comment (str): Error description or operational comment
|
||||
- metadata_id (str | None, optional): Id of the `prediction_metadata` Mongo
|
||||
document to reference for this row; `None`/absent when there's nothing to
|
||||
reference
|
||||
|
||||
Returns:
|
||||
dict: Formatted default prediction data with error indicators
|
||||
@@ -645,6 +834,7 @@ class Gates(MinioManager):
|
||||
'prediction_confidence': [input_data['prediction_confidence']],
|
||||
'prediction_status': ['Bad'],
|
||||
'comments': [input_data['comment']],
|
||||
'metadata_id': [input_data.get('metadata_id')],
|
||||
}
|
||||
)
|
||||
|
||||
@@ -652,7 +842,7 @@ class Gates(MinioManager):
|
||||
return data.to_dict()
|
||||
|
||||
@activity.defn(name='format_retrain_report')
|
||||
async def format_retrain_report(self, input_data: dict[str, Any]) -> dict:
|
||||
def format_retrain_report(self, input_data: dict[str, Any]) -> dict:
|
||||
"""
|
||||
Format retrain report data for storage and audit trail maintenance.
|
||||
|
||||
@@ -722,7 +912,7 @@ class Gates(MinioManager):
|
||||
return report.to_dict()
|
||||
|
||||
@activity.defn(name='write_metrics')
|
||||
async def write_metrics(self, input_data: dict[str, Any]):
|
||||
def write_metrics(self, input_data: dict[str, Any]):
|
||||
"""
|
||||
Write prediction performance metrics to Prometheus monitoring system.
|
||||
|
||||
@@ -759,19 +949,19 @@ class Gates(MinioManager):
|
||||
'model_name': metadata['model_name'],
|
||||
'workflow_name': metadata['workflow_name'],
|
||||
}
|
||||
await self.emit_metric(
|
||||
self.emit_metric_sync(
|
||||
metric_object=metrics.PREDICTIONS_WRITTEN_COUNT,
|
||||
tags=core_tags,
|
||||
)
|
||||
|
||||
await self.emit_metric(
|
||||
self.emit_metric_sync(
|
||||
metric_object=metrics.PREDICTION_CONFIDENCE_MONITOR,
|
||||
method='set',
|
||||
tags=core_tags,
|
||||
value=prediction_confidence,
|
||||
)
|
||||
|
||||
await self.emit_metric(
|
||||
self.emit_metric_sync(
|
||||
metric_object=metrics.PREDICTION_RESPONSE_TIME_MONITOR,
|
||||
method='observe',
|
||||
tags=core_tags,
|
||||
@@ -781,7 +971,7 @@ class Gates(MinioManager):
|
||||
for server_id, tags in opc_metrics.items():
|
||||
for tag, response_time in tags.items():
|
||||
if response_time is not None:
|
||||
await self.emit_metric(
|
||||
self.emit_metric_sync(
|
||||
metric_object=metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR,
|
||||
method='observe',
|
||||
tags={
|
||||
@@ -792,7 +982,7 @@ class Gates(MinioManager):
|
||||
value=response_time,
|
||||
)
|
||||
|
||||
await self.emit_metric(
|
||||
self.emit_metric_sync(
|
||||
metric_object=metrics.PREDICTION_OPC_WRITING_COUNT,
|
||||
tags={
|
||||
**core_tags,
|
||||
|
||||
@@ -1,16 +1,23 @@
|
||||
from temporalio import activity, workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
import tempfile
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from shutil import rmtree
|
||||
from typing import Any
|
||||
|
||||
import mlflow
|
||||
import numpy as np
|
||||
from pandas import to_datetime
|
||||
import pandas as pd
|
||||
from pandas import DataFrame, to_datetime
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
from sientia_do.observability.logger import Logger
|
||||
from sientia_do.observability.metrics_controller import MetricsController
|
||||
from sientia_do.repository.minio_repository import MinioRepository
|
||||
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
|
||||
from sientia_do.repository.minio_repository_sync import MinioRepository
|
||||
from sientia_do.temporal.constants import (
|
||||
DATETIME_FORMAT,
|
||||
DATETIME_FORMAT_MS_WITH_TZ,
|
||||
@@ -18,58 +25,84 @@ with workflow.unsafe.imports_passed_through():
|
||||
now,
|
||||
)
|
||||
from sientia_do.utils.formatters import create_sample_dict
|
||||
from sientia_model.model_repository.mlflow_repository import SientiaMLflowRepository
|
||||
from sientia_model.model_repository.plugin_store import PluginStore
|
||||
|
||||
from laborious.utils.dataframe_debug import build_dataframe_debug_message
|
||||
from laborious.utils.messages import MISSING_MINIO_REPOSITORY
|
||||
from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
|
||||
from laborious.utils.repository.minio_manager import MinioManager
|
||||
from laborious.utils.repository.model_repository import MLFlowRepository
|
||||
|
||||
|
||||
class MLFlow(MinioManager):
|
||||
class MLFlow(SientiaMonitoring):
|
||||
"""
|
||||
MLFlow integration activities for model inference operations.
|
||||
Temporal activities that talk to MLflow through ``SientiaMLflowRepository`` and ``SientiaModel`` wrappers.
|
||||
|
||||
This class provides activities for interacting with MLFlow models, including
|
||||
data transformation and prediction operations. It handles authentication,
|
||||
data preprocessing, and model management with configurable retention policies.
|
||||
Models are resolved by registered name and the ``production`` alias (not by legacy stages or
|
||||
separate transform/predict flavors). ``get_cached_model`` loads or reuses a wrapper; inference
|
||||
uses ``wrapper.transform`` / ``wrapper.predict``; retrain uses ``wrapper.retrain`` or
|
||||
``wrapper.train`` plus ``store_model`` and registry promotion via ``promote_to_alias``.
|
||||
|
||||
The class implements comprehensive error handling and logging for all
|
||||
MLFlow operations, ensuring reliable model inference in production environments.
|
||||
Large inputs and outputs flow through ``MinioDataFramePayload`` when workflows offload parquet
|
||||
to MinIO. On failure, transform/predict still return a payload with ``success: False`` and
|
||||
error details for downstream gates.
|
||||
|
||||
Attributes:
|
||||
model_monitoring_repository (MLFlowRepository): Repository for MLFlow operations
|
||||
mlflow_repository: Client for tracking, registry, artifact download, and run lifecycle.
|
||||
plugin_store: Reference to the store (runtime is installed on the worker; reserved for
|
||||
future store-backed helpers).
|
||||
"""
|
||||
|
||||
_MAX_DEBUG_DATAFRAME_ROWS = 100
|
||||
_DEFAULT_MODEL_ALIAS = 'production'
|
||||
_REFERENCE_ARTIFACT_CANDIDATES = ('evaluation_data.csv', 'test_data.csv')
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
mlflow_repository: MLFlowRepository,
|
||||
mlflow_repository: SientiaMLflowRepository,
|
||||
plugin_store: PluginStore,
|
||||
minio_repository: MinioRepository | None = None,
|
||||
logger: Logger | None = None,
|
||||
notification_handler: NotificationHandler | None = None,
|
||||
metrics_controller: MetricsController | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize MLFlow activities with an already-built repository.
|
||||
Attach shared MLflow and MinIO clients used by all ML activities in this mixin.
|
||||
|
||||
Args:
|
||||
mlflow_repository: Repository for MLFlow operations, built and owned by the caller
|
||||
(shared with other activities rather than rebuilt here)
|
||||
logger: Logger instance for observability and debugging
|
||||
notification_handler: Notification handler for alerts and monitoring
|
||||
- mlflow_repository: Repository built by ``Activities`` (or injected in tests).
|
||||
- plugin_store: Plugin store instance from worker bootstrap.
|
||||
- minio_repository: MinIO client for ``MinioDataFramePayload`` upload/download.
|
||||
- logger: Structured logger.
|
||||
- notification_handler: Notifications on hard failures where applicable.
|
||||
- metrics_controller: Shared metrics controller.
|
||||
|
||||
Return:
|
||||
None
|
||||
"""
|
||||
MinioManager.__init__(
|
||||
self, minio_repository, logger, notification_handler, metrics_controller
|
||||
|
||||
self.minio_repository = minio_repository
|
||||
SientiaMonitoring.__init__(
|
||||
self,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
self.model_monitoring_repository = mlflow_repository
|
||||
self.mlflow_repository = mlflow_repository
|
||||
self.plugin_store = plugin_store
|
||||
|
||||
def close(self) -> None:
|
||||
"""
|
||||
Close the MLFlow activity and clean up resources.
|
||||
Release MinIO manager resources held by the mixin.
|
||||
|
||||
Return:
|
||||
None
|
||||
"""
|
||||
MinioManager.close(self)
|
||||
if self.minio_repository is not None:
|
||||
try:
|
||||
self.minio_repository.close()
|
||||
finally:
|
||||
self.minio_repository = None
|
||||
SientiaMonitoring.shutdown(self)
|
||||
|
||||
def __del__(self):
|
||||
self.close()
|
||||
@@ -92,53 +125,133 @@ class MLFlow(MinioManager):
|
||||
metadata,
|
||||
)
|
||||
|
||||
@activity.defn(name='request_transform')
|
||||
async def request_transform(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
|
||||
def _detect_and_parse_datetime_index(self, data: pd.DataFrame, metadata: dict) -> pd.DataFrame:
|
||||
"""
|
||||
Transform input data using MLFlow models.
|
||||
Ensure the transform output index is homogeneous and encoded as ``DATETIME_FORMAT_WITH_TZ`` strings.
|
||||
|
||||
This activity processes input data through MLFlow model transformation,
|
||||
including data preprocessing, format conversion, and validation. It handles
|
||||
data deduplication, pivoting, and cleanup to ensure optimal model performance.
|
||||
|
||||
The transformation process includes:
|
||||
1. Data deduplication based on variable and timestamp
|
||||
2. Data pivoting for model input format
|
||||
3. Null value handling and cleanup
|
||||
4. MLFlow model transformation request
|
||||
5. Response validation and logging
|
||||
Accepts an all-string index (validated against the format), or all-``datetime`` /
|
||||
``Timestamp`` (naive timestamps are localized to UTC before formatting). Mixed element types
|
||||
or unsupported types raise ``ValueError`` with a message logged at info level.
|
||||
|
||||
Args:
|
||||
input_data: Configuration and data for transformation
|
||||
Required keys:
|
||||
- metadata (dict): Workflow execution metadata
|
||||
- data (dict): Input data for transformation
|
||||
- model_name (str): Name of the MLFlow model to use
|
||||
- model_retention (int): Model retention period in minutes
|
||||
- data: DataFrame whose index carries the time dimension after transform.
|
||||
- metadata: Workflow metadata for log correlation.
|
||||
|
||||
Returns:
|
||||
dict: Transformed data from MLFlow model
|
||||
Return:
|
||||
``pd.DataFrame``: Same frame with a normalized string index; empty frames are returned as-is.
|
||||
"""
|
||||
|
||||
if data.empty:
|
||||
self.info('Data is empty, skipping datetime index detection and parsing', metadata)
|
||||
return data
|
||||
|
||||
index = data.index
|
||||
index_type = type(index[0])
|
||||
|
||||
self.info(f'Index type: {index_type}', metadata)
|
||||
|
||||
message = (
|
||||
f'Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, '
|
||||
f'string in format {DATETIME_FORMAT_WITH_TZ}.'
|
||||
)
|
||||
|
||||
if not all(isinstance(i, index_type) for i in index):
|
||||
types = map(str, map(type, index))
|
||||
raise ValueError(f'{message}. Elements are {",".join(types)}')
|
||||
|
||||
if index_type is str:
|
||||
try:
|
||||
pd.to_datetime(data.index, format=DATETIME_FORMAT_WITH_TZ)
|
||||
except ValueError as e:
|
||||
raise ValueError(f'{message}. Unable to parse given date format: {e}') from e
|
||||
|
||||
elif index_type is datetime or index_type is pd.Timestamp:
|
||||
idx = data.index
|
||||
if hasattr(idx, 'tz') and idx.tz is None:
|
||||
data.index = idx.tz_localize('UTC') # type: ignore[attr-defined]
|
||||
|
||||
data.index = data.index.strftime(DATETIME_FORMAT_WITH_TZ) # type: ignore[attr-defined]
|
||||
else:
|
||||
raise ValueError(f'{message}. Got {index_type}.')
|
||||
|
||||
return data
|
||||
|
||||
def _resolve_model_version_for_run(self, run_id: str) -> str:
|
||||
"""
|
||||
Map an MLflow ``run_id`` to the latest registered model version that produced that run.
|
||||
|
||||
``search_model_versions`` may return multiple versions if the model was registered more than
|
||||
once for the same run; the highest numeric ``version`` wins so promotion targets the newest
|
||||
artifact set.
|
||||
|
||||
Args:
|
||||
- run_id: Run UUID from ``retrain_model`` / experiment payload.
|
||||
|
||||
Return:
|
||||
str: Registry version string acceptable by ``promote_to_alias``.
|
||||
|
||||
Raises:
|
||||
Exception: If transformation fails or MLFlow model is unavailable
|
||||
ValueError: If the filter returns no versions (model not registered for this run).
|
||||
"""
|
||||
|
||||
versions = self.mlflow_repository._client.search_model_versions(
|
||||
filter_string=f"run_id='{run_id}'"
|
||||
)
|
||||
if not versions:
|
||||
raise ValueError(f'No registered model version found for run_id={run_id}')
|
||||
latest = max(versions, key=lambda v: int(v.version))
|
||||
return str(latest.version)
|
||||
|
||||
def _resolve_model_alias(self, model_config: dict[str, Any] | None = None) -> str:
|
||||
"""
|
||||
Resolve which MLflow alias should be used for model lookup/promotion.
|
||||
|
||||
Args:
|
||||
- model_config: Optional model configuration that may include ``alias``.
|
||||
|
||||
Return:
|
||||
str: Alias name trimmed and normalized; defaults to ``production``.
|
||||
"""
|
||||
if not model_config:
|
||||
return self._DEFAULT_MODEL_ALIAS
|
||||
alias = str(model_config.get('alias', self._DEFAULT_MODEL_ALIAS)).strip()
|
||||
return alias or self._DEFAULT_MODEL_ALIAS
|
||||
|
||||
@activity.defn(name='request_transform')
|
||||
def request_transform(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
|
||||
"""
|
||||
Pivot long-format sensor rows, load the production wrapper, and run ``wrapper.transform``.
|
||||
|
||||
Expected tabular shape after load: columns including ``variable``, ``timestamp``, ``value``,
|
||||
and ``created_at`` for deduplication. Data are sorted by ``created_at``, de-duplicated per
|
||||
``(variable, timestamp)``, pivoted wide, then passed to the model. ``model_config`` may
|
||||
include ``retention_minutes`` for wrapper cache TTL.
|
||||
|
||||
Args:
|
||||
- input_data: Dict with ``metadata``, ``model_name``, ``data`` (``MinioDataFramePayload``
|
||||
dict or inline dataframe dict), and optional ``model_config``.
|
||||
|
||||
Return:
|
||||
``MinioDataFramePayload`` with transformed frame and ``success: True``, or a payload
|
||||
with ``success: False`` and exception details in ``status`` if transform fails.
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
self.info('Transforming data...', metadata)
|
||||
|
||||
payload = MinioDataFramePayload.from_dict(input_data['data'])
|
||||
data = await payload.retrieve(self.minio_repository, metadata)
|
||||
data = payload.retrieve(self.minio_repository, metadata)
|
||||
|
||||
model_name = input_data['model_name']
|
||||
model_config = input_data.get('model_config', {})
|
||||
model_alias = self._resolve_model_alias(model_config)
|
||||
|
||||
self._debug_dataframe('Raw input data:', data, metadata)
|
||||
|
||||
# Sort by created_at in descending order and keep first occurrence of each variable/timestamp pair
|
||||
# Long → wide: keep newest row per (variable, timestamp), then pivot for the wrapper API.
|
||||
data = data.sort_values('created_at', ascending=False).drop_duplicates(
|
||||
subset=['variable', 'timestamp'], keep='first'
|
||||
)
|
||||
|
||||
# Pivot data for model input format
|
||||
data = data.pivot(index='timestamp', columns='variable', values='value')
|
||||
data.fillna(np.nan, inplace=True)
|
||||
|
||||
@@ -149,10 +262,25 @@ class MLFlow(MinioManager):
|
||||
|
||||
self._debug_dataframe('Processed input data:', data, metadata)
|
||||
|
||||
# Request transformation from MLFlow model
|
||||
response_data = await self.model_monitoring_repository.transform(
|
||||
model_name, data, model_config, metadata
|
||||
)
|
||||
try:
|
||||
wrapper = self.mlflow_repository.get_cached_model(
|
||||
model_name=model_name,
|
||||
alias=model_alias,
|
||||
retention_minutes=model_config.get('retention_minutes', 0),
|
||||
metadata=metadata,
|
||||
)
|
||||
transformed_df, transform_meta = wrapper.transform(data)
|
||||
if transform_meta:
|
||||
self.info(f'Wrapper transform metadata: {transform_meta}', metadata)
|
||||
|
||||
transformed_df = self._detect_and_parse_datetime_index(transformed_df, metadata)
|
||||
|
||||
response_data: dict[str, Any] = {'success': True, 'content': transformed_df}
|
||||
except Exception as e:
|
||||
response_data = {
|
||||
'success': False,
|
||||
'content': {'message': str(e), 'traceback': traceback.format_exc()},
|
||||
}
|
||||
|
||||
self.debug(
|
||||
f'Transform raw response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
|
||||
@@ -167,7 +295,7 @@ class MLFlow(MinioManager):
|
||||
self.info('Data transformed successfully', metadata)
|
||||
|
||||
if not response_data.get('success', False):
|
||||
return await MinioDataFramePayload.from_dataframe(
|
||||
return MinioDataFramePayload.from_dataframe(
|
||||
dataframe=None,
|
||||
minio_repo=self.minio_repository,
|
||||
model_name=model_name,
|
||||
@@ -178,7 +306,7 @@ class MLFlow(MinioManager):
|
||||
logger=self.logger,
|
||||
)
|
||||
|
||||
return await MinioDataFramePayload.from_dataframe(
|
||||
return MinioDataFramePayload.from_dataframe(
|
||||
dataframe=response_data['content'],
|
||||
minio_repo=self.minio_repository,
|
||||
model_name=model_name,
|
||||
@@ -192,62 +320,78 @@ class MLFlow(MinioManager):
|
||||
)
|
||||
|
||||
@activity.defn(name='request_predict')
|
||||
async def request_predict(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
|
||||
def request_predict(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
|
||||
"""
|
||||
Execute predictions using MLFlow models.
|
||||
Load the production wrapper and call ``wrapper.predict`` on the prepared feature frame.
|
||||
|
||||
This activity performs ML model inference using MLFlow models with the
|
||||
transformed data. It handles data format conversion, null value processing,
|
||||
and model prediction requests with comprehensive error handling.
|
||||
|
||||
The prediction process includes:
|
||||
1. Data format validation and cleanup
|
||||
2. Null value handling for model compatibility
|
||||
3. MLFlow model prediction request
|
||||
4. Response validation and logging
|
||||
5. Performance monitoring and metrics
|
||||
The activity normalizes ``NaN`` to ``None`` for JSON-friendly columns, sets the row index
|
||||
the same way as ``retrain_model`` (UTC ``DatetimeIndex`` from ``DATETIME_FORMAT_WITH_TZ``),
|
||||
restores that index on the prediction frame, normalizes the prediction index to
|
||||
``DATETIME_FORMAT_WITH_TZ`` strings like ``request_transform``, and records ``response_time``.
|
||||
Non-DataFrame predictions are coerced to a single ``prediction`` column.
|
||||
|
||||
Args:
|
||||
input_data: Configuration and data for prediction
|
||||
Required keys:
|
||||
- metadata (dict): Workflow execution metadata
|
||||
- data (dict): Transformed data for prediction
|
||||
- model_name (str): Name of the MLFlow model to use
|
||||
- model_retention (int): Model retention period in minutes
|
||||
Optional keys (in model_config):
|
||||
- skip_predict_timestamp (bool): If True, don't reinsert the
|
||||
timestamp column from the index before predicting. Default False.
|
||||
- input_data: Same envelope as ``request_transform`` (``metadata``, ``model_name``,
|
||||
``data``, optional ``model_config`` with ``retention_minutes``).
|
||||
|
||||
Returns:
|
||||
dict: Prediction results from MLFlow model
|
||||
|
||||
Raises:
|
||||
Exception: If prediction fails or MLFlow model is unavailable
|
||||
Return:
|
||||
``MinioDataFramePayload`` with predictions or error status mirroring transform behaviour.
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
self.info('Predicting data...', metadata)
|
||||
|
||||
payload = MinioDataFramePayload.from_dict(input_data['data'])
|
||||
data = await payload.retrieve(self.minio_repository, metadata)
|
||||
data = payload.retrieve(self.minio_repository, metadata)
|
||||
|
||||
model_name = input_data['model_name']
|
||||
model_config = input_data.get('model_config', {})
|
||||
model_alias = self._resolve_model_alias(model_config)
|
||||
|
||||
self._debug_dataframe('Input data for prediction:', data, metadata)
|
||||
|
||||
# Convert numpy.nan to None for model compatibility
|
||||
data.replace(np.nan, None, inplace=True)
|
||||
|
||||
if not model_config.get('skip_predict_timestamp', False):
|
||||
data['timestamp'] = data.index
|
||||
data['timestamp'] = to_datetime(
|
||||
data['timestamp'], format=DATETIME_FORMAT_WITH_TZ
|
||||
).dt.strftime(DATETIME_FORMAT)
|
||||
|
||||
# Request prediction from MLFlow model
|
||||
response_data = await self.model_monitoring_repository.predict(
|
||||
model_name, data, model_config, metadata
|
||||
data.index = pd.DatetimeIndex(
|
||||
to_datetime(data.index, format=DATETIME_FORMAT_WITH_TZ, utc=True)
|
||||
)
|
||||
input_index = data.index
|
||||
|
||||
try:
|
||||
wrapper = self.mlflow_repository.get_cached_model(
|
||||
model_name=model_name,
|
||||
alias=model_alias,
|
||||
retention_minutes=model_config.get('retention_minutes', 0),
|
||||
metadata=metadata,
|
||||
)
|
||||
start_time = datetime.now()
|
||||
predict_data, pred_meta = wrapper.predict({}, data)
|
||||
end_time = datetime.now()
|
||||
|
||||
if pred_meta:
|
||||
self.info(f'Wrapper predict metadata: {pred_meta}', metadata)
|
||||
|
||||
if isinstance(predict_data, DataFrame):
|
||||
self._debug_dataframe(
|
||||
'Data received from model prediction:', predict_data, metadata
|
||||
)
|
||||
predict_data.columns = pd.Index(['prediction'])
|
||||
else:
|
||||
self.debug(
|
||||
f'Data received from model prediction (not a DataFrame): {predict_data}',
|
||||
metadata,
|
||||
)
|
||||
predict_data = pd.DataFrame(predict_data, columns=['prediction'])
|
||||
|
||||
predict_data.index = input_index
|
||||
predict_data['response_time'] = (end_time - start_time).total_seconds()
|
||||
predict_data = self._detect_and_parse_datetime_index(predict_data, metadata)
|
||||
|
||||
response_data: dict[str, Any] = {'success': True, 'content': predict_data}
|
||||
except Exception as e:
|
||||
response_data = {
|
||||
'success': False,
|
||||
'content': {'message': str(e), 'traceback': traceback.format_exc()},
|
||||
}
|
||||
|
||||
self.debug(
|
||||
f'Prediction response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
|
||||
@@ -257,7 +401,7 @@ class MLFlow(MinioManager):
|
||||
self.info('Data predicted successfully', metadata)
|
||||
|
||||
if not response_data.get('success', False):
|
||||
return await MinioDataFramePayload.from_dataframe(
|
||||
return MinioDataFramePayload.from_dataframe(
|
||||
dataframe=None,
|
||||
minio_repo=self.minio_repository,
|
||||
model_name=model_name,
|
||||
@@ -268,7 +412,7 @@ class MLFlow(MinioManager):
|
||||
logger=self.logger,
|
||||
)
|
||||
|
||||
return await MinioDataFramePayload.from_dataframe(
|
||||
return MinioDataFramePayload.from_dataframe(
|
||||
dataframe=response_data['content'],
|
||||
minio_repo=self.minio_repository,
|
||||
model_name=model_name,
|
||||
@@ -282,36 +426,23 @@ class MLFlow(MinioManager):
|
||||
)
|
||||
|
||||
@activity.defn(name='retrain_model')
|
||||
async def retrain_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
def retrain_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Retrain MLFlow models with updated training data.
|
||||
Fit an updated wrapper from historical data, then log and register in MLflow.
|
||||
|
||||
This activity orchestrates the complete model retraining process,
|
||||
including data preparation, model retraining execution, and result
|
||||
validation. It handles data preprocessing, column cleanup, and
|
||||
comprehensive error handling for production model management.
|
||||
|
||||
The retraining process includes:
|
||||
1. Data timestamp extraction and validation
|
||||
2. Column cleanup and data preparation
|
||||
3. Data pivoting for model input format
|
||||
4. MLFlow model retraining execution
|
||||
5. Result validation and error handling
|
||||
Flow: load long-format data from MinIO → dedupe/pivot like inference prep → require
|
||||
``model_config['target']`` → read current ``production`` version for ``source_run_id`` tag →
|
||||
run ``wrapper.retrain`` outside run timing → ``start_run`` with retrain tags → log input
|
||||
CSV artifact → ``store_model`` and ``log_params``. Does not promote; the workflow calls
|
||||
``update_production_model`` after validation.
|
||||
|
||||
Args:
|
||||
input_data (dict): Input data containing:
|
||||
- metadata (dict): Workflow execution metadata
|
||||
- data (dict[str, Any]): Training data for model retraining
|
||||
- model_name (str): Name of the MLFlow model to retrain
|
||||
- input_data: Must include ``metadata``, ``model_name``, ``data`` (payload), and
|
||||
``model_config`` with at least ``target``.
|
||||
|
||||
Returns:
|
||||
dict: Retraining results containing:
|
||||
- status (str): Retraining operation status
|
||||
- timestamp (str): Timestamp of the retraining operation
|
||||
- experiment (str): MLFlow experiment identifier
|
||||
|
||||
Raises:
|
||||
Exception: If retraining fails or encounters critical errors
|
||||
Return:
|
||||
On success: ``success``, ``experiment`` (``run_id``, ``experiment_id``, ``experiment_name``),
|
||||
``message``, ``timestamp``. On failure: ``success: False``, error fields, and optional trace.
|
||||
"""
|
||||
|
||||
if self.minio_repository is None:
|
||||
@@ -320,13 +451,12 @@ class MLFlow(MinioManager):
|
||||
metadata = input_data['metadata']
|
||||
|
||||
try:
|
||||
# Payload-based retrain input (inline dict or MinIO offloaded).
|
||||
payload = MinioDataFramePayload.from_dict(input_data['data'])
|
||||
data = await payload.retrieve(self.minio_repository, metadata)
|
||||
data = payload.retrieve(self.minio_repository, metadata)
|
||||
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
await self.send_notification_async(
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='ERROR_LOADING_RETRAIN_DATA',
|
||||
message=f'Error loading retrain data: {e}',
|
||||
@@ -352,7 +482,6 @@ class MLFlow(MinioManager):
|
||||
timestamp = data['timestamp'].max()
|
||||
self.debug(f'Timestamp: {timestamp}', metadata)
|
||||
|
||||
# Sort by created_at in descending order and keep first occurrence of each variable/timestamp pair
|
||||
if 'created_at' in data.columns:
|
||||
data = data.sort_values('created_at', ascending=False).drop_duplicates(
|
||||
subset=['variable', 'timestamp'], keep='first'
|
||||
@@ -363,74 +492,131 @@ class MLFlow(MinioManager):
|
||||
data.drop(columns=['model_id'], inplace=True, errors='ignore')
|
||||
data.drop(columns=['created_at'], inplace=True, errors='ignore')
|
||||
|
||||
# Pivot data for model input format
|
||||
data = data.pivot(index='timestamp', columns='variable', values='value')
|
||||
data.fillna(np.nan, inplace=True)
|
||||
# data.reset_index(inplace=True)
|
||||
data.columns.name = None
|
||||
data.index.name = None
|
||||
|
||||
data['timestamp'] = data.index
|
||||
data['timestamp'] = to_datetime(
|
||||
data['timestamp'], format=DATETIME_FORMAT_WITH_TZ
|
||||
).dt.strftime(DATETIME_FORMAT)
|
||||
data['timestamp'] = to_datetime(data['timestamp'], format=DATETIME_FORMAT)
|
||||
|
||||
data.columns.name = None
|
||||
|
||||
retrain_output = await self.model_monitoring_repository.retrain_model(
|
||||
data=data, model_name=model_name, model_config=model_config, metadata=metadata
|
||||
data.index = pd.DatetimeIndex(
|
||||
to_datetime(data.index, format=DATETIME_FORMAT_WITH_TZ, utc=True)
|
||||
)
|
||||
|
||||
if not retrain_output['success']:
|
||||
trace = retrain_output['traceback']
|
||||
await self.send_notification_async(
|
||||
metadata=metadata,
|
||||
notification_id='RETRAIN_MODEL_ERROR',
|
||||
message=f'Error retraining model {model_name}: {retrain_output["message"]}',
|
||||
block='retrain_model',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace,
|
||||
)
|
||||
self.error(trace, metadata=metadata)
|
||||
target = model_config.get('target')
|
||||
if target is None:
|
||||
msg = 'model_config must include "target" for retraining'
|
||||
self.info(msg, metadata)
|
||||
return {
|
||||
'success': False,
|
||||
'experiment': None,
|
||||
'message': msg,
|
||||
'traceback': '',
|
||||
'timestamp': str(timestamp),
|
||||
}
|
||||
|
||||
return {**retrain_output, 'timestamp': timestamp}
|
||||
try:
|
||||
model_alias = self._resolve_model_alias(model_config)
|
||||
mv_src = self.mlflow_repository._client.get_model_version_by_alias(
|
||||
name=model_name,
|
||||
alias=model_alias,
|
||||
)
|
||||
source_run_id = mv_src.run_id
|
||||
|
||||
wrapper = self.mlflow_repository.get_cached_model(
|
||||
model_name=model_name,
|
||||
alias=model_alias,
|
||||
retention_minutes=0,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
# Keep heavy model fitting outside MLflow run timing.
|
||||
prediction_data = wrapper.retrain(data).rename(columns={target: 'prediction'})
|
||||
|
||||
# Both frames are indexed by the pivoted timestamp: expose it as an explicit column
|
||||
# and join on it. `validate` fails loudly instead of silently multiplying rows if the
|
||||
# wrapper ever returns duplicated timestamps.
|
||||
evaluation_data = (
|
||||
data.reset_index(names='timestamp')
|
||||
.merge(
|
||||
prediction_data.reset_index(names='timestamp'),
|
||||
on='timestamp',
|
||||
how='left',
|
||||
validate='one_to_one',
|
||||
)
|
||||
.rename(columns={target: 'target'})
|
||||
.sort_values(by='timestamp', ascending=True, ignore_index=True)
|
||||
)
|
||||
|
||||
run_name = f'{model_name}-retrain-{datetime.now().strftime("%Y%m%d%H%M%S")}'
|
||||
|
||||
with self.mlflow_repository.start_run(
|
||||
model_name=model_name,
|
||||
run_name=run_name,
|
||||
experiment_name=model_name,
|
||||
tags={'retrain': 'true', 'source_run_id': source_run_id},
|
||||
metadata=metadata,
|
||||
) as run_info:
|
||||
tmp_dir = tempfile.mkdtemp(prefix='laborious_retrain_')
|
||||
try:
|
||||
raw_csv = Path(tmp_dir) / 'retrain_input.csv'
|
||||
evaluation_csv = Path(tmp_dir) / 'evaluation_data.csv'
|
||||
data.to_csv(raw_csv, index=False)
|
||||
evaluation_data.to_csv(evaluation_csv, index=False)
|
||||
|
||||
mlflow.log_artifact(str(raw_csv))
|
||||
mlflow.log_artifact(str(evaluation_csv))
|
||||
finally:
|
||||
rmtree(tmp_dir, ignore_errors=True)
|
||||
|
||||
wrapper.store_model(name=model_name)
|
||||
|
||||
self.mlflow_repository.log_params(
|
||||
{
|
||||
'retrain': 'true',
|
||||
'retrain_date': datetime.now().isoformat(),
|
||||
'source_run_id': source_run_id,
|
||||
'retrain_samples': str(data.shape),
|
||||
}
|
||||
)
|
||||
|
||||
experiment_payload = {
|
||||
'run_id': run_info.run_id,
|
||||
'experiment_id': run_info.experiment_id,
|
||||
'experiment_name': model_name,
|
||||
}
|
||||
|
||||
return {
|
||||
'success': True,
|
||||
'experiment': experiment_payload,
|
||||
'message': 'Model retrained successfully.',
|
||||
'timestamp': str(timestamp),
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f'Error retraining model {model_name}: {e}'
|
||||
self.info(error_msg, metadata)
|
||||
return {
|
||||
'success': False,
|
||||
'experiment': None,
|
||||
'message': error_msg,
|
||||
'traceback': traceback.format_exc(),
|
||||
'timestamp': str(timestamp),
|
||||
}
|
||||
|
||||
@activity.defn(name='update_production_model')
|
||||
async def update_production_model(self, input_data: dict[str, Any]) -> dict[Any, Any]:
|
||||
def update_production_model(self, input_data: dict[str, Any]) -> dict[Any, Any]:
|
||||
"""
|
||||
Update production model with newly trained model version.
|
||||
Point the ``production`` alias at the model version registered for the retrain run.
|
||||
|
||||
This activity manages the critical process of updating production
|
||||
models with newly trained versions. It handles model deployment,
|
||||
status tracking, and comprehensive reporting for operational
|
||||
visibility and audit trails.
|
||||
|
||||
The update process includes:
|
||||
1. Production model update execution
|
||||
2. Status and metadata tracking
|
||||
3. Comprehensive reporting and logging
|
||||
4. Error handling and notification
|
||||
5. Audit trail maintenance
|
||||
Resolves the highest numeric registry version whose ``run_id`` matches
|
||||
``experiment['run_id']``, then calls ``promote_to_alias``. On failure, sends a notification
|
||||
and re-raises so the workflow can surface the error.
|
||||
|
||||
Args:
|
||||
input_data (dict): Input data containing:
|
||||
- metadata (dict): Workflow execution metadata
|
||||
- model_name (str): Name of the MLFlow model to update
|
||||
- experiment (str): MLFlow experiment identifier
|
||||
- model_id (str): Unique identifier for the model version
|
||||
- timestamp (str): Timestamp of the update operation
|
||||
- status (str): Current status of the model update
|
||||
- input_data: ``metadata``, ``model_name``, and ``experiment`` with ``run_id`` and
|
||||
``experiment_id`` (as returned from ``retrain_model``).
|
||||
|
||||
Returns:
|
||||
dict[Any, Any]: Comprehensive update report containing:
|
||||
- model_id (str): Model version identifier
|
||||
- model_name (str): Name of the updated model
|
||||
- timestamp (str): Update operation timestamp
|
||||
- status (str): Update operation status
|
||||
- Additional MLFlow response metadata
|
||||
|
||||
Raises:
|
||||
Exception: If production model update fails
|
||||
Return:
|
||||
Dict with ``model_name``, promoted ``version``, ``mlflow_run_id``, ``mlflow_experiment_id``.
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
model_name = input_data['model_name']
|
||||
@@ -440,16 +626,30 @@ class MLFlow(MinioManager):
|
||||
)
|
||||
|
||||
try:
|
||||
response = await self.model_monitoring_repository.update_production_model(
|
||||
experiment=experiment, model_name=model_name, metadata=metadata
|
||||
run_id = experiment['run_id']
|
||||
experiment_id = experiment['experiment_id']
|
||||
|
||||
version = self._resolve_model_version_for_run(run_id)
|
||||
|
||||
promote_alias = self._resolve_model_alias(input_data.get('model_config'))
|
||||
self.mlflow_repository.promote_to_alias(
|
||||
model_name=model_name,
|
||||
version=version,
|
||||
alias=promote_alias,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
self.info(f'Production model {model_name} updated successfully', metadata)
|
||||
return response
|
||||
return {
|
||||
'model_name': model_name,
|
||||
'version': version,
|
||||
'mlflow_run_id': run_id,
|
||||
'mlflow_experiment_id': experiment_id,
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
await self.send_notification_async(
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='UPDATE_PRODUCTION_MODEL_ERROR',
|
||||
message=f'Error updating production model {model_name}: {e}',
|
||||
@@ -460,50 +660,98 @@ class MLFlow(MinioManager):
|
||||
self.error(trace, metadata=metadata)
|
||||
raise e
|
||||
|
||||
@activity.defn(name='get_reference_data')
|
||||
async def get_reference_data(self, input_data: dict[str, Any]) -> list[dict] | None:
|
||||
def _resolve_reference_artifact_name(self, run_id: str) -> str | None:
|
||||
"""
|
||||
Get reference data from the MLflow Model Registry.
|
||||
Pick the first available reference CSV artifact path from the MLflow run.
|
||||
|
||||
This method retrieves evaluation reference data stored as artifacts in the
|
||||
MLflow Model Registry. The reference data is typically used for model
|
||||
drift detection, performance comparison, and quality validation. The method
|
||||
loads the data from a CSV artifact file and formats timestamps for
|
||||
consistent processing.
|
||||
|
||||
The method handles:
|
||||
1. Loading evaluation data artifact from MLflow Model Registry
|
||||
2. Timestamp parsing and formatting for consistency
|
||||
3. Data conversion to dictionary format for workflow consumption
|
||||
4. Graceful handling of missing reference data
|
||||
Candidates are checked in priority order: ``evaluation_data.csv``, then ``test_data.csv``.
|
||||
A path matches when it equals the candidate or ends with ``/<candidate>`` for nested layouts.
|
||||
|
||||
Args:
|
||||
input_data (dict): Input data containing:
|
||||
- metadata (dict): Workflow execution metadata
|
||||
- model_name (str): Name of the MLFlow model to get reference data from
|
||||
- run_id: MLflow run UUID linked to the production model version.
|
||||
|
||||
Returns:
|
||||
list[dict[Hashable, Any]] | None: Reference data from the MLflow Model Registry
|
||||
as a list of dictionaries. Returns None if reference data is not found
|
||||
or if the artifact does not exist.
|
||||
|
||||
Raises:
|
||||
Exception: If artifact loading fails or encounters errors during processing
|
||||
Return:
|
||||
Artifact path string for ``download_artifacts``, or ``None`` if no candidate exists.
|
||||
"""
|
||||
listed = self.mlflow_repository._client.list_artifacts(run_id)
|
||||
paths = [file_info.path for file_info in listed]
|
||||
for candidate in self._REFERENCE_ARTIFACT_CANDIDATES:
|
||||
for path in paths:
|
||||
if path == candidate or path.endswith(f'/{candidate}'):
|
||||
return path
|
||||
return None
|
||||
|
||||
def _find_downloaded_csv(self, tmpdir: str, artifact_name: str) -> Path | None:
|
||||
"""
|
||||
Locate a downloaded reference CSV in the temp directory.
|
||||
|
||||
Args:
|
||||
- tmpdir: Directory where ``download_artifacts`` wrote files.
|
||||
- artifact_name: Basename of the resolved artifact (e.g. ``retrain_input.csv``).
|
||||
|
||||
Return:
|
||||
``Path`` to the CSV file if found, else ``None``.
|
||||
"""
|
||||
direct = Path(tmpdir) / artifact_name
|
||||
if direct.exists():
|
||||
return direct
|
||||
matches = list(Path(tmpdir).rglob(artifact_name))
|
||||
return matches[0] if matches else None
|
||||
|
||||
@activity.defn(name='get_reference_data')
|
||||
def get_reference_data(self, input_data: dict[str, Any]) -> list[dict] | None:
|
||||
"""
|
||||
Download reference training CSV from the MLflow run linked to the production alias.
|
||||
|
||||
Resolves ``evaluation_data.csv`` or ``test_data.csv`` via artifact listing before download.
|
||||
``evaluation_data.csv`` is preferred when both exist. Used by drift workflows to compare
|
||||
live data against the reference distribution logged with the model.
|
||||
Timestamps are normalized to ``DATETIME_FORMAT`` string columns before returning records.
|
||||
|
||||
Args:
|
||||
- input_data: ``metadata``, ``model_name``, and optional ``model_config`` with ``alias``.
|
||||
|
||||
Return:
|
||||
List of row dicts with normalized timestamps, or ``None`` if resolution or load fails.
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
model_name = input_data['model_name']
|
||||
artifact = 'evaluation_data.csv'
|
||||
|
||||
reference_data = await self.model_monitoring_repository.load_artifact_dataframe(
|
||||
model_name=model_name, artifact_path=artifact, metadata=metadata
|
||||
)
|
||||
try:
|
||||
model_alias = self._resolve_model_alias(input_data.get('model_config'))
|
||||
mv = self.mlflow_repository._client.get_model_version_by_alias(
|
||||
name=model_name,
|
||||
alias=model_alias,
|
||||
)
|
||||
run_id = mv.run_id
|
||||
|
||||
if reference_data is None:
|
||||
self.warning(f'Reference data not found for model {model_name}', metadata)
|
||||
artifact_path = self._resolve_reference_artifact_name(run_id)
|
||||
if artifact_path is None:
|
||||
self.warning(f'Reference data not found for model {model_name}', metadata)
|
||||
return None
|
||||
|
||||
artifact_name = Path(artifact_path).name
|
||||
tmpdir = tempfile.mkdtemp(prefix='laborious_eval_')
|
||||
try:
|
||||
self.mlflow_repository.download_artifacts(
|
||||
run_id=run_id,
|
||||
artifact_path=artifact_path,
|
||||
dst_path=tmpdir,
|
||||
metadata=metadata,
|
||||
)
|
||||
csv_path = self._find_downloaded_csv(tmpdir, artifact_name)
|
||||
if csv_path is None:
|
||||
self.warning(f'Reference data not found for model {model_name}', metadata)
|
||||
return None
|
||||
reference_data = pd.read_csv(csv_path)
|
||||
finally:
|
||||
rmtree(tmpdir, ignore_errors=True)
|
||||
|
||||
reference_data['timestamp'] = to_datetime(reference_data['timestamp'])
|
||||
reference_data['timestamp'] = reference_data['timestamp'].dt.strftime(DATETIME_FORMAT)
|
||||
|
||||
return reference_data.to_dict(orient='records')
|
||||
|
||||
except Exception as e:
|
||||
self.warning(f'Reference data not found for model {model_name}: {e}', metadata)
|
||||
return None
|
||||
|
||||
reference_data['timestamp'] = to_datetime(reference_data['timestamp'])
|
||||
reference_data['timestamp'] = reference_data['timestamp'].dt.strftime(DATETIME_FORMAT)
|
||||
|
||||
return reference_data.to_dict(orient='records')
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -7,14 +7,15 @@ with workflow.unsafe.imports_passed_through():
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
from pandas import DataFrame, Index, to_datetime
|
||||
from sientia.ModelAnalysis import ModelAnalysis
|
||||
import pandas as pd
|
||||
from pandas import DataFrame, Index, Series, to_datetime
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
from sientia_do.observability.logger import Logger
|
||||
from sientia_do.observability.metrics_controller import MetricsController
|
||||
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
|
||||
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
|
||||
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
|
||||
from sientia_model.analytics.drift_analysis import DriftAnalysis, DriftInsufficientDataError
|
||||
|
||||
from laborious import metrics
|
||||
from laborious.utils.dataframe_debug import build_dataframe_debug_message
|
||||
@@ -27,9 +28,12 @@ warnings.filterwarnings(
|
||||
|
||||
class ModelMetrics(SientiaMonitoring):
|
||||
"""
|
||||
Metrics activities for the Laborious system.
|
||||
Metrics and statistical analysis activities for the Laborious pipeline.
|
||||
|
||||
This class provides activities for writing metrics to the Prometheus monitoring system.
|
||||
This class centralizes drift/statistical computations and model-quality
|
||||
aggregates used by scheduled workflows. Besides producing tabular outputs
|
||||
for persistence, it also emits operational metrics (count, lag, error)
|
||||
through ``SientiaMonitoring`` so execution health is observable in runtime.
|
||||
"""
|
||||
|
||||
_MAX_DEBUG_DATAFRAME_ROWS = 100
|
||||
@@ -44,7 +48,10 @@ class ModelMetrics(SientiaMonitoring):
|
||||
|
||||
def close(self) -> None:
|
||||
"""
|
||||
Close the model metrics activity and clean up resources.
|
||||
Shutdown monitoring resources associated with model metrics activities.
|
||||
|
||||
This is invoked during worker teardown to flush/close metric controller
|
||||
internals and prevent dangling telemetry tasks.
|
||||
"""
|
||||
SientiaMonitoring.shutdown(self)
|
||||
|
||||
@@ -69,7 +76,26 @@ class ModelMetrics(SientiaMonitoring):
|
||||
metadata,
|
||||
)
|
||||
|
||||
async def get_drift_metrics(
|
||||
def _drift_analyze_stage_error(
|
||||
self,
|
||||
exc: Exception,
|
||||
context: str,
|
||||
metadata: dict[str, Any],
|
||||
core_labels: dict[str, Any],
|
||||
) -> None:
|
||||
"""
|
||||
Log analyzer failure for a drift stage and increment the analyze error metric.
|
||||
|
||||
Args:
|
||||
- exc (Exception): Failure raised by ``sientia_model``.
|
||||
- context (str): Short label for the log line (e.g. univariate detection).
|
||||
- metadata (dict[str, Any]): Workflow metadata for logging.
|
||||
- core_labels (dict[str, Any]): Tags from ``get_core_labels`` for metrics.
|
||||
"""
|
||||
self.error(f'{context}: {exc}', metadata)
|
||||
self.emit_metric_sync(metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels)
|
||||
|
||||
def get_drift_metrics(
|
||||
self,
|
||||
reference_data: DataFrame,
|
||||
target_data: DataFrame,
|
||||
@@ -80,24 +106,38 @@ class ModelMetrics(SientiaMonitoring):
|
||||
metadata: dict[str, Any],
|
||||
) -> DataFrame:
|
||||
"""
|
||||
Calculate univariate drift metrics for a model.
|
||||
Compute univariate and multivariate drift outputs and merge them into one dataframe.
|
||||
|
||||
The method orchestrates three analysis stages (univariate drift,
|
||||
multivariate drift, and dataframe projection), emitting lag/count/error
|
||||
metrics for each stage independently so failures are attributable.
|
||||
|
||||
Args:
|
||||
model_analysis (ModelAnalysis): Model analysis object
|
||||
reference_data (DataFrame): Reference data
|
||||
target_data (DataFrame): Target data
|
||||
reference_columns (list[str]): Reference columns
|
||||
drift_metrics (list[str]): Drift metrics
|
||||
metadata (dict[str, Any]): Workflow execution metadata
|
||||
- reference_data (DataFrame): Baseline dataset representing expected behavior.
|
||||
- target_data (DataFrame): Current analysis dataset to compare against reference.
|
||||
- target_name (str): Target column name used by ``DriftAnalysis`` config.
|
||||
- reference_columns (Index): Feature columns evaluated for drift.
|
||||
- drift_metrics (list[str]): Enabled univariate methods.
|
||||
- chunk_period (str): Time bucket granularity used by analysis methods.
|
||||
- metadata (dict[str, Any]): Workflow metadata for logs and notifications.
|
||||
|
||||
Return:
|
||||
DataFrame: Consolidated drift dataframe from ``get_drift_metrics_dataframe`` using
|
||||
``method`` / ``value`` (and optional ``threshold``, ``drift_type``), ready for
|
||||
activity-level formatting before Postgres export.
|
||||
"""
|
||||
# ``DriftAnalysis`` uses truthiness checks on ``features`` (e.g. ``if not features``);
|
||||
# a pandas ``Index`` is ambiguous in boolean context — normalize to a list.
|
||||
feature_names: list[str] = list(reference_columns)
|
||||
|
||||
config = {
|
||||
'target': target_name,
|
||||
'prediction': 'prediction',
|
||||
'timestamp': 'timestamp',
|
||||
'features': reference_columns,
|
||||
'features': feature_names,
|
||||
}
|
||||
|
||||
model_analysis = ModelAnalysis(config=config)
|
||||
drift_analysis = DriftAnalysis(config=config)
|
||||
|
||||
self._debug_dataframe(
|
||||
f'Reference data: Size {reference_data.shape}', reference_data, metadata
|
||||
@@ -108,64 +148,88 @@ class ModelMetrics(SientiaMonitoring):
|
||||
core_labels = self.get_core_labels(metadata, operation_type='detect_univariate_drift')
|
||||
start_time = time.time()
|
||||
try:
|
||||
univariate_drift = model_analysis.detect_univariate_drift(
|
||||
univariate_drift = drift_analysis.detect_univariate_drift(
|
||||
reference_df=reference_data,
|
||||
analysis_df=target_data,
|
||||
features=reference_columns,
|
||||
features=feature_names,
|
||||
timestamp_col=config['timestamp'],
|
||||
methods=drift_metrics,
|
||||
chunk_period=chunk_period,
|
||||
)
|
||||
except Exception as e:
|
||||
self.error(f'Error detecting univariate drift: {e}', metadata)
|
||||
await self.emit_metric(
|
||||
metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels
|
||||
if isinstance(e, DriftInsufficientDataError):
|
||||
raise
|
||||
self._drift_analyze_stage_error(
|
||||
e, 'Error detecting univariate drift', metadata, core_labels
|
||||
)
|
||||
raise e
|
||||
await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels)
|
||||
await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels)
|
||||
raise
|
||||
self.observe_lag_sync(start_time, metrics.MODEL_ANALYZE_LAG, core_labels)
|
||||
self.emit_metric_sync(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels)
|
||||
|
||||
core_labels = self.get_core_labels(metadata, operation_type='detect_multivariate_drift')
|
||||
start_time = time.time()
|
||||
try:
|
||||
multivariate_drift = model_analysis.detect_multivariate_drift(
|
||||
multivariate_drift = drift_analysis.detect_multivariate_drift(
|
||||
reference_df=reference_data,
|
||||
analysis_df=target_data,
|
||||
features=reference_columns,
|
||||
features=feature_names,
|
||||
timestamp_col=config['timestamp'],
|
||||
chunk_period=chunk_period,
|
||||
)
|
||||
except Exception as e:
|
||||
self.error(f'Error detecting multivariate drift: {e}', metadata)
|
||||
await self.emit_metric(
|
||||
metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels
|
||||
if isinstance(e, DriftInsufficientDataError):
|
||||
raise
|
||||
self._drift_analyze_stage_error(
|
||||
e, 'Error detecting multivariate drift', metadata, core_labels
|
||||
)
|
||||
raise e
|
||||
await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels)
|
||||
await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels)
|
||||
raise
|
||||
self.observe_lag_sync(start_time, metrics.MODEL_ANALYZE_LAG, core_labels)
|
||||
self.emit_metric_sync(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels)
|
||||
|
||||
start_time = time.time()
|
||||
core_labels = self.get_core_labels(metadata, operation_type='get_drift_metrics_dataframe')
|
||||
try:
|
||||
drift_df = model_analysis.get_drift_metrics_dataframe(
|
||||
drift_df = drift_analysis.get_drift_metrics_dataframe(
|
||||
univariate_drift=univariate_drift,
|
||||
multivariate_drift=multivariate_drift,
|
||||
)
|
||||
except Exception as e:
|
||||
self.error(f'Error getting drift metrics: {e}', metadata)
|
||||
await self.emit_metric(
|
||||
metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels
|
||||
if isinstance(e, DriftInsufficientDataError):
|
||||
raise
|
||||
self._drift_analyze_stage_error(
|
||||
e, 'Error building drift metrics dataframe', metadata, core_labels
|
||||
)
|
||||
raise e
|
||||
await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels)
|
||||
await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels)
|
||||
|
||||
self._debug_dataframe(f'Drift dataframe: Size {drift_df.shape}', drift_df, metadata)
|
||||
raise
|
||||
self.observe_lag_sync(start_time, metrics.MODEL_ANALYZE_LAG, core_labels)
|
||||
self.emit_metric_sync(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels)
|
||||
|
||||
return drift_df
|
||||
|
||||
@staticmethod
|
||||
def _to_naive_utc(series: Series) -> Series:
|
||||
"""
|
||||
Parse ``series`` as datetime and return a TZ-naive UTC copy.
|
||||
|
||||
``sientia_model.analytics.drift_analysis.DriftAnalysis`` preserves the
|
||||
timezone of the input dataframe in its outputs, while target rows
|
||||
loaded from PostgreSQL come in with ``+00:00``. Forcing both sides of
|
||||
a comparison to TZ-naive UTC keeps ``isin`` / ``floor`` operations
|
||||
deterministic regardless of how the analyzer constructs its
|
||||
timestamps.
|
||||
|
||||
Args:
|
||||
- series (Series): Input series containing datetime-parseable values.
|
||||
|
||||
Return:
|
||||
Series: Datetime64 series with ``tz=None`` representing UTC instants.
|
||||
"""
|
||||
parsed = to_datetime(series)
|
||||
if getattr(parsed.dt, 'tz', None) is not None:
|
||||
parsed = parsed.dt.tz_convert('UTC').dt.tz_localize(None)
|
||||
return parsed
|
||||
|
||||
@activity.defn(name='calculate_drift')
|
||||
async def calculate_drift(self, input_data: dict[str, Any]) -> list[dict]:
|
||||
def calculate_drift(self, input_data: dict[str, Any]) -> list[dict]:
|
||||
"""
|
||||
Calculate drift metrics for a model.
|
||||
|
||||
@@ -195,14 +259,17 @@ class ModelMetrics(SientiaMonitoring):
|
||||
|
||||
target_data = target_data.pivot(index='timestamp', columns='variable', values='value')
|
||||
target_data['timestamp'] = target_data.index
|
||||
# Keep timestamps as datetime: DriftAnalysis._chunk_dataframe relies on
|
||||
# ``pd.Grouper(freq=...)`` which rejects string timestamp columns.
|
||||
target_data['timestamp'] = to_datetime(target_data['timestamp'])
|
||||
target_data['timestamp'] = target_data['timestamp'].dt.strftime(DATETIME_FORMAT)
|
||||
target_data = target_data.reset_index(drop=True)
|
||||
target_data.dropna(inplace=True)
|
||||
|
||||
if reference_raw_data is not None:
|
||||
self.info('Using reference data', metadata)
|
||||
reference_data = DataFrame(reference_raw_data)
|
||||
if 'timestamp' in reference_data.columns:
|
||||
reference_data['timestamp'] = to_datetime(reference_data['timestamp'])
|
||||
accurate = True
|
||||
else:
|
||||
# Get 30% first rows of target_data
|
||||
@@ -211,7 +278,7 @@ class ModelMetrics(SientiaMonitoring):
|
||||
reference_data = target_data.head(int(len(target_data) * 0.3))
|
||||
accurate = False
|
||||
|
||||
await self.send_notification_async(
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='MODEL_METRICS_REFERENCE_DATA_WARNING',
|
||||
message='Using 30% first rows of target data as reference data',
|
||||
@@ -225,7 +292,7 @@ class ModelMetrics(SientiaMonitoring):
|
||||
).columns
|
||||
|
||||
try:
|
||||
drift_df = await self.get_drift_metrics(
|
||||
drift_df = self.get_drift_metrics(
|
||||
reference_data=reference_data,
|
||||
target_data=target_data,
|
||||
target_name=target_name,
|
||||
@@ -235,32 +302,30 @@ class ModelMetrics(SientiaMonitoring):
|
||||
metadata=metadata,
|
||||
)
|
||||
except Exception as e:
|
||||
self.error(f'Error getting drift metrics: {e}', metadata)
|
||||
await self.send_notification_async(
|
||||
if isinstance(e, DriftInsufficientDataError):
|
||||
self.error(str(e), metadata)
|
||||
notification_id = e.notification_id
|
||||
notification_message = str(e)
|
||||
else:
|
||||
self.error(f'Error getting drift metrics: {e}', metadata)
|
||||
notification_id = 'MODEL_METRICS_GET_DRIFT_METRICS_ERROR'
|
||||
notification_message = f'Error getting drift metrics: {e}'
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='MODEL_METRICS_GET_DRIFT_METRICS_ERROR',
|
||||
message=f'Error getting drift metrics: {e}',
|
||||
notification_id=notification_id,
|
||||
message=notification_message,
|
||||
block='model_metrics',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=traceback.format_exc(),
|
||||
)
|
||||
return []
|
||||
raise
|
||||
|
||||
if drift_df.empty:
|
||||
self.warning('No drift metrics found', metadata)
|
||||
return []
|
||||
|
||||
# Drop unnecessary columns
|
||||
drift_df.drop(columns=['p_value'], inplace=True)
|
||||
|
||||
# Extract timestamps only until minutes
|
||||
if chunk_period == 'min':
|
||||
target_timestamps = target_data['timestamp'].apply(lambda x: x[:16])
|
||||
else:
|
||||
target_timestamps = target_data['timestamp']
|
||||
|
||||
# Drop rows where timestamp is not in target data, to avoid save drift from reference
|
||||
drift_df = drift_df[drift_df['timestamp'].isin(target_timestamps)]
|
||||
# Drop chunks whose floored timestamp does not appear in the analysis window.
|
||||
# ``DriftAnalysis`` chunks over ``analysis_df``; this only excludes rows that
|
||||
# do not belong to the current target window (e.g. stray merged reference rows).
|
||||
target_floor = self._to_naive_utc(target_data['timestamp']).dt.floor(chunk_period)
|
||||
drift_floor = self._to_naive_utc(drift_df['timestamp']).dt.floor(chunk_period)
|
||||
drift_df = drift_df[drift_floor.isin(target_floor)]
|
||||
|
||||
if drift_df.empty:
|
||||
self.warning(
|
||||
@@ -269,33 +334,36 @@ class ModelMetrics(SientiaMonitoring):
|
||||
)
|
||||
return []
|
||||
|
||||
# Rename columns to match database columns
|
||||
drift_df.rename(
|
||||
columns={
|
||||
'metric': 'method',
|
||||
'statistic': 'value',
|
||||
},
|
||||
inplace=True,
|
||||
)
|
||||
# Analyzer emits diagnostic columns that are not stored in ``sientia_data.drift_metrics``.
|
||||
drift_df = drift_df.drop(columns=['threshold', 'drift_type'], errors='ignore')
|
||||
|
||||
# Drop duplicates
|
||||
drift_df.drop_duplicates(
|
||||
subset=['timestamp', 'method', 'feature'], keep='first', inplace=True
|
||||
)
|
||||
|
||||
drift_df['model_id'] = model_id
|
||||
drift_df['model_id'] = str(model_id)
|
||||
drift_df['accurate'] = accurate
|
||||
|
||||
drift_df['timestamp'] = to_datetime(drift_df['timestamp'])
|
||||
drift_df['timestamp'] = drift_df['timestamp'].dt.tz_localize('UTC')
|
||||
drift_df['timestamp'] = drift_df['timestamp'].dt.strftime(DATETIME_FORMAT_WITH_TZ)
|
||||
# ``timestamp`` is overridden with the most recent target instant so
|
||||
# every persisted row shares a single business timestamp (the run's
|
||||
# logical "now"), matching what downstream consumers expect.
|
||||
latest_target_timestamp = self._to_naive_utc(target_data['timestamp']).max()
|
||||
drift_df['timestamp'] = (
|
||||
pd.Timestamp(latest_target_timestamp)
|
||||
.tz_localize('UTC')
|
||||
.strftime(DATETIME_FORMAT_WITH_TZ)
|
||||
)
|
||||
|
||||
# ``chunk_start_date`` / ``chunk_end_date`` may carry nanosecond
|
||||
# precision (beyond ``timestamptz`` microseconds), so serialize as ISO
|
||||
# text for the ``text`` Postgres columns.
|
||||
for column in ('chunk_start_date', 'chunk_end_date'):
|
||||
drift_df[column] = drift_df[column].apply(
|
||||
lambda value: pd.Timestamp(value).isoformat() if pd.notna(value) else None
|
||||
)
|
||||
|
||||
self._debug_dataframe(f'Drift dataframe: Size {drift_df.shape}', drift_df, metadata)
|
||||
|
||||
return drift_df.to_dict(orient='records')
|
||||
|
||||
@activity.defn(name='calculate_simple_metrics')
|
||||
async def calculate_simple_metrics(self, input_data: dict[str, Any]) -> list[dict]:
|
||||
def calculate_simple_metrics(self, input_data: dict[str, Any]) -> list[dict]:
|
||||
"""
|
||||
Calculate simple metrics for a model. Metrics available are:
|
||||
- rmse
|
||||
@@ -319,7 +387,7 @@ class ModelMetrics(SientiaMonitoring):
|
||||
metadata = input_data['metadata']
|
||||
model_id = input_data['model_id']
|
||||
target_data = DataFrame(input_data['target_data'])
|
||||
metrics = input_data['metrics']
|
||||
metric_names = input_data['metrics']
|
||||
interval_minutes = input_data['interval_minutes']
|
||||
|
||||
data_size = target_data.shape[0]
|
||||
@@ -329,9 +397,9 @@ class ModelMetrics(SientiaMonitoring):
|
||||
diff = target_data['target'] - target_data['prediction']
|
||||
diff_squared = diff**2
|
||||
|
||||
self.info(f'Calculating simple metrics for model {model_id}: {metrics}', metadata)
|
||||
self.info(f'Calculating simple metrics for model {model_id}: {metric_names}', metadata)
|
||||
|
||||
for metric in metrics:
|
||||
for metric in metric_names:
|
||||
if metric == 'rmse':
|
||||
output_data.append({'metric': 'rmse', 'value': np.sqrt(np.mean(diff_squared))})
|
||||
elif metric == 'mse':
|
||||
|
||||
@@ -82,15 +82,13 @@ class OPC(SientiaMonitoring):
|
||||
notification_handler: NotificationHandler,
|
||||
metrics_controller: MetricsController,
|
||||
):
|
||||
self.logger = logger
|
||||
self.notification_handler = notification_handler
|
||||
self.opc_servers = opc_servers
|
||||
|
||||
SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
|
||||
|
||||
self.opc_repository: dict[str, OpcRepository] = {}
|
||||
|
||||
async def init_opc(self):
|
||||
def init_opc(self):
|
||||
"""
|
||||
Initialize OPC server connections and establish communication channels.
|
||||
|
||||
@@ -106,26 +104,18 @@ class OPC(SientiaMonitoring):
|
||||
4. Logs connection status for operational visibility
|
||||
|
||||
Raises:
|
||||
ValueError: A server whose `cert_path` is set but whose `server_uri`
|
||||
(or private key) is missing fails `OpcRepository.set_security()`
|
||||
with no try/except between here and `main()`. This is a
|
||||
configuration error, not a connectivity one, and is intentionally
|
||||
fatal to the whole worker process — fix the deployment config
|
||||
rather than catch it here.
|
||||
Exception: If OPC repository initialization fails or connection
|
||||
establishment encounters critical errors
|
||||
|
||||
Note:
|
||||
Connection failures (server unreachable, timeout, etc.) are logged
|
||||
and reported but do not prevent the initialization of other OPC
|
||||
servers — each server's `connect()` call catches those on its own.
|
||||
This isolation does NOT extend to configuration errors raised before
|
||||
a connection is attempted (see `Raises` above): a single
|
||||
misconfigured server intentionally stops every other server, and the
|
||||
whole worker, from starting.
|
||||
Connection failures are logged and reported but do not prevent
|
||||
the initialization of other OPC servers. Each server is handled
|
||||
independently to ensure maximum availability.
|
||||
"""
|
||||
self.info('Initializing OPC servers...')
|
||||
for opc_id, server in self.opc_servers.items():
|
||||
self.opc_repository[opc_id] = OpcRepository(
|
||||
opc_id=server['id'],
|
||||
opc_id=opc_id,
|
||||
server_name=server['server_name'],
|
||||
url=server['url'],
|
||||
logger=self.logger,
|
||||
@@ -134,12 +124,12 @@ class OPC(SientiaMonitoring):
|
||||
private_key_path=server['private_key_path'],
|
||||
server_cert_path=server['server_cert_path'],
|
||||
notification_handler=self.notification_handler,
|
||||
reconnection_interval=server['reconnection_interval'],
|
||||
reconnection_interval=server.get('reconnection_interval', 60),
|
||||
metrics_controller=self.metrics_controller,
|
||||
)
|
||||
is_connected, error_data = await self.opc_repository[opc_id].connect()
|
||||
is_connected, error_data = self.opc_repository[opc_id].connect()
|
||||
if not is_connected:
|
||||
await self.send_notification_async(
|
||||
self.send_notification(
|
||||
metadata={
|
||||
'model_id': '-',
|
||||
'model_name': '-',
|
||||
@@ -155,7 +145,7 @@ class OPC(SientiaMonitoring):
|
||||
else:
|
||||
self.info(f'OPC server {opc_id}:{server["server_name"]} connected successfully.')
|
||||
|
||||
async def write_data(
|
||||
def write_data(
|
||||
self,
|
||||
server_id: str,
|
||||
tag: str,
|
||||
@@ -173,11 +163,11 @@ class OPC(SientiaMonitoring):
|
||||
"""
|
||||
|
||||
try:
|
||||
is_success, info_data = await self.opc_repository[server_id].write_data(
|
||||
is_success, info_data = self.opc_repository[server_id].write_data(
|
||||
tag, data, data_type, metadata
|
||||
)
|
||||
if not is_success:
|
||||
await self.send_notification_async(
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id=info_data['notification_id'],
|
||||
message=info_data['message'],
|
||||
@@ -189,7 +179,7 @@ class OPC(SientiaMonitoring):
|
||||
return info_data['response_time'], None
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
await self.send_notification_async(
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id=f'WRITE_OPC_{tag_type.upper()}_ERROR',
|
||||
message=f'Error writing data to OPC server: {e}',
|
||||
@@ -197,9 +187,9 @@ class OPC(SientiaMonitoring):
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace,
|
||||
)
|
||||
raise e
|
||||
raise
|
||||
|
||||
async def validate_server(self, server_id: str, metadata: dict[str, Any]) -> bool:
|
||||
def validate_server(self, server_id: str, metadata: dict[str, Any]) -> bool:
|
||||
"""
|
||||
Validate that an OPC server is available and configured for write operations.
|
||||
|
||||
@@ -222,7 +212,7 @@ class OPC(SientiaMonitoring):
|
||||
"""
|
||||
if self.opc_repository.get(server_id) is None:
|
||||
message = f'OPC server {server_id} not found to perform write operation.'
|
||||
await self.send_notification_async(
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='OPC_SERVER_NOT_FOUND',
|
||||
message=message,
|
||||
@@ -233,7 +223,7 @@ class OPC(SientiaMonitoring):
|
||||
return False
|
||||
return True
|
||||
|
||||
async def _write_tags_from_config(
|
||||
def _write_tags_from_config(
|
||||
self,
|
||||
server_id: str,
|
||||
tags_config: dict[str, dict[str, Any]],
|
||||
@@ -264,7 +254,7 @@ class OPC(SientiaMonitoring):
|
||||
reconnect_in_progress_seen = False
|
||||
|
||||
for tag, tag_config in tags_config.items():
|
||||
response_time, error_info = await self.write_data(
|
||||
response_time, error_info = self.write_data(
|
||||
server_id=server_id,
|
||||
tag=tag,
|
||||
data=data.head(1)[data_column].values[0],
|
||||
@@ -289,7 +279,7 @@ class OPC(SientiaMonitoring):
|
||||
|
||||
return response_times, session_bad_seen, session_bad_status, reconnect_in_progress_seen
|
||||
|
||||
async def manage_output_tags(
|
||||
def manage_output_tags(
|
||||
self,
|
||||
server_id: str,
|
||||
config: dict[str, Any],
|
||||
@@ -339,7 +329,7 @@ class OPC(SientiaMonitoring):
|
||||
group_session_bad,
|
||||
group_status,
|
||||
group_reconnect,
|
||||
) = await self._write_tags_from_config(
|
||||
) = self._write_tags_from_config(
|
||||
server_id=server_id,
|
||||
tags_config=config[config_key],
|
||||
data=data,
|
||||
@@ -365,7 +355,7 @@ class OPC(SientiaMonitoring):
|
||||
)
|
||||
|
||||
@activity.defn(name='write_opc_data')
|
||||
async def write_opc_data(
|
||||
def write_opc_data(
|
||||
self, input_data: dict[str, Any]
|
||||
) -> tuple[dict[Hashable, Any], dict[str, dict[str, float | None]]]:
|
||||
"""
|
||||
@@ -396,10 +386,10 @@ class OPC(SientiaMonitoring):
|
||||
session_bad_status: str | None = None
|
||||
reconnect_in_progress_seen = False
|
||||
|
||||
metrics: dict[str, dict[str, float | None]] = {}
|
||||
opc_metrics: dict[str, dict[str, float | None]] = {}
|
||||
|
||||
for server_id, config in opc_output_config.items():
|
||||
if not await self.validate_server(server_id, metadata):
|
||||
if not self.validate_server(server_id, metadata):
|
||||
success = False
|
||||
continue
|
||||
|
||||
@@ -409,8 +399,8 @@ class OPC(SientiaMonitoring):
|
||||
local_session_bad,
|
||||
local_status,
|
||||
local_reconnect_in_progress,
|
||||
) = await self.manage_output_tags(server_id, config, data, metadata)
|
||||
metrics[server_id] = local_response_times
|
||||
) = self.manage_output_tags(server_id, config, data, metadata)
|
||||
opc_metrics[server_id] = local_response_times
|
||||
local_count = len(local_response_times)
|
||||
success = success and local_success
|
||||
if local_session_bad:
|
||||
@@ -419,8 +409,10 @@ class OPC(SientiaMonitoring):
|
||||
if local_reconnect_in_progress:
|
||||
reconnect_in_progress_seen = True
|
||||
|
||||
n_pred = len(config.get('prediction_tags') or {})
|
||||
n_conf = len(config.get('confidence_tags') or {})
|
||||
self.info(
|
||||
f'Process completed for OPC server {server_id}: {local_count} of {len(config.get("prediction_tags", []))} prediction tags and {len(config.get("confidence_tags", []))} confidence tags',
|
||||
f'Process completed for OPC server {server_id}: {local_count} of {n_pred} prediction tags and {n_conf} confidence tags',
|
||||
metadata,
|
||||
)
|
||||
|
||||
@@ -433,7 +425,7 @@ class OPC(SientiaMonitoring):
|
||||
opc_status=session_bad_status,
|
||||
reconnect_in_progress=reconnect_in_progress_seen,
|
||||
),
|
||||
metrics,
|
||||
opc_metrics,
|
||||
)
|
||||
|
||||
def process_confidence(
|
||||
@@ -497,7 +489,7 @@ class OPC(SientiaMonitoring):
|
||||
|
||||
return data.to_dict()
|
||||
|
||||
async def aclose(self):
|
||||
def close(self):
|
||||
"""
|
||||
Gracefully shutdown all OPC server connections and cleanup resources.
|
||||
|
||||
@@ -518,4 +510,5 @@ class OPC(SientiaMonitoring):
|
||||
their current state and provides a clean shutdown experience.
|
||||
"""
|
||||
for opc in self.opc_repository.values():
|
||||
await opc.disconnect()
|
||||
opc.disconnect()
|
||||
self.opc_repository.clear()
|
||||
|
||||
132
laborious/activities/prediction_metadata.py
Normal file
132
laborious/activities/prediction_metadata.py
Normal file
@@ -0,0 +1,132 @@
|
||||
from temporalio import activity, workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
from typing import Any
|
||||
|
||||
from bson import ObjectId
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||
from sientia_do.observability.logger import Logger
|
||||
from sientia_do.observability.metrics_controller import MetricsController
|
||||
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
|
||||
from sientia_do.repository.mongodb_repository_sync import MongoDBRepository
|
||||
|
||||
|
||||
class PredictionMetadata(SientiaMonitoring):
|
||||
"""
|
||||
Centralized, best-effort writer for prediction gate metadata (triggered filter, policy,
|
||||
counter) to MongoDB.
|
||||
|
||||
Gates only compute and return a metadata document; they never write it themselves. Keeping
|
||||
the write here, instead of in each gate, puts the one real write point in the subworkflow
|
||||
that's about to persist a `predictions` row — the only place that reliably knows whether one
|
||||
is being written at all — and lets other metadata sources reuse the same activity instead of
|
||||
duplicating write logic wherever they originate.
|
||||
"""
|
||||
|
||||
mongo_repository: MongoDBRepository | None = None
|
||||
_PREDICTION_METADATA_COLLECTION = 'prediction_metadata'
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
mongo_repository: MongoDBRepository | None = None,
|
||||
logger: Logger | None = None,
|
||||
notification_handler: NotificationHandler | None = None,
|
||||
metrics_controller: MetricsController | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize the prediction metadata writer.
|
||||
|
||||
Args:
|
||||
mongo_repository: MongoDB, where gate metadata (triggered filter, policy, counter) is
|
||||
written best-effort. `None` is a valid state (see `write_prediction_metadata`).
|
||||
|
||||
Raises:
|
||||
Exception: If BaseActivity initialization fails
|
||||
"""
|
||||
self.mongo_repository = mongo_repository
|
||||
SientiaMonitoring.__init__(
|
||||
self,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
def close(self) -> None:
|
||||
"""
|
||||
Close the prediction metadata writer and clean up resources.
|
||||
"""
|
||||
if self.mongo_repository is not None:
|
||||
try:
|
||||
self.mongo_repository.close()
|
||||
finally:
|
||||
self.mongo_repository = None
|
||||
SientiaMonitoring.shutdown(self)
|
||||
|
||||
def __del__(self):
|
||||
self.close()
|
||||
|
||||
@activity.defn(name='write_prediction_metadata')
|
||||
def write_prediction_metadata(self, input_data: dict[str, Any]) -> str | None:
|
||||
"""
|
||||
Best-effort write of one or more gate metadata documents to Mongo, combined into a single
|
||||
`prediction_metadata` document.
|
||||
|
||||
Implements both business rules for this feature in one place:
|
||||
- The prediction write must never be blocked or failed by a metadata write. A failed
|
||||
`insert(...)` is swallowed, never re-raised, logged here with call-site context (which
|
||||
`mongo_repository.insert(...)` itself can't know) — matching how `storage.py`'s
|
||||
`cleanup_minio_objects_expired` also logs again at the call site even though the
|
||||
repository already logs internally and emits its own Prometheus metric
|
||||
(`MONGODB_WRITE_ERROR_COUNT`); the generic library log is for the metric/trace, this
|
||||
one is for whoever is reading logs for this specific call.
|
||||
- Nothing is written when there's nothing relevant to report. An empty/missing
|
||||
`documents` list (or one full of falsy/empty entries) skips Mongo entirely, so callers
|
||||
can call this unconditionally and get `None` back for "no metadata this time" instead
|
||||
of deciding upfront whether there's something worth saving.
|
||||
|
||||
Takes a *list* of documents, not a single one: each caller contributes its own document
|
||||
(shaped like `{'filters': [...]}`), and this flattens every `filters` entry across all of
|
||||
them into one `prediction_metadata` document before the single insert. Today only one
|
||||
gate ever contributes a document per prediction cycle (the gate chain short-circuits on
|
||||
the first non-`None` `path_flag`), but the list shape means this activity won't need to
|
||||
change again once more than one source contributes metadata for the same cycle.
|
||||
|
||||
`insert(...)` doesn't hand back the inserted document's `_id`, so it is generated here,
|
||||
client-side, before the insert, as a plain string (not a BSON ObjectId) — matching this
|
||||
codebase's convention of string `_id`s (e.g. the `models` collection) and the 24-hex-char
|
||||
size of the `predictions.metadata_id` Postgres column that will reference it.
|
||||
|
||||
Args:
|
||||
input_data: Dict with:
|
||||
- documents (list[dict[str, Any]]): Metadata payloads to combine, each shaped
|
||||
like `{'filters': [...]}`. Missing/empty, or entries with no `filters`, skip
|
||||
the write entirely.
|
||||
- metadata (dict): Workflow metadata for contextual logging.
|
||||
|
||||
Return:
|
||||
The generated id, or `None` when there's nothing to write, Mongo isn't configured, or
|
||||
the insert failed.
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
documents = input_data.get('documents') or []
|
||||
|
||||
filters = [
|
||||
entry for document in documents if document for entry in document.get('filters', [])
|
||||
]
|
||||
|
||||
if not filters or self.mongo_repository is None:
|
||||
return None
|
||||
|
||||
generated_id = str(ObjectId())
|
||||
|
||||
try:
|
||||
self.mongo_repository.insert(
|
||||
self._PREDICTION_METADATA_COLLECTION,
|
||||
{'filters': filters, '_id': generated_id},
|
||||
metadata,
|
||||
)
|
||||
except Exception as e:
|
||||
self.error(f'Failed to write prediction metadata: {e}', metadata)
|
||||
return None
|
||||
|
||||
return generated_id
|
||||
@@ -1,9 +1,8 @@
|
||||
from temporalio import activity, workflow
|
||||
|
||||
from laborious.utils.repository.minio_manager import MinioManager
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
# Extend the Temporal Postgres activities for convenient query -> MinIO export
|
||||
import time
|
||||
import traceback
|
||||
from datetime import timedelta
|
||||
from typing import Any
|
||||
@@ -11,11 +10,14 @@ with workflow.unsafe.imports_passed_through():
|
||||
import pandas as pd
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
from sientia_do.observability import metrics
|
||||
from sientia_do.observability.logger import Logger
|
||||
from sientia_do.observability.metrics_controller import MetricsController
|
||||
from sientia_do.repository.minio_repository import MinioRepository
|
||||
from sientia_do.temporal.activities.postgres import Postgres
|
||||
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
|
||||
from sientia_do.repository.minio_repository_sync import MinioRepository
|
||||
from sientia_do.temporal.activities.postgres_sync import Postgres
|
||||
from sientia_do.temporal.constants import now
|
||||
from sqlalchemy import text
|
||||
|
||||
from laborious.utils.messages import MISSING_MINIO_REPOSITORY
|
||||
from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
|
||||
@@ -24,10 +26,12 @@ with workflow.unsafe.imports_passed_through():
|
||||
_LOAD_QUERY_OFFLOAD_SKIP_KEYS = frozenset({'model_name', 'key_prefix', 'size_threshold_bytes'})
|
||||
|
||||
|
||||
class Storage(Postgres, MinioManager):
|
||||
class Storage(Postgres, SientiaMonitoring):
|
||||
"""
|
||||
Extensions for Postgres activities with a helper to export query results
|
||||
directly to MinIO as Parquet and return the object name.
|
||||
directly to MinIO as Parquet and return the object name, and with
|
||||
``repeat_last_prediction``, a query specific to this repo's REPEAT path
|
||||
(moved here from the shared library — nothing else consumes it).
|
||||
"""
|
||||
|
||||
minio_repository: MinioRepository | None = None
|
||||
@@ -62,14 +66,16 @@ class Storage(Postgres, MinioManager):
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
MinioManager.__init__(
|
||||
self, minio_repository, logger, notification_handler, metrics_controller
|
||||
self.minio_repository = minio_repository
|
||||
SientiaMonitoring.__init__(
|
||||
self,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
@activity.defn(name='load_query_with_minio_offload')
|
||||
async def load_query_with_minio_offload(
|
||||
self, input_data: dict[str, Any]
|
||||
) -> MinioDataFramePayload:
|
||||
def load_query_with_minio_offload(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
|
||||
"""
|
||||
Run the custom SQL load, then return a MinIO-aware dataframe wire dict.
|
||||
|
||||
@@ -90,7 +96,7 @@ class Storage(Postgres, MinioManager):
|
||||
metadata: dict = input_data.get('metadata', {})
|
||||
model_name = input_data['model_name']
|
||||
|
||||
rows = await self.load_custom_query(
|
||||
rows = self.load_custom_query(
|
||||
input_data,
|
||||
)
|
||||
if not rows:
|
||||
@@ -101,7 +107,7 @@ class Storage(Postgres, MinioManager):
|
||||
else:
|
||||
dataframe = pd.DataFrame(rows)
|
||||
|
||||
return await MinioDataFramePayload.from_dataframe(
|
||||
return MinioDataFramePayload.from_dataframe(
|
||||
dataframe,
|
||||
minio_repo=self.minio_repository,
|
||||
workflow_metadata=metadata,
|
||||
@@ -111,31 +117,144 @@ class Storage(Postgres, MinioManager):
|
||||
)
|
||||
|
||||
@activity.defn(name='export_payload_to_postgres')
|
||||
async def export_payload_to_postgres(self, input_data: dict[str, Any]) -> dict:
|
||||
def export_payload_to_postgres(self, input_data: dict[str, Any]) -> dict:
|
||||
"""
|
||||
Export a payload to PostgreSQL.
|
||||
Resolve a MinIO-aware payload into a DataFrame and persist it into PostgreSQL.
|
||||
|
||||
This activity accepts the serialized payload produced by previous steps
|
||||
(inline dict or MinIO object reference), reconstructs the tabular data,
|
||||
and delegates the final write to ``export_data_to_postgres`` using the
|
||||
same input contract expected by the Postgres activity mixin.
|
||||
|
||||
Args:
|
||||
- input_data (dict[str, Any]): Activity input containing ``data`` as a
|
||||
``MinioDataFramePayload``-compatible dict plus database write options
|
||||
(schema/table/on_conflict/metadata and related fields).
|
||||
|
||||
Return:
|
||||
dict: Result dictionary returned by ``export_data_to_postgres``, including
|
||||
success status and optional write diagnostics.
|
||||
"""
|
||||
metadata = input_data.get('metadata')
|
||||
payload = MinioDataFramePayload.from_dict(input_data['data'])
|
||||
data = await payload.retrieve(self.minio_repository, metadata)
|
||||
data = payload.retrieve(self.minio_repository, metadata)
|
||||
|
||||
return await self.export_data_to_postgres(
|
||||
return self.export_data_to_postgres(
|
||||
{
|
||||
**input_data,
|
||||
'data': data,
|
||||
}
|
||||
)
|
||||
|
||||
def _build_repeat_last_prediction_query(self, schema: str, table_name: str):
|
||||
table_identifier = f'"{schema}".{table_name}'
|
||||
return text( # noqa: S608
|
||||
f"""
|
||||
INSERT INTO {table_identifier} (
|
||||
model_id,
|
||||
prediction,
|
||||
timestamp,
|
||||
response_time,
|
||||
prediction_status,
|
||||
prediction_confidence,
|
||||
metadata_id,
|
||||
created_at
|
||||
)
|
||||
SELECT
|
||||
model_id,
|
||||
prediction,
|
||||
:last_timestamp,
|
||||
response_time,
|
||||
prediction_status,
|
||||
prediction_confidence,
|
||||
:metadata_id,
|
||||
NOW()
|
||||
FROM {table_identifier}
|
||||
WHERE model_id = :model
|
||||
ORDER BY timestamp DESC
|
||||
LIMIT 1;
|
||||
""" # nosec B608
|
||||
)
|
||||
|
||||
@activity.defn(name='repeat_last_prediction')
|
||||
def repeat_last_prediction(self, query_items: dict[str, Any]):
|
||||
"""
|
||||
Insert a new `predictions` row that copies the last one for a model forward in time.
|
||||
|
||||
REPEAT persists a new row (unlike STOP, which persists nothing), reusing the previous
|
||||
row's `prediction`/`response_time`/`prediction_status`/`prediction_confidence` values with
|
||||
a fresh `timestamp`/`created_at`. `metadata_id` references the `prediction_metadata`
|
||||
document for this cycle, if any (see `PredictionMetadata.write_prediction_metadata`);
|
||||
`None` is a valid value — the column is nullable.
|
||||
|
||||
Args (query_items):
|
||||
metadata (dict): Workflow metadata for logging and metrics.
|
||||
schema (str): Postgres schema of the `predictions` table.
|
||||
table_name (str): `predictions` table name.
|
||||
model (str): `model_id` to repeat the last prediction for.
|
||||
last_timestamp: Timestamp to stamp the new row with.
|
||||
metadata_id (str | None): `prediction_metadata` document id, or `None`.
|
||||
|
||||
Raises:
|
||||
Exception: Re-raised after logging/metrics if the insert fails.
|
||||
"""
|
||||
metadata = query_items['metadata']
|
||||
schema = query_items['schema']
|
||||
table_name = query_items['table_name']
|
||||
model = query_items['model']
|
||||
last_timestamp = query_items['last_timestamp']
|
||||
metadata_id = query_items.get('metadata_id')
|
||||
|
||||
repeat_query = self._build_repeat_last_prediction_query(schema, table_name)
|
||||
self.info(f'Repeating last prediction for model {model}', metadata)
|
||||
self.debug(
|
||||
f'Query parameters: schema={schema}, model={model}, last_timestamp={last_timestamp}',
|
||||
metadata,
|
||||
)
|
||||
core_labels = {
|
||||
**self.get_core_labels(metadata, 'execute_sql'),
|
||||
'conflict_mode': 'error',
|
||||
}
|
||||
|
||||
with self.session_factory.begin() as session:
|
||||
try:
|
||||
start_time = time.time()
|
||||
session.execute(
|
||||
repeat_query,
|
||||
{
|
||||
'last_timestamp': last_timestamp,
|
||||
'model': model,
|
||||
'metadata_id': metadata_id,
|
||||
},
|
||||
)
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='ERROR_REPEATING_LAST_PREDICTION',
|
||||
message=f'Error repeating last prediction: {e}',
|
||||
block='repeat_last_prediction',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace,
|
||||
)
|
||||
self.error(trace, metadata)
|
||||
self.emit_metric_sync(
|
||||
metric_object=metrics.POSTGRES_WRITE_ERROR_COUNT,
|
||||
tags=core_labels,
|
||||
)
|
||||
raise
|
||||
else:
|
||||
self.observe_lag_sync(start_time, metrics.POSTGRES_WRITE_LAG, core_labels)
|
||||
self.emit_metric_sync(
|
||||
metric_object=metrics.POSTGRES_WRITE_COUNT,
|
||||
tags=core_labels,
|
||||
)
|
||||
|
||||
@activity.defn(name='cleanup_minio_objects_expired')
|
||||
async def cleanup_minio_objects_expired(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
def cleanup_minio_objects_expired(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Delete objects under the given prefixes that are older than the retention window.
|
||||
|
||||
See also: `ModelImport.ensure_import_bundle_retention`
|
||||
(laborious/activities/model_import.py) — a newer, unrelated lifecycle-policy expiry
|
||||
mechanism for the import bundle prefix. The two are independent by design for now; see
|
||||
docs/minio-object-expiry.md.
|
||||
|
||||
Args (input_data):
|
||||
metadata (dict): Workflow metadata for logging and metrics.
|
||||
prefixes (list[str]): Key prefixes to scan (one level or subtree per prefix).
|
||||
@@ -159,7 +278,7 @@ class Storage(Postgres, MinioManager):
|
||||
'deleted_count': 0,
|
||||
}
|
||||
try:
|
||||
keys = await self.minio_repository.list_objects(
|
||||
keys = self.minio_repository.list_objects(
|
||||
prefix=prefix,
|
||||
recursive=True,
|
||||
metadata=metadata,
|
||||
@@ -171,7 +290,7 @@ class Storage(Postgres, MinioManager):
|
||||
continue
|
||||
if ts >= cutoff:
|
||||
continue
|
||||
await self.minio_repository.delete_file(
|
||||
self.minio_repository.delete_file(
|
||||
object_name=key,
|
||||
metadata=metadata,
|
||||
)
|
||||
@@ -189,7 +308,7 @@ class Storage(Postgres, MinioManager):
|
||||
report['deleted_count'] += 1
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
await self.send_notification_async(
|
||||
self.send_notification(
|
||||
metadata=metadata,
|
||||
notification_id='ERROR_CLEANUP_MINIO_OBJECTS_EXPIRED',
|
||||
message=f'Error cleaning up MinIO objects: {e}',
|
||||
@@ -206,10 +325,17 @@ class Storage(Postgres, MinioManager):
|
||||
return report
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close Storage resources (MinIO client and Postgres engine)."""
|
||||
"""
|
||||
Shutdown Storage resources in deterministic order.
|
||||
|
||||
The method first closes Postgres resources via ``Postgres.close`` (engine,
|
||||
sessions, and monitoring hooks), then closes the optional MinIO repository
|
||||
and clears the local reference to avoid accidental reuse after shutdown.
|
||||
"""
|
||||
if hasattr(self, 'engine'):
|
||||
Postgres.close(self)
|
||||
MinioManager.close(self)
|
||||
|
||||
def __del__(self):
|
||||
self.close()
|
||||
if self.minio_repository is not None:
|
||||
try:
|
||||
self.minio_repository.close()
|
||||
finally:
|
||||
self.minio_repository = None
|
||||
|
||||
@@ -183,45 +183,6 @@ MODEL_ANALYZE_ERROR_COUNT = Counter(
|
||||
)
|
||||
|
||||
|
||||
# ================== Model import metrics ==================
|
||||
#
|
||||
# The failing step and the error code travel in `operation_type`, the one label the shared
|
||||
# `CORE_LABELS` leaves free. Nothing else about a failure is a label: an error reason is a sentence
|
||||
# and an exception is a log line, neither of which belongs in a metric's cardinality.
|
||||
|
||||
MODEL_IMPORT_STARTED_COUNT = Counter(
|
||||
'laborious_model_import_started_count',
|
||||
'Number of model imports that claimed their import log row',
|
||||
CORE_LABELS,
|
||||
)
|
||||
|
||||
MODEL_IMPORT_COMPLETED_COUNT = Counter(
|
||||
'laborious_model_import_completed_count',
|
||||
'Number of model imports that reached a registered version and a model document',
|
||||
CORE_LABELS,
|
||||
)
|
||||
|
||||
MODEL_IMPORT_ERROR_COUNT = Counter(
|
||||
'laborious_model_import_error_count',
|
||||
'Number of failed model imports, with the failing step or error code in operation_type',
|
||||
CORE_LABELS,
|
||||
)
|
||||
|
||||
MODEL_IMPORT_KDF_DURATION = Histogram(
|
||||
'laborious_model_import_kdf_duration_seconds',
|
||||
'Time spent deriving the bundle key with Argon2id',
|
||||
CORE_LABELS,
|
||||
buckets=[0.5, 1.0, 2.0, 3.0, 5.0, 8.0, 13.0, 21.0],
|
||||
)
|
||||
|
||||
MODEL_IMPORT_STEP_DURATION = Histogram(
|
||||
'laborious_model_import_step_duration_seconds',
|
||||
'Duration of each model import step, named in operation_type',
|
||||
CORE_LABELS,
|
||||
buckets=[0.1, 0.5, 1.0, 5.0, 10.0, 30.0, 60.0, 300.0],
|
||||
)
|
||||
|
||||
|
||||
# ================== PI Web API metrics ==================
|
||||
|
||||
PI_WEB_API_LABELS = [*CORE_LABELS, 'tag_name']
|
||||
|
||||
@@ -1,27 +0,0 @@
|
||||
"""Reading `.sientia` bundles: the wire format, the seven gates and the step vocabulary.
|
||||
|
||||
There is no writer here. The producer is the Streamlit app in `sientia-projects-templates`
|
||||
(`app/src/operations/model_export/`), and a second production writer for a format with one producer
|
||||
would be a fork waiting to happen. What stands in for it in the tests is
|
||||
`tests/helpers/bundle_factory.py` plus the golden fixtures under `tests/fixtures/bundle/`.
|
||||
"""
|
||||
|
||||
from laborious.utils.bundle.format import BundleFormatError, BundleHeader
|
||||
from laborious.utils.bundle.reader import (
|
||||
BundleGateError,
|
||||
BundleLimits,
|
||||
BundleReader,
|
||||
OpenedBundle,
|
||||
)
|
||||
from laborious.utils.bundle.steps import GATE_TO_STEP, ImportStep
|
||||
|
||||
__all__ = [
|
||||
'GATE_TO_STEP',
|
||||
'BundleFormatError',
|
||||
'BundleGateError',
|
||||
'BundleHeader',
|
||||
'BundleLimits',
|
||||
'BundleReader',
|
||||
'ImportStep',
|
||||
'OpenedBundle',
|
||||
]
|
||||
@@ -1,216 +0,0 @@
|
||||
"""The `.sientia` wire format, as the producer writes it.
|
||||
|
||||
Every constant here is copied from the producer that emits the bundles this reader opens — the
|
||||
Streamlit app in `sientia-projects-templates`, modules
|
||||
`app/src/operations/model_export/crypto.py` and `app/src/operations/model_export/bundle.py`. The
|
||||
producer file each value came from is named next to it. The golden fixtures under
|
||||
`tests/fixtures/bundle/` are the anti-drift device: if the producer's layout moves, the tests here
|
||||
fail against the committed bytes instead of failing in production.
|
||||
|
||||
This module deliberately imports nothing from Laborious, Temporal, MLflow or libsodium: it is pure
|
||||
`struct` and `hashlib` over bytes, so it stays liftable into a shared package if the format ever
|
||||
gets one (the cancelled QTZPOC-15). `reader.py` is where the crypto lives.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import struct
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
# --- Producer: app/src/operations/model_export/crypto.py ---
|
||||
MAGIC: Final[bytes] = b'SIENTIA1'
|
||||
FORMAT_VERSION: Final[int] = 1
|
||||
AEAD_ID_XCHACHA20POLY1305_SECRETSTREAM: Final[int] = 1
|
||||
KDF_ID_ARGON2ID: Final[int] = 1
|
||||
SALT_SIZE: Final[int] = 16
|
||||
DIGEST_SIZE: Final[int] = 32
|
||||
RESERVED_SIGNATURE_SIZE: Final[int] = 64
|
||||
# `KDF_MEMLIMIT_BYTES` / `KDF_OPSLIMIT` on the producer: what it writes is also the ceiling the
|
||||
# reader refuses to exceed, so header bytes cannot buy an arbitrary allocation on the worker.
|
||||
MAX_KDF_MEMLIMIT: Final[int] = 256 * 1024 * 1024 # 268435456
|
||||
MAX_KDF_OPSLIMIT: Final[int] = 3
|
||||
# `KDF_PARALLELISM_RESERVED`: libsodium's `crypto_pwhash` has no lane parameter, so the field is
|
||||
# reserved and must be exactly 1.
|
||||
PARALLELISM: Final[int] = 1
|
||||
CHUNK_SIZE: Final[int] = 64 * 1024 # 65536, fixed by the format
|
||||
KEY_SIZE: Final[int] = 32 # crypto_secretstream_xchacha20poly1305_KEYBYTES
|
||||
|
||||
# magic(8s) version(B) aead_id(B) kdf_id(B) salt(16s) memlimit(Q) opslimit(Q)
|
||||
# parallelism(B) digest(32s) reserved_signature(64s) — producer's `_HEADER_STRUCT`.
|
||||
HEADER_STRUCT: Final[struct.Struct] = struct.Struct('>8sBBB16sQQB32s64s')
|
||||
HEADER_SIZE: Final[int] = HEADER_STRUCT.size # 140
|
||||
|
||||
# --- libsodium constants the framing depends on ---
|
||||
# Not producer constants, but constants of the primitive the producer uses. Kept as literals so
|
||||
# this module stays import-free; `tests/laborious/utils/bundle/test_format.py` asserts they equal
|
||||
# the values `nacl.bindings` reports, which is what makes the literals safe.
|
||||
STREAM_HEADER_SIZE: Final[int] = 24 # crypto_secretstream_xchacha20poly1305_HEADERBYTES
|
||||
ABYTES: Final[int] = 17 # crypto_secretstream_xchacha20poly1305_ABYTES
|
||||
|
||||
# The pre-upload check prefix of `08-security-and-encryption.md` § 5.3: enough bytes to hold the
|
||||
# header, the stream header and one full chunk. Derived from constants, never read from the file.
|
||||
PREFIX_SIZE: Final[int] = HEADER_SIZE + STREAM_HEADER_SIZE + CHUNK_SIZE + ABYTES # 65717
|
||||
|
||||
# The digest covers the chunk stream only. The producer computes `sha256(ciphertext)` over the
|
||||
# joined chunks and *then* writes `header + stream_header + ciphertext`, so on disk the covered
|
||||
# range starts after both headers. Verified against a bundle produced by the producer's own code:
|
||||
# `sha256(file[164:])` matches the header field and `sha256(file[140:])` does not.
|
||||
DIGEST_OFFSET: Final[int] = HEADER_SIZE + STREAM_HEADER_SIZE # 164
|
||||
|
||||
MIN_FILE_SIZE: Final[int] = HEADER_SIZE + STREAM_HEADER_SIZE
|
||||
|
||||
SUPPORTED_FORMAT_VERSIONS: Final[frozenset[int]] = frozenset({FORMAT_VERSION})
|
||||
SUPPORTED_AEAD_IDS: Final[frozenset[int]] = frozenset({AEAD_ID_XCHACHA20POLY1305_SECRETSTREAM})
|
||||
SUPPORTED_KDF_IDS: Final[frozenset[int]] = frozenset({KDF_ID_ARGON2ID})
|
||||
|
||||
_HASH_BLOCK_SIZE: Final[int] = 1024 * 1024
|
||||
|
||||
|
||||
class BundleFormatError(Exception):
|
||||
"""The bytes are not a `.sientia` bundle this reader implements.
|
||||
|
||||
Raised for every header-level rejection: wrong magic, unknown version, unknown primitive ids,
|
||||
cost parameters above the ceiling, a reserved field carrying a value, a file too short to hold
|
||||
a header, or a chunk stream whose digest does not match the header's. `reader.py` turns this
|
||||
into a gate-2 rejection; nothing here knows about gates or import steps.
|
||||
"""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BundleHeader:
|
||||
"""The 140-byte plaintext header, parsed and validated."""
|
||||
|
||||
format_version: int
|
||||
aead_id: int
|
||||
kdf_id: int
|
||||
salt: bytes
|
||||
kdf_memlimit_bytes: int
|
||||
kdf_opslimit: int
|
||||
kdf_parallelism: int
|
||||
ciphertext_digest: bytes
|
||||
reserved_signature: bytes
|
||||
|
||||
|
||||
def parse_header(raw: bytes) -> BundleHeader:
|
||||
"""Parse and validate the plaintext header.
|
||||
|
||||
Everything the reader trusts about the file's shape is decided here, before any offset past
|
||||
the header is used and before a key is derived.
|
||||
|
||||
Args:
|
||||
raw: at least `HEADER_SIZE` bytes read from the start of the object.
|
||||
|
||||
Returns:
|
||||
BundleHeader: the validated header.
|
||||
|
||||
Raises:
|
||||
BundleFormatError: on any invalid or unsupported field.
|
||||
"""
|
||||
if len(raw) < HEADER_SIZE:
|
||||
raise BundleFormatError(
|
||||
f'file is shorter than the {HEADER_SIZE}-byte header ({len(raw)} bytes available)'
|
||||
)
|
||||
|
||||
(
|
||||
magic,
|
||||
format_version,
|
||||
aead_id,
|
||||
kdf_id,
|
||||
salt,
|
||||
kdf_memlimit_bytes,
|
||||
kdf_opslimit,
|
||||
kdf_parallelism,
|
||||
ciphertext_digest,
|
||||
reserved_signature,
|
||||
) = HEADER_STRUCT.unpack(raw[:HEADER_SIZE])
|
||||
|
||||
if magic != MAGIC:
|
||||
raise BundleFormatError('file does not start with the expected bundle marker')
|
||||
if format_version not in SUPPORTED_FORMAT_VERSIONS:
|
||||
raise BundleFormatError(f'unsupported bundle format version: {format_version}')
|
||||
if aead_id not in SUPPORTED_AEAD_IDS:
|
||||
raise BundleFormatError(f'unsupported encryption identifier: {aead_id}')
|
||||
if kdf_id not in SUPPORTED_KDF_IDS:
|
||||
raise BundleFormatError(f'unsupported key derivation identifier: {kdf_id}')
|
||||
if kdf_memlimit_bytes > MAX_KDF_MEMLIMIT:
|
||||
raise BundleFormatError(
|
||||
f'declared key derivation memory {kdf_memlimit_bytes} exceeds the '
|
||||
f'{MAX_KDF_MEMLIMIT} ceiling'
|
||||
)
|
||||
if kdf_opslimit > MAX_KDF_OPSLIMIT:
|
||||
raise BundleFormatError(
|
||||
f'declared key derivation passes {kdf_opslimit} exceeds the {MAX_KDF_OPSLIMIT} ceiling'
|
||||
)
|
||||
if kdf_memlimit_bytes <= 0 or kdf_opslimit <= 0:
|
||||
raise BundleFormatError('key derivation parameters must be positive')
|
||||
if kdf_parallelism != PARALLELISM:
|
||||
raise BundleFormatError(
|
||||
f'reserved parallelism field must be {PARALLELISM}, found {kdf_parallelism}'
|
||||
)
|
||||
if reserved_signature != bytes(RESERVED_SIGNATURE_SIZE):
|
||||
raise BundleFormatError('reserved signature field must be empty in this format version')
|
||||
|
||||
return BundleHeader(
|
||||
format_version=format_version,
|
||||
aead_id=aead_id,
|
||||
kdf_id=kdf_id,
|
||||
salt=salt,
|
||||
kdf_memlimit_bytes=kdf_memlimit_bytes,
|
||||
kdf_opslimit=kdf_opslimit,
|
||||
kdf_parallelism=kdf_parallelism,
|
||||
ciphertext_digest=ciphertext_digest,
|
||||
reserved_signature=reserved_signature,
|
||||
)
|
||||
|
||||
|
||||
def read_header(path: Path | str) -> BundleHeader:
|
||||
"""Read and validate the header of the file at `path`.
|
||||
|
||||
Raises:
|
||||
BundleFormatError: when the file cannot hold a header plus a stream header, or when the
|
||||
header itself is invalid.
|
||||
"""
|
||||
file_path = Path(path)
|
||||
size = file_path.stat().st_size
|
||||
if size < MIN_FILE_SIZE:
|
||||
raise BundleFormatError(
|
||||
f'file is {size} bytes, smaller than the {MIN_FILE_SIZE} bytes a bundle needs for its '
|
||||
'header and encrypted stream header'
|
||||
)
|
||||
with file_path.open('rb') as handle:
|
||||
return parse_header(handle.read(HEADER_SIZE))
|
||||
|
||||
|
||||
def compute_ciphertext_digest(path: Path | str) -> bytes:
|
||||
"""SHA-256 of the chunk stream — the bytes from `DIGEST_OFFSET` to the end of the file."""
|
||||
digest = hashlib.sha256()
|
||||
with Path(path).open('rb') as handle:
|
||||
handle.seek(DIGEST_OFFSET)
|
||||
while block := handle.read(_HASH_BLOCK_SIZE):
|
||||
digest.update(block)
|
||||
return digest.digest()
|
||||
|
||||
|
||||
def verify_ciphertext_digest(path: Path | str, header: BundleHeader) -> None:
|
||||
"""Compare the file's chunk-stream digest against the header's.
|
||||
|
||||
Raises:
|
||||
BundleFormatError: when they differ — a truncated, overwritten or corrupted transfer. This
|
||||
runs before any key derivation, so a damaged file never costs an Argon2id pass.
|
||||
"""
|
||||
if compute_ciphertext_digest(path) != header.ciphertext_digest:
|
||||
raise BundleFormatError(
|
||||
'the encrypted content does not match the digest recorded in the bundle header'
|
||||
)
|
||||
|
||||
|
||||
def file_digest(path: Path | str) -> str:
|
||||
"""Hex SHA-256 of the whole file, for the `expected_digest` comparison of gate 1."""
|
||||
digest = hashlib.sha256()
|
||||
with Path(path).open('rb') as handle:
|
||||
while block := handle.read(_HASH_BLOCK_SIZE):
|
||||
digest.update(block)
|
||||
return digest.hexdigest()
|
||||
@@ -1,495 +0,0 @@
|
||||
"""Reading a `.sientia` bundle: the seven gates, in order, in one function.
|
||||
|
||||
`BundleReader.open()` is the only entry point. Gate ordering is a property of that function rather
|
||||
than of the caller, so it cannot be reordered by accident and can be unit-tested without Temporal:
|
||||
|
||||
1. object-level limits (`check_object_limits`, called by the download activity before the body is
|
||||
fetched — the only gate that does not need the file)
|
||||
2. header validation, including the chunk-stream digest, before any key derivation
|
||||
3. AEAD decryption — the integrity check; a wrong password and a tampered file fail identically
|
||||
4. archive inspection over the central directory, writing nothing to disk
|
||||
5. extraction into a fresh, isolated directory
|
||||
6. structure and schema of the extracted tree
|
||||
7. content policy over the file set
|
||||
|
||||
Nothing here logs, notifies or emits a metric: rejections are raised as `BundleGateError` and the
|
||||
activity layer owns the one failure channel.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import stat
|
||||
import time
|
||||
import unicodedata
|
||||
import zipfile
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import nacl.bindings as sodium
|
||||
import nacl.exceptions
|
||||
|
||||
from laborious.utils.bundle import format as fmt
|
||||
from laborious.utils.bundle.steps import GATE_TO_STEP, ImportStep
|
||||
|
||||
# Gate 5 — chunk size for the byte-counted extraction copy loop.
|
||||
_EXTRACT_CHUNK_BYTES = 1024 * 1024
|
||||
|
||||
# Gate 6 — the layout `bundle.build_zip` writes and `model_repository.py` reads back.
|
||||
METADATA_FILE_NAME = 'metadata.json'
|
||||
ARTIFACTS_DIR_NAME = 'artifacts'
|
||||
REQUIRED_TOP_LEVEL_KEYS = frozenset({'parameters', 'metadata'})
|
||||
REQUIRED_METADATA_FIELDS = (
|
||||
'model_name',
|
||||
'experiment_name',
|
||||
'run_id',
|
||||
'model_version',
|
||||
'model_project',
|
||||
'export_timestamp',
|
||||
)
|
||||
REQUIRED_PARAMETERS = ('target_variable',)
|
||||
REQUIRED_MODEL_DIRS = ('prediction_model', 'data_model')
|
||||
REQUIRED_MODEL_FILES = ('MLmodel', 'model.pkl')
|
||||
|
||||
# The import record's `experiment_name` column is `VARCHAR(50)` with a length-3 floor, and the name
|
||||
# is written onto the record before provisioning starts — so the column's limits are enforced here,
|
||||
# not three activities later as a constraint violation.
|
||||
EXPERIMENT_NAME_MIN_LENGTH = 3
|
||||
EXPERIMENT_NAME_MAX_LENGTH = 50
|
||||
|
||||
# Gate 7 — an MLflow 2.x artifact tree for this platform: the model directories' known files, the
|
||||
# model card, the CSVs the training template logs and the bundle's own metadata.
|
||||
ALLOWED_FILE_NAMES = frozenset(
|
||||
{
|
||||
'MLmodel',
|
||||
'metadata.json',
|
||||
'model.pkl',
|
||||
'conda.yaml',
|
||||
'python_env.yaml',
|
||||
'requirements.txt',
|
||||
'model_card.json',
|
||||
'model_card.svg',
|
||||
}
|
||||
)
|
||||
ALLOWED_FILE_SUFFIXES = frozenset({'.csv', '.json', '.yaml', '.yml', '.txt', '.pkl', '.svg', '.md'})
|
||||
|
||||
_DECRYPT_FAILURE_MESSAGE = (
|
||||
'could not open this bundle: wrong password, or the file is damaged or was altered'
|
||||
)
|
||||
|
||||
|
||||
class BundleGateError(Exception):
|
||||
"""A gate rejected the bundle.
|
||||
|
||||
Carries the gate number and the import step that gate maps onto, so the workflow never has to
|
||||
guess which step to record.
|
||||
"""
|
||||
|
||||
def __init__(self, message: str, *, gate: int, step: ImportStep | None = None) -> None:
|
||||
super().__init__(message)
|
||||
self.gate = gate
|
||||
self.step = step if step is not None else GATE_TO_STEP[gate]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BundleLimits:
|
||||
"""Ceilings enforced by gates 1 and 4. Values come from `build_import_config()`."""
|
||||
|
||||
max_object_bytes: int = 1024 * 1024 * 1024
|
||||
max_entries: int = 5000
|
||||
max_uncompressed_bytes: int = 4 * 1024 * 1024 * 1024
|
||||
max_compression_ratio: float = 200.0
|
||||
object_prefix: str = 'imported_models/'
|
||||
object_suffix: str = '.sientia'
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class OpenedBundle:
|
||||
"""What a bundle yields once every gate has passed. Nothing secret is carried."""
|
||||
|
||||
extracted_dir: Path
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
parameters: dict[str, Any] = field(default_factory=dict)
|
||||
digest: str = ''
|
||||
header: fmt.BundleHeader | None = None
|
||||
# How long Argon2id took, so the worker's most expensive step is measurable without the
|
||||
# activity layer having to reach inside the reader.
|
||||
kdf_seconds: float = 0.0
|
||||
|
||||
|
||||
class BundleReader:
|
||||
"""Opens `.sientia` bundles produced by the exporter in `sientia-projects-templates`."""
|
||||
|
||||
def __init__(self, limits: BundleLimits | None = None) -> None:
|
||||
self.limits = limits or BundleLimits()
|
||||
|
||||
# ------------------------------------------------------------------ gate 1
|
||||
|
||||
def check_object_limits(self, *, size: int, key: str) -> None:
|
||||
"""Gate 1, object level: reject before the object's body is fetched.
|
||||
|
||||
Args:
|
||||
size: the object's size as reported by `stat_object`.
|
||||
key: the object key.
|
||||
|
||||
Raises:
|
||||
BundleGateError: oversized object, key outside the import prefix, or wrong suffix.
|
||||
"""
|
||||
if size > self.limits.max_object_bytes:
|
||||
raise BundleGateError(
|
||||
f'uploaded object is {size} bytes, above the '
|
||||
f'{self.limits.max_object_bytes} byte limit',
|
||||
gate=1,
|
||||
)
|
||||
if size <= 0:
|
||||
raise BundleGateError('uploaded object is empty', gate=1)
|
||||
if not key.startswith(self.limits.object_prefix):
|
||||
raise BundleGateError('uploaded object is not under the import prefix', gate=1)
|
||||
if not key.endswith(self.limits.object_suffix):
|
||||
raise BundleGateError(
|
||||
f'uploaded object does not end in {self.limits.object_suffix}', gate=1
|
||||
)
|
||||
|
||||
# --------------------------------------------------------------------- open
|
||||
|
||||
def open(
|
||||
self,
|
||||
encrypted_path: Path | str,
|
||||
password: str,
|
||||
extract_root: Path | str,
|
||||
work_dir: Path | str | None = None,
|
||||
) -> OpenedBundle:
|
||||
"""Run gates 2 to 7 in order and return the opened bundle.
|
||||
|
||||
Args:
|
||||
encrypted_path: the downloaded `.sientia` file.
|
||||
password: the plaintext bundle password. Never stored, never logged, never returned.
|
||||
extract_root: the directory to extract into. It MUST NOT exist yet (gate 5).
|
||||
work_dir: where the decrypted zip is written. Defaults to `extract_root`'s parent.
|
||||
|
||||
Returns:
|
||||
OpenedBundle: extracted tree, `metadata`, `parameters`, digest and header.
|
||||
|
||||
Raises:
|
||||
BundleGateError: with the gate that rejected the bundle and its import step.
|
||||
"""
|
||||
encrypted = Path(encrypted_path)
|
||||
root = Path(extract_root)
|
||||
staging = Path(work_dir) if work_dir is not None else root.parent
|
||||
|
||||
header = self._gate_2_header(encrypted)
|
||||
zip_path, kdf_seconds = self._gate_3_decrypt(encrypted, password, header, staging)
|
||||
try:
|
||||
entries = self._gate_4_inspect(zip_path)
|
||||
self._gate_5_extract(zip_path, root, entries)
|
||||
metadata, parameters = self._gate_6_structure(root)
|
||||
self._gate_7_content_policy(root)
|
||||
except BaseException:
|
||||
# A rejection after decryption leaves nothing extracted behind (gate 5's contract) and
|
||||
# never leaves a decrypted zip on the worker for the next import to trip over.
|
||||
shutil.rmtree(root, ignore_errors=True)
|
||||
zip_path.unlink(missing_ok=True)
|
||||
raise
|
||||
zip_path.unlink(missing_ok=True)
|
||||
|
||||
return OpenedBundle(
|
||||
extracted_dir=root,
|
||||
metadata=metadata,
|
||||
parameters=parameters,
|
||||
digest=header.ciphertext_digest.hex(),
|
||||
header=header,
|
||||
kdf_seconds=kdf_seconds,
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------ gate 2
|
||||
|
||||
def _gate_2_header(self, encrypted: Path) -> fmt.BundleHeader:
|
||||
"""Parse the header and check the chunk-stream digest, before deriving anything."""
|
||||
try:
|
||||
header = fmt.read_header(encrypted)
|
||||
fmt.verify_ciphertext_digest(encrypted, header)
|
||||
except fmt.BundleFormatError as error:
|
||||
raise BundleGateError(str(error), gate=2) from error
|
||||
except OSError as error:
|
||||
raise BundleGateError('downloaded bundle could not be read', gate=2) from error
|
||||
return header
|
||||
|
||||
# ------------------------------------------------------------------ gate 3
|
||||
|
||||
def _gate_3_decrypt(
|
||||
self, encrypted: Path, password: str, header: fmt.BundleHeader, staging: Path
|
||||
) -> tuple[Path, float]:
|
||||
"""Derive the key with the header's parameters and decrypt the stream to a file.
|
||||
|
||||
The AEAD *is* the integrity check: a wrong password and a tampered byte fail here with the
|
||||
same message, which never claims to distinguish them.
|
||||
"""
|
||||
staging.mkdir(parents=True, exist_ok=True)
|
||||
zip_path = staging / 'bundle.zip'
|
||||
try:
|
||||
started = time.monotonic()
|
||||
key = self._derive_key(password, header)
|
||||
kdf_seconds = time.monotonic() - started
|
||||
self._pull_stream(encrypted, key, zip_path)
|
||||
except BundleGateError:
|
||||
zip_path.unlink(missing_ok=True)
|
||||
raise
|
||||
except (nacl.exceptions.CryptoError, ValueError, RuntimeError) as error:
|
||||
zip_path.unlink(missing_ok=True)
|
||||
raise BundleGateError(_DECRYPT_FAILURE_MESSAGE, gate=3) from error
|
||||
except OSError as error:
|
||||
zip_path.unlink(missing_ok=True)
|
||||
raise BundleGateError('decrypted bundle could not be written', gate=3) from error
|
||||
return zip_path, kdf_seconds
|
||||
|
||||
@staticmethod
|
||||
def _derive_key(password: str, header: fmt.BundleHeader) -> bytes:
|
||||
"""Argon2id over the NFC-normalised password, with the header's own cost parameters.
|
||||
|
||||
The password is normalised and **not** trimmed, matching the producer's `_derive_key`: the
|
||||
same password typed on another OS must derive the same key, while a trailing space is part
|
||||
of the secret.
|
||||
"""
|
||||
normalized = unicodedata.normalize('NFC', password)
|
||||
return sodium.crypto_pwhash_alg(
|
||||
fmt.KEY_SIZE,
|
||||
normalized.encode('utf-8'),
|
||||
header.salt,
|
||||
header.kdf_opslimit,
|
||||
header.kdf_memlimit_bytes,
|
||||
sodium.crypto_pwhash_ALG_ARGON2ID13,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _pull_stream(encrypted: Path, key: bytes, zip_path: Path) -> None:
|
||||
"""Decrypt the chunk stream into `zip_path`, requiring the FINAL tag."""
|
||||
chunk_size = fmt.CHUNK_SIZE + fmt.ABYTES
|
||||
state = sodium.crypto_secretstream_xchacha20poly1305_state()
|
||||
with encrypted.open('rb') as source, zip_path.open('wb') as target:
|
||||
source.seek(fmt.HEADER_SIZE)
|
||||
stream_header = source.read(fmt.STREAM_HEADER_SIZE)
|
||||
sodium.crypto_secretstream_xchacha20poly1305_init_pull(state, stream_header, key)
|
||||
last_tag: int | None = None
|
||||
while chunk := source.read(chunk_size):
|
||||
plaintext, last_tag = sodium.crypto_secretstream_xchacha20poly1305_pull(
|
||||
state, chunk
|
||||
)
|
||||
target.write(plaintext)
|
||||
if last_tag != sodium.crypto_secretstream_xchacha20poly1305_TAG_FINAL:
|
||||
# A stream that stops without its FINAL tag was cut short. The digest already catches
|
||||
# truncation of a stored file; this catches a stream that was framed to look complete.
|
||||
raise BundleGateError(_DECRYPT_FAILURE_MESSAGE, gate=3)
|
||||
|
||||
# ------------------------------------------------------------------ gate 4
|
||||
|
||||
def _gate_4_inspect(self, zip_path: Path) -> list[zipfile.ZipInfo]:
|
||||
"""Inspect the central directory only. Nothing is written to disk by this gate."""
|
||||
try:
|
||||
with zipfile.ZipFile(zip_path) as archive:
|
||||
entries = archive.infolist()
|
||||
except zipfile.BadZipFile as error:
|
||||
raise BundleGateError('bundle content is not a readable archive', gate=4) from error
|
||||
|
||||
if len(entries) > self.limits.max_entries:
|
||||
raise BundleGateError(
|
||||
f'archive declares {len(entries)} entries, above the '
|
||||
f'{self.limits.max_entries} limit',
|
||||
gate=4,
|
||||
)
|
||||
if not entries:
|
||||
raise BundleGateError('archive is empty', gate=4)
|
||||
|
||||
uncompressed = sum(entry.file_size for entry in entries)
|
||||
compressed = sum(entry.compress_size for entry in entries)
|
||||
if uncompressed > self.limits.max_uncompressed_bytes:
|
||||
raise BundleGateError(
|
||||
f'archive declares {uncompressed} uncompressed bytes, above the '
|
||||
f'{self.limits.max_uncompressed_bytes} limit',
|
||||
gate=4,
|
||||
)
|
||||
ratio = uncompressed / max(compressed, 1)
|
||||
if ratio > self.limits.max_compression_ratio:
|
||||
raise BundleGateError(
|
||||
f'archive compression ratio {ratio:.1f} is above the '
|
||||
f'{self.limits.max_compression_ratio} limit',
|
||||
gate=4,
|
||||
)
|
||||
|
||||
for entry in entries:
|
||||
self._check_entry_name(entry.filename)
|
||||
self._check_entry_mode(entry)
|
||||
|
||||
return entries
|
||||
|
||||
@staticmethod
|
||||
def _check_entry_name(name: str) -> None:
|
||||
"""Reject absolute names and traversal, from the declared name alone."""
|
||||
normalised = name.replace('\\', '/')
|
||||
if normalised.startswith('/') or (len(normalised) > 1 and normalised[1] == ':'):
|
||||
raise BundleGateError('archive declares an absolute entry path', gate=4)
|
||||
if any(part == '..' for part in normalised.split('/')):
|
||||
raise BundleGateError('archive declares an entry escaping its own tree', gate=4)
|
||||
|
||||
@staticmethod
|
||||
def _check_entry_mode(entry: zipfile.ZipInfo) -> None:
|
||||
"""Reject symlinks and anything that is neither a regular file nor a directory."""
|
||||
mode = entry.external_attr >> 16
|
||||
if stat.S_IFMT(mode) == 0:
|
||||
# No file-type bits stored. This is the normal case for two kinds of entry: a zip
|
||||
# written by a tool that records no Unix mode at all (`mode == 0`), and
|
||||
# `ZipFile.writestr`, which stores permissions only — the producer's `metadata.json`
|
||||
# arrives as `0o600 << 16`. Nothing to check: the member is extracted as a plain file
|
||||
# with the worker's own permissions either way.
|
||||
return
|
||||
if stat.S_ISLNK(mode):
|
||||
raise BundleGateError('archive declares a symbolic link', gate=4)
|
||||
if not (stat.S_ISREG(mode) or stat.S_ISDIR(mode)):
|
||||
raise BundleGateError('archive declares an entry that is not a regular file', gate=4)
|
||||
|
||||
# ------------------------------------------------------------------ gate 5
|
||||
|
||||
def _gate_5_extract(self, zip_path: Path, root: Path, entries: list[zipfile.ZipInfo]) -> None:
|
||||
"""Extract into a fresh directory, containing every member by resolved path.
|
||||
|
||||
A pre-existing extraction root is a failure rather than something to clear: it means
|
||||
another import, or a previous attempt, owns that path. `entries` is the central-directory
|
||||
listing gate 4 already read; the archive is still opened here to read member bytes, but its
|
||||
`infolist()` is not read a second time.
|
||||
"""
|
||||
if root.exists():
|
||||
raise BundleGateError('extraction directory already exists', gate=5)
|
||||
root.mkdir(parents=True)
|
||||
resolved_root = root.resolve()
|
||||
try:
|
||||
with zipfile.ZipFile(zip_path) as archive:
|
||||
for entry in entries:
|
||||
self._extract_member(archive, entry, root, resolved_root)
|
||||
except BundleGateError:
|
||||
shutil.rmtree(root, ignore_errors=True)
|
||||
raise
|
||||
except (OSError, zipfile.BadZipFile) as error:
|
||||
shutil.rmtree(root, ignore_errors=True)
|
||||
raise BundleGateError('bundle content could not be unpacked', gate=5) from error
|
||||
|
||||
@staticmethod
|
||||
def _extract_member(
|
||||
archive: zipfile.ZipFile, entry: zipfile.ZipInfo, root: Path, resolved_root: Path
|
||||
) -> None:
|
||||
"""Write one member, re-checking containment on the resolved path."""
|
||||
name = entry.filename.replace('\\', '/')
|
||||
target = root / name
|
||||
resolved = Path(os.path.realpath(target))
|
||||
if resolved != resolved_root and resolved_root not in resolved.parents:
|
||||
raise BundleGateError('archive member resolves outside the extraction root', gate=5)
|
||||
if entry.is_dir():
|
||||
resolved.mkdir(parents=True, exist_ok=True)
|
||||
return
|
||||
if resolved.exists():
|
||||
raise BundleGateError('archive member would overwrite an existing file', gate=5)
|
||||
resolved.parent.mkdir(parents=True, exist_ok=True)
|
||||
# Written as a plain file: the archive's stored permissions are never restored, so no
|
||||
# member can arrive executable.
|
||||
#
|
||||
# Copied in chunks, counting bytes actually written, instead of `shutil.copyfileobj`:
|
||||
# `entry.file_size` is central-directory metadata gate 4 already checked in aggregate, but a
|
||||
# DEFLATE stream can be crafted to decompress to far more bytes than its own header
|
||||
# declares. Aborting as soon as the running count exceeds `file_size` catches that member
|
||||
# before the excess ever reaches disk, rather than trusting the declared size alone.
|
||||
written = 0
|
||||
with archive.open(entry) as source, resolved.open('wb') as sink:
|
||||
while chunk := source.read(_EXTRACT_CHUNK_BYTES):
|
||||
written += len(chunk)
|
||||
if written > entry.file_size:
|
||||
raise BundleGateError(
|
||||
f'archive member {entry.filename!r} decompresses to more than its '
|
||||
f'declared {entry.file_size} bytes',
|
||||
gate=5,
|
||||
)
|
||||
sink.write(chunk)
|
||||
|
||||
# ------------------------------------------------------------------ gate 6
|
||||
|
||||
def _gate_6_structure(self, root: Path) -> tuple[dict[str, Any], dict[str, Any]]:
|
||||
"""Validate the extracted layout and the metadata schema."""
|
||||
metadata_path = root / METADATA_FILE_NAME
|
||||
artifacts_dir = root / ARTIFACTS_DIR_NAME
|
||||
if not metadata_path.is_file():
|
||||
raise BundleGateError('bundle has no metadata document at its root', gate=6)
|
||||
if not artifacts_dir.is_dir():
|
||||
raise BundleGateError('bundle has no artifacts directory', gate=6)
|
||||
|
||||
document = self._load_metadata_document(metadata_path)
|
||||
metadata = document['metadata']
|
||||
parameters = document['parameters']
|
||||
|
||||
self._check_metadata_fields(metadata)
|
||||
self._check_parameters(parameters)
|
||||
self._check_model_directories(artifacts_dir)
|
||||
return metadata, parameters
|
||||
|
||||
@staticmethod
|
||||
def _load_metadata_document(metadata_path: Path) -> dict[str, Any]:
|
||||
"""Read `metadata.json` and check its two top-level keys."""
|
||||
try:
|
||||
document = json.loads(metadata_path.read_text(encoding='utf-8'))
|
||||
except (json.JSONDecodeError, UnicodeDecodeError, OSError) as error:
|
||||
raise BundleGateError('bundle metadata is not readable JSON', gate=6) from error
|
||||
if not isinstance(document, dict) or set(document) != REQUIRED_TOP_LEVEL_KEYS:
|
||||
raise BundleGateError(
|
||||
'bundle metadata does not carry exactly the parameters and metadata blocks', gate=6
|
||||
)
|
||||
if not isinstance(document['metadata'], dict) or not isinstance(
|
||||
document['parameters'], dict
|
||||
):
|
||||
raise BundleGateError('bundle metadata blocks are not objects', gate=6)
|
||||
return document
|
||||
|
||||
@staticmethod
|
||||
def _check_metadata_fields(metadata: dict[str, Any]) -> None:
|
||||
"""Every origin field the import records must be present and usable."""
|
||||
for field_name in REQUIRED_METADATA_FIELDS:
|
||||
value = metadata.get(field_name)
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
raise BundleGateError(f'bundle metadata is missing the {field_name} field', gate=6)
|
||||
experiment_name = metadata['experiment_name']
|
||||
if not (EXPERIMENT_NAME_MIN_LENGTH <= len(experiment_name) <= EXPERIMENT_NAME_MAX_LENGTH):
|
||||
raise BundleGateError(
|
||||
"bundle experiment name does not fit the platform's limits", gate=6
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _check_parameters(parameters: dict[str, Any]) -> None:
|
||||
"""`target_variable` is what `model_config.target` is built from; nothing else supplies it."""
|
||||
for name in REQUIRED_PARAMETERS:
|
||||
value = parameters.get(name)
|
||||
if value is None or not str(value).strip():
|
||||
raise BundleGateError(f'bundle parameters are missing {name}', gate=6)
|
||||
|
||||
@staticmethod
|
||||
def _check_model_directories(artifacts_dir: Path) -> None:
|
||||
"""`prediction_model` and `data_model` must each be a loadable MLflow model directory."""
|
||||
for directory in REQUIRED_MODEL_DIRS:
|
||||
model_dir = artifacts_dir / directory
|
||||
if not model_dir.is_dir():
|
||||
raise BundleGateError(f'bundle has no {directory} artifacts', gate=6)
|
||||
for required in REQUIRED_MODEL_FILES:
|
||||
if not (model_dir / required).is_file():
|
||||
raise BundleGateError(f'bundle {directory} artifacts are incomplete', gate=6)
|
||||
|
||||
# ------------------------------------------------------------------ gate 7
|
||||
|
||||
def _gate_7_content_policy(self, root: Path) -> None:
|
||||
"""Only the file kinds an MLflow artifact tree calls for may be present."""
|
||||
for path in root.rglob('*'):
|
||||
if path.is_dir():
|
||||
continue
|
||||
if not path.is_file():
|
||||
raise BundleGateError('bundle contains an entry that is not a file', gate=7)
|
||||
if path.name in ALLOWED_FILE_NAMES:
|
||||
continue
|
||||
if path.suffix.lower() in ALLOWED_FILE_SUFFIXES:
|
||||
continue
|
||||
raise BundleGateError('bundle contains a file a model export should not carry', gate=7)
|
||||
@@ -1,55 +0,0 @@
|
||||
"""The import pipeline's step vocabulary, and the gate-to-step mapping.
|
||||
|
||||
`ImportStep` is the join between three things: the machine-readable step written into the import
|
||||
record's JSON detail, the key of the error catalog (one code and one sentence per step) and the
|
||||
value the frontend switches on. It lives in the bundle package rather than in the activities module
|
||||
because the reader raises with it — `laborious/activities/model_import.py` imports it from here so
|
||||
there is exactly one definition.
|
||||
|
||||
It holds **pipeline** steps only. The two failures that are not steps of the pipeline — failing to
|
||||
claim the record and failing to write it — are reported by the failure channel under their own
|
||||
markers (`not_claimable`, `status_write`), deliberately outside this enum so the
|
||||
catalog-completeness test keeps meaning "every pipeline step has a sentence".
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from enum import StrEnum
|
||||
from typing import Final
|
||||
|
||||
|
||||
class ImportStep(StrEnum):
|
||||
"""Where an import stopped. Values are what gets written to the record."""
|
||||
|
||||
RECEIVED = 'received'
|
||||
DOWNLOAD = 'download'
|
||||
DECRYPTION = 'decryption'
|
||||
ARCHIVE_INSPECTION = 'archive_inspection'
|
||||
EXTRACTION = 'extraction'
|
||||
STRUCTURE_VALIDATION = 'structure_validation'
|
||||
CONTENT_POLICY = 'content_policy'
|
||||
EXPERIMENT_CREATION = 'experiment_creation'
|
||||
ARTIFACT_UPLOAD = 'artifact_upload'
|
||||
REGISTRATION = 'registration'
|
||||
MODEL_DOCUMENT = 'model_document'
|
||||
CLEANUP = 'cleanup'
|
||||
|
||||
|
||||
# The seven gates map onto the first seven steps, with gates 1 and 2 sharing one: both say "these
|
||||
# are not the bytes that were sent" — an object over the size ceiling, a key with the wrong shape, a
|
||||
# digest that disagrees with what the uploader computed, a wrong magic, an unsupported version, a
|
||||
# chunk stream that does not match the header's digest. The catalog's `download` sentence ("The
|
||||
# uploaded file could not be read, or it is not the file that was sent. Upload it again.") is true
|
||||
# for all of them, and none of them may borrow the `decryption` sentence, which implicates the
|
||||
# password.
|
||||
GATE_TO_STEP: Final[dict[int, ImportStep]] = {
|
||||
1: ImportStep.DOWNLOAD,
|
||||
2: ImportStep.DOWNLOAD,
|
||||
3: ImportStep.DECRYPTION,
|
||||
4: ImportStep.ARCHIVE_INSPECTION,
|
||||
5: ImportStep.EXTRACTION,
|
||||
6: ImportStep.STRUCTURE_VALIDATION,
|
||||
7: ImportStep.CONTENT_POLICY,
|
||||
}
|
||||
|
||||
GATE_COUNT: Final[int] = 7
|
||||
@@ -1,34 +1,67 @@
|
||||
import json
|
||||
import tempfile
|
||||
from os import getenv, path
|
||||
from os import getenv
|
||||
from typing import Any
|
||||
|
||||
|
||||
def build_mlflow_config() -> dict[str, Any]:
|
||||
"""
|
||||
Build MLFlow server configuration from environment variables.
|
||||
Read MLflow tracking and registry credentials from the environment.
|
||||
|
||||
This function constructs an MLFlow configuration dictionary from
|
||||
environment variables with sensible defaults for local development.
|
||||
It handles server connection and authentication parameters.
|
||||
Used by ``Activities`` when constructing ``SientiaMLflowRepository``. The ``url`` value is the
|
||||
same string workers and notebooks should use for ``MLFLOW_TRACKING_URI``-style clients.
|
||||
|
||||
Environment Variables:
|
||||
MLFLOW_HOST: MLFlow server hostname (default: http://localhost)
|
||||
MLFLOW_PORT: MLFlow server port (default: 5080)
|
||||
MLFLOW_USERNAME: MLFlow username (default: aignosi)
|
||||
MLFLOW_PASSWORD: MLFlow password (default: aignosi)
|
||||
MLFLOW_URL: Host with scheme
|
||||
MLFLOW_USERNAME: Basic-auth or service user (default: aignosi)
|
||||
MLFLOW_PASSWORD: Password or token (default: aignosi)
|
||||
|
||||
Returns:
|
||||
dict: MLFlow configuration dictionary with all required parameters
|
||||
Return:
|
||||
dict[str, Any]: ``url``, ``username``, ``password``.
|
||||
"""
|
||||
return {
|
||||
'host': getenv('MLFLOW_HOST', 'http://localhost'),
|
||||
'port': int(getenv('MLFLOW_PORT', '5080')),
|
||||
'url': getenv('MLFLOW_URL', 'http://localhost:5080'),
|
||||
'username': getenv('MLFLOW_USERNAME', 'aignosi'),
|
||||
'password': getenv('MLFLOW_PASSWORD', 'aignosi'),
|
||||
}
|
||||
|
||||
|
||||
def build_plugin_store_config() -> dict[str, Any]:
|
||||
"""
|
||||
Collect settings for ``PluginStore`` (Git-backed catalog + runtime install via pip).
|
||||
|
||||
Mirrors the model-manager service: the worker passes these kwargs into ``PluginStore`` after
|
||||
``install_runtime`` resolves wheels from the configured PyPI index. Missing optional env vars
|
||||
become ``None`` so the store can run without auth in local dev.
|
||||
|
||||
Environment Variables:
|
||||
STORE_BASE_URL: Git HTTP(S) server (e.g. Gitea) base URL (default: http://localhost:3000)
|
||||
STORE_OWNER: Namespace or org owning the store repo (default: sientia)
|
||||
STORE_REPO: Repository name (default: model-library-store)
|
||||
STORE_BRANCH: Checkout branch; unset lets the client use default
|
||||
STORE_USERNAME / STORE_PASSWORD: HTTP basic credentials for Git fetch
|
||||
STORE_CACHE_TTL_SECONDS: Optional integer seconds for metadata cache TTL
|
||||
PYPI_SERVER: Index URL for ``pip install`` during runtime install (default: http://localhost:5000)
|
||||
PYPI_USERNAME / PYPI_PASSWORD: Optional index authentication
|
||||
|
||||
Return:
|
||||
dict[str, Any]: Keys aligned with ``PluginStore`` constructor parameter names.
|
||||
"""
|
||||
|
||||
cache_ttl_seconds = getenv('STORE_CACHE_TTL_SECONDS')
|
||||
return {
|
||||
'base_url': getenv('STORE_BASE_URL', 'http://localhost:3000'),
|
||||
'owner': getenv('STORE_OWNER', 'sientia'),
|
||||
'repo': getenv('STORE_REPO', 'model-library-store'),
|
||||
'username': getenv('STORE_USERNAME'),
|
||||
'password': getenv('STORE_PASSWORD'),
|
||||
'branch': getenv('STORE_BRANCH'),
|
||||
'cache_ttl_seconds': int(cache_ttl_seconds) if cache_ttl_seconds else None,
|
||||
'pypi_index_url': getenv('PYPI_SERVER', 'http://localhost:5000'),
|
||||
'pypi_username': getenv('PYPI_USERNAME'),
|
||||
'pypi_password': getenv('PYPI_PASSWORD'),
|
||||
}
|
||||
|
||||
|
||||
def build_opc_config() -> dict[str, Any]:
|
||||
"""
|
||||
Build OPC server configuration from environment variables.
|
||||
@@ -74,14 +107,15 @@ def build_minio_config() -> dict[str, Any]:
|
||||
Build MinIO (S3-compatible) configuration from environment variables.
|
||||
|
||||
Environment Variables:
|
||||
MINIO_ENDPOINT: MinIO endpoint including scheme (default: http://localhost:9000)
|
||||
MINIO_ENDPOINT_URL: Host:port or URL for the S3 API (default: http://localhost:9000)
|
||||
MINIO_ACCESS_KEY: Access key (default: minioadmin)
|
||||
MINIO_SECRET_KEY: Secret key (default: minioadmin)
|
||||
MINIO_REGION: Region name for S3 client (default: us-east-1)
|
||||
MINIO_BUCKET_DEFAULT: Default bucket for uploads (default: laborious)
|
||||
MINIO_SECURE: Whether to use HTTPS (default: false)
|
||||
Returns:
|
||||
dict: MinIO configuration dictionary
|
||||
MINIO_DEFAULT_BUCKET: Default bucket for Laborious payloads (default: laborious)
|
||||
MINIO_RETENTION_HOURS: Offloaded object retention window (default: 24)
|
||||
MINIO_SECURE: If ``true``, use HTTPS (default: false)
|
||||
|
||||
Return:
|
||||
dict[str, Any]: Keys consumed by ``Activities`` / ``MinioRepository``.
|
||||
"""
|
||||
return {
|
||||
'endpoint_url': getenv('MINIO_ENDPOINT_URL', 'http://localhost:9000'),
|
||||
@@ -91,89 +125,3 @@ def build_minio_config() -> dict[str, Any]:
|
||||
'retention_hours': int(getenv('MINIO_RETENTION_HOURS', '24')),
|
||||
'secure': getenv('MINIO_SECURE', 'false') == 'true',
|
||||
}
|
||||
|
||||
|
||||
def build_import_config() -> dict[str, Any]:
|
||||
"""
|
||||
Build the `.sientia` model-import configuration from environment variables.
|
||||
|
||||
The password key is read as the raw base64 string and **not** validated here: the worker must
|
||||
start whether or not `IMPORT_MODEL_ENABLED` is set, and a missing or malformed key is an import
|
||||
failure at the decryption step (`laborious.utils.import_password.load_envelope_key`), not a
|
||||
boot failure.
|
||||
|
||||
Environment Variables:
|
||||
IMPORT_PASSWORD_KEY: base64 of the 32-byte AES-256-GCM key that wraps bundle passwords
|
||||
IMPORT_BUNDLE_BUCKET: bucket holding the uploaded bundles (default: sientia)
|
||||
IMPORT_BUNDLE_PREFIX: key prefix the uploads must sit under (default: imported_models/)
|
||||
IMPORT_WORK_DIR: worker scratch directory (default: <tempdir>/sientia-import)
|
||||
IMPORT_MAX_OBJECT_BYTES: gate 1 object size ceiling (default: 1 GiB)
|
||||
IMPORT_MAX_ARCHIVE_ENTRIES: gate 4 entry count ceiling (default: 5000)
|
||||
IMPORT_MAX_UNCOMPRESSED_BYTES: gate 4 uncompressed size ceiling (default: 4 GiB)
|
||||
IMPORT_MAX_COMPRESSION_RATIO: gate 4 compression ratio ceiling (default: 200)
|
||||
IMPORT_MODELS_COLLECTION: MongoDB collection holding the model listing (default: models)
|
||||
IMPORT_MODEL_ID_COUNTERS_COLLECTION: MongoDB collection holding the atomic id counters that
|
||||
assign `models` document ids (default: model_id_counters). **External dependency**: this
|
||||
collection and its seed value (`seq = max(int(id) for id in models)`) are not created by
|
||||
this repository — a migration must add them to
|
||||
`sientia-dataops-database-migrations/mongodb/migrations/`, following the `models`
|
||||
baseline's `ensureCollection`/`ensureIndexes` convention, before any environment's counter
|
||||
is first incremented. See `openspec/changes/model-import-fix-bugs/design.md`, Decision 3.
|
||||
IMPORT_BUNDLE_RETENTION_DAYS: lifecycle expiry for the uploaded bundle (default: 7)
|
||||
|
||||
Returns:
|
||||
dict: import configuration dictionary
|
||||
"""
|
||||
prefix = getenv('IMPORT_BUNDLE_PREFIX', 'imported_models/')
|
||||
return {
|
||||
'password_key': getenv('IMPORT_PASSWORD_KEY'),
|
||||
'bucket': getenv('IMPORT_BUNDLE_BUCKET', 'sientia'),
|
||||
'prefix': prefix if prefix.endswith('/') else f'{prefix}/',
|
||||
'work_dir': getenv('IMPORT_WORK_DIR', path.join(tempfile.gettempdir(), 'sientia-import')),
|
||||
'max_object_bytes': int(getenv('IMPORT_MAX_OBJECT_BYTES', str(1024 * 1024 * 1024))),
|
||||
'max_entries': int(getenv('IMPORT_MAX_ARCHIVE_ENTRIES', '5000')),
|
||||
'max_uncompressed_bytes': int(
|
||||
getenv('IMPORT_MAX_UNCOMPRESSED_BYTES', str(4 * 1024 * 1024 * 1024))
|
||||
),
|
||||
'max_compression_ratio': float(getenv('IMPORT_MAX_COMPRESSION_RATIO', '200')),
|
||||
'models_collection': getenv('IMPORT_MODELS_COLLECTION', 'models'),
|
||||
'counters_collection': getenv('IMPORT_MODEL_ID_COUNTERS_COLLECTION', 'model_id_counters'),
|
||||
'retention_days': int(getenv('IMPORT_BUNDLE_RETENTION_DAYS', '7')),
|
||||
}
|
||||
|
||||
|
||||
def build_import_status_config() -> dict[str, Any]:
|
||||
"""
|
||||
Build the connection to the import log database — the POSTGRES_* server, another database.
|
||||
|
||||
The import log is `public.experiment_run` in the database owned by the Spring Boot BFF
|
||||
(`sientia-core-mlops-bff`), which sits on the same Postgres server as the `sientia` database
|
||||
Laborious itself uses. Only the database name differs, so only the database name has a variable
|
||||
of its own; host, port, user and password come from POSTGRES_*. The table and schema are
|
||||
constants of the code, never inputs.
|
||||
|
||||
The four IMPORT_STATUS_DB_* connection variables remain as overrides, for a deployment where the
|
||||
log genuinely lives on another server — the e2e suite is exactly that, two containers.
|
||||
|
||||
Environment Variables:
|
||||
IMPORT_STATUS_DB_NAME: database (default: sientia-core-mlops-bff)
|
||||
POSTGRES_HOST: server host, unless IMPORT_STATUS_DB_HOST overrides it (default: localhost)
|
||||
POSTGRES_PORT: server port, unless IMPORT_STATUS_DB_PORT overrides it (default: 5432)
|
||||
POSTGRES_USER: user, unless IMPORT_STATUS_DB_USER overrides it (default: sientia)
|
||||
POSTGRES_PASSWORD: password, unless IMPORT_STATUS_DB_PASSWORD overrides it (default: empty)
|
||||
|
||||
Returns:
|
||||
dict: import status database configuration dictionary
|
||||
"""
|
||||
# Written as `override or fallback or default` rather than
|
||||
# `getenv(override) or getenv(fallback, default)`: inside an `or`, the two-argument `getenv`
|
||||
# widens to `str | None` for mypy, so the `int(...)` below fails to type-check. Putting the
|
||||
# default last as its own term keeps every value a plain `str` and reads in precedence order.
|
||||
return {
|
||||
'host': getenv('IMPORT_STATUS_DB_HOST') or getenv('POSTGRES_HOST') or 'localhost',
|
||||
'port': int(getenv('IMPORT_STATUS_DB_PORT') or getenv('POSTGRES_PORT') or '5432'),
|
||||
'dbname': getenv('IMPORT_STATUS_DB_NAME') or 'sientia-core-mlops-bff',
|
||||
'user': getenv('IMPORT_STATUS_DB_USER') or getenv('POSTGRES_USER') or 'sientia',
|
||||
'password': getenv('IMPORT_STATUS_DB_PASSWORD') or getenv('POSTGRES_PASSWORD') or '',
|
||||
'schema': 'public',
|
||||
}
|
||||
|
||||
@@ -1,97 +0,0 @@
|
||||
"""The bundle password envelope: AES-256-GCM, opened inside the activity that uses the password.
|
||||
|
||||
The workflow only ever carries `password_envelope`, so no plaintext password reaches the Temporal
|
||||
event history. This module is imported by the bundle-opening activity and by nothing else — in
|
||||
particular not by `laborious/workflows/import_model.py`, which a test asserts.
|
||||
|
||||
Envelope layout, chosen because WebCrypto gives the Angular frontend AES-GCM natively:
|
||||
|
||||
base64( nonce(12 bytes) || ciphertext || GCM tag(16 bytes) )
|
||||
|
||||
The key is a 32-byte value held in a Kubernetes Secret and exposed as `IMPORT_PASSWORD_KEY`,
|
||||
base64-encoded. Neither the key, the envelope nor the recovered password is ever logged, notified,
|
||||
used as a metric label or written to the import record.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
from typing import Final
|
||||
|
||||
from cryptography.exceptions import InvalidTag
|
||||
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
|
||||
|
||||
NONCE_SIZE: Final[int] = 12
|
||||
TAG_SIZE: Final[int] = 16
|
||||
KEY_SIZE: Final[int] = 32
|
||||
MIN_ENVELOPE_SIZE: Final[int] = NONCE_SIZE + TAG_SIZE
|
||||
|
||||
|
||||
class EnvelopeError(Exception):
|
||||
"""The envelope or its key is unusable.
|
||||
|
||||
Deliberately one type for every cause — missing key, wrong key, malformed base64, altered
|
||||
bytes — because the failure is reported at the decryption step either way. What it must never
|
||||
be confused with is a wrong *bundle* password: that one comes from the reader's gate 3, and the
|
||||
two carry different sentences.
|
||||
"""
|
||||
|
||||
|
||||
def load_envelope_key(raw_key: str | None) -> bytes:
|
||||
"""Decode the configured envelope key.
|
||||
|
||||
Args:
|
||||
raw_key: the value of `IMPORT_PASSWORD_KEY` — base64 of exactly 32 bytes.
|
||||
|
||||
Returns:
|
||||
bytes: the 32-byte key.
|
||||
|
||||
Raises:
|
||||
EnvelopeError: when the value is absent, not base64, or not 32 bytes long. There is no
|
||||
fallback interpretation: an unset key never means "the envelope is plaintext".
|
||||
"""
|
||||
if not raw_key:
|
||||
raise EnvelopeError('the import password key is not configured')
|
||||
try:
|
||||
key = base64.b64decode(raw_key, validate=True)
|
||||
except ValueError as error:
|
||||
raise EnvelopeError('the import password key is not valid base64') from error
|
||||
if len(key) != KEY_SIZE:
|
||||
raise EnvelopeError(f'the import password key must be {KEY_SIZE} bytes')
|
||||
return key
|
||||
|
||||
|
||||
def decrypt_password_envelope(envelope: str, key: bytes) -> str:
|
||||
"""Recover the bundle password from its envelope.
|
||||
|
||||
Args:
|
||||
envelope: base64 of `nonce || ciphertext || tag`.
|
||||
key: the 32-byte key from `load_envelope_key`.
|
||||
|
||||
Returns:
|
||||
str: the plaintext password, exactly as it was typed — not trimmed, not case-folded. NFC
|
||||
normalisation happens at key derivation, where the producer does it.
|
||||
|
||||
Raises:
|
||||
EnvelopeError: malformed envelope, wrong key or altered bytes (AES-GCM authentication).
|
||||
"""
|
||||
if not envelope:
|
||||
raise EnvelopeError('the import request carries no password envelope')
|
||||
try:
|
||||
payload = base64.b64decode(envelope, validate=True)
|
||||
except ValueError as error:
|
||||
raise EnvelopeError('the password envelope is not valid base64') from error
|
||||
if len(payload) < MIN_ENVELOPE_SIZE:
|
||||
raise EnvelopeError('the password envelope is too short to hold a nonce and a tag')
|
||||
|
||||
nonce, sealed = payload[:NONCE_SIZE], payload[NONCE_SIZE:]
|
||||
try:
|
||||
plaintext = AESGCM(key).decrypt(nonce, sealed, None)
|
||||
except InvalidTag as error:
|
||||
raise EnvelopeError('the password envelope could not be authenticated') from error
|
||||
except ValueError as error:
|
||||
raise EnvelopeError('the password envelope could not be opened') from error
|
||||
try:
|
||||
return plaintext.decode('utf-8')
|
||||
except UnicodeDecodeError as error:
|
||||
raise EnvelopeError('the recovered password is not valid UTF-8') from error
|
||||
@@ -1,546 +0,0 @@
|
||||
"""The import's failure vocabulary: a closed catalog, typed errors and one reporting channel.
|
||||
|
||||
Three things live here and nowhere else:
|
||||
|
||||
- **the catalog** — one stable `error_code` and one plain-language `error_reason` per pipeline step.
|
||||
The reason is written for a person who does not know this system exists: what happened, what state
|
||||
the model is in, what to do next. It never carries a path, a table name, an object key, a Python
|
||||
exception class or a traceback, and it is never built from `str(exception)`.
|
||||
- **the typed errors** — the two failures that must not be retried, because repeating them cannot
|
||||
change the answer: a record this workflow does not own, and a statement the database refuses.
|
||||
- **the channel** — `report_import_failure`, the single funnel every failure passes through. It is
|
||||
the only place in the import path that logs an error or sends a notification, which a test
|
||||
asserts, so no failure is reported twice or half-reported.
|
||||
"""
|
||||
|
||||
from temporalio import workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
import traceback
|
||||
from collections import deque
|
||||
from collections.abc import Iterator, Mapping
|
||||
from dataclasses import dataclass
|
||||
from enum import StrEnum
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, Protocol
|
||||
|
||||
from temporalio.exceptions import ApplicationError
|
||||
|
||||
from laborious import metrics
|
||||
from laborious.utils.bundle.steps import ImportStep
|
||||
|
||||
|
||||
class ImportErrorCode(StrEnum):
|
||||
"""The machine-readable half of a failure: the frontend's i18n key and support's triage key.
|
||||
|
||||
Written under `import.code` in the record's JSON detail — not in a column of its own, so the
|
||||
import adds nothing to a table it shares with the training flow.
|
||||
"""
|
||||
|
||||
REQUEST_INVALID = 'IMPORT_REQUEST_INVALID'
|
||||
FILE_UNREADABLE = 'IMPORT_FILE_UNREADABLE'
|
||||
BUNDLE_CANNOT_BE_OPENED = 'IMPORT_BUNDLE_CANNOT_BE_OPENED'
|
||||
BUNDLE_REJECTED = 'IMPORT_BUNDLE_REJECTED'
|
||||
BUNDLE_UNPACK_FAILED = 'IMPORT_BUNDLE_UNPACK_FAILED'
|
||||
BUNDLE_INCOMPLETE = 'IMPORT_BUNDLE_INCOMPLETE'
|
||||
BUNDLE_UNEXPECTED_CONTENT = 'IMPORT_BUNDLE_UNEXPECTED_CONTENT'
|
||||
STORAGE_PREPARATION_FAILED = 'IMPORT_STORAGE_PREPARATION_FAILED'
|
||||
MODEL_FILES_NOT_STORED = 'IMPORT_MODEL_FILES_NOT_STORED'
|
||||
MODEL_NOT_REGISTERED = 'IMPORT_MODEL_NOT_REGISTERED'
|
||||
MODEL_SETTINGS_NOT_SAVED = 'IMPORT_MODEL_SETTINGS_NOT_SAVED'
|
||||
UNEXPECTED_ERROR = 'IMPORT_UNEXPECTED_ERROR'
|
||||
# Codes that never reach a record's `error_message`, and are therefore absent from the catalog:
|
||||
# one names the failure to own a row, one the failure to write one, one the temporary files left
|
||||
# on the worker after an import that otherwise succeeded. All three live on the channel — log,
|
||||
# metric, notification, Temporal error — and the third is additionally flagged in the record's
|
||||
# detail as `cleanup_failed`, without a code and without changing the status.
|
||||
LOG_ROW_NOT_CLAIMABLE = 'IMPORT_LOG_ROW_NOT_CLAIMABLE'
|
||||
STATUS_NOT_RECORDED = 'IMPORT_STATUS_NOT_RECORDED'
|
||||
CLEANUP_INCOMPLETE = 'IMPORT_CLEANUP_INCOMPLETE'
|
||||
|
||||
|
||||
class ImportResponsibility(StrEnum):
|
||||
"""Who can act on a step's failure: the user (re-export, re-type, re-submit) or only support."""
|
||||
|
||||
USER = 'user'
|
||||
SUPPORT = 'support'
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ImportFailure:
|
||||
"""One classification of one failure: the step, its code, its sentence and whose to fix."""
|
||||
|
||||
step: ImportStep | None
|
||||
code: ImportErrorCode
|
||||
reason: str
|
||||
responsibility: ImportResponsibility = ImportResponsibility.SUPPORT
|
||||
|
||||
|
||||
_CONTACT_SUPPORT = "Contact support with this import's identifier."
|
||||
_RE_EXPORT = 'Export the model again from the origin platform.'
|
||||
|
||||
IMPORT_FAILURE_CATALOG: Final[
|
||||
Mapping[ImportStep, tuple[ImportErrorCode, str, ImportResponsibility]]
|
||||
] = MappingProxyType(
|
||||
{
|
||||
ImportStep.RECEIVED: (
|
||||
ImportErrorCode.REQUEST_INVALID,
|
||||
'The import could not be started because the request was incomplete or '
|
||||
'contradictory. Nothing was created. Submit the import again.',
|
||||
ImportResponsibility.USER,
|
||||
),
|
||||
ImportStep.DOWNLOAD: (
|
||||
ImportErrorCode.FILE_UNREADABLE,
|
||||
'The uploaded file could not be read, or it is not the file that was sent. '
|
||||
'Upload it again.',
|
||||
ImportResponsibility.USER,
|
||||
),
|
||||
ImportStep.DECRYPTION: (
|
||||
ImportErrorCode.BUNDLE_CANNOT_BE_OPENED,
|
||||
'This file could not be opened: the password is wrong, or the file is damaged or '
|
||||
'was altered. Confirm the password with whoever exported the model, then try '
|
||||
'again.',
|
||||
ImportResponsibility.USER,
|
||||
),
|
||||
ImportStep.ARCHIVE_INSPECTION: (
|
||||
ImportErrorCode.BUNDLE_REJECTED,
|
||||
"The file's contents did not pass the platform's safety checks, so it was not "
|
||||
f'opened. {_RE_EXPORT}',
|
||||
ImportResponsibility.USER,
|
||||
),
|
||||
ImportStep.EXTRACTION: (
|
||||
ImportErrorCode.BUNDLE_UNPACK_FAILED,
|
||||
f'The file could not be unpacked safely and nothing from it was kept. {_RE_EXPORT}',
|
||||
ImportResponsibility.USER,
|
||||
),
|
||||
ImportStep.STRUCTURE_VALIDATION: (
|
||||
ImportErrorCode.BUNDLE_INCOMPLETE,
|
||||
'This file is not a complete model export — part of what the platform needs is '
|
||||
f'missing from it. {_RE_EXPORT}',
|
||||
ImportResponsibility.USER,
|
||||
),
|
||||
ImportStep.CONTENT_POLICY: (
|
||||
ImportErrorCode.BUNDLE_UNEXPECTED_CONTENT,
|
||||
'The file contains something a model export should not contain, so it was '
|
||||
f'rejected. {_RE_EXPORT}',
|
||||
ImportResponsibility.USER,
|
||||
),
|
||||
ImportStep.EXPERIMENT_CREATION: (
|
||||
ImportErrorCode.STORAGE_PREPARATION_FAILED,
|
||||
'The platform could not prepare a place to keep this model. The import stopped '
|
||||
f'and no model was created. {_CONTACT_SUPPORT}',
|
||||
ImportResponsibility.SUPPORT,
|
||||
),
|
||||
ImportStep.ARTIFACT_UPLOAD: (
|
||||
ImportErrorCode.MODEL_FILES_NOT_STORED,
|
||||
"The model's files could not be stored on this platform. The import stopped and "
|
||||
f'the model is not available. {_CONTACT_SUPPORT}',
|
||||
ImportResponsibility.SUPPORT,
|
||||
),
|
||||
ImportStep.REGISTRATION: (
|
||||
ImportErrorCode.MODEL_NOT_REGISTERED,
|
||||
"The model's files were stored, but the model itself could not be registered, so "
|
||||
f'it cannot be used. {_CONTACT_SUPPORT}',
|
||||
ImportResponsibility.SUPPORT,
|
||||
),
|
||||
ImportStep.MODEL_DOCUMENT: (
|
||||
ImportErrorCode.MODEL_SETTINGS_NOT_SAVED,
|
||||
'The model was registered, but the settings that tell the platform how to run it '
|
||||
f'could not be saved. {_CONTACT_SUPPORT}',
|
||||
ImportResponsibility.SUPPORT,
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
UNCLASSIFIED_FAILURE: Final[tuple[ImportErrorCode, str]] = (
|
||||
ImportErrorCode.UNEXPECTED_ERROR,
|
||||
f'The import stopped for an unexpected reason and no model was created. {_CONTACT_SUPPORT}',
|
||||
)
|
||||
|
||||
# The three explicit-code failures are not pipeline steps, so `classify_import_failure` falls
|
||||
# through to `UNCLASSIFIED_FAILURE`'s reason for all of them. That sentence says "no model was
|
||||
# created", which is false for a cleanup-only failure (the model was imported) and misleading for
|
||||
# the other two (neither ever attempted the pipeline). Each gets its own reason here instead.
|
||||
_OVERRIDE_REASONS: Final[Mapping[ImportErrorCode, str]] = MappingProxyType(
|
||||
{
|
||||
ImportErrorCode.CLEANUP_INCOMPLETE: (
|
||||
'The model was imported successfully; only its temporary files could not be removed '
|
||||
'from the worker. No action needed.'
|
||||
),
|
||||
ImportErrorCode.LOG_ROW_NOT_CLAIMABLE: 'The import request could not be started.',
|
||||
ImportErrorCode.STATUS_NOT_RECORDED: 'The import outcome could not be recorded.',
|
||||
}
|
||||
)
|
||||
|
||||
# `cleanup` deletes temporary files on the worker. It is never the reported failure: when every
|
||||
# other step succeeded and only cleanup failed, the import is a success and the leak is a log line.
|
||||
# So it deliberately has no user-facing sentence, and the catalog-completeness test knows it.
|
||||
STEPS_WITHOUT_A_USER_FACING_REASON: Final[frozenset[ImportStep]] = frozenset({ImportStep.CLEANUP})
|
||||
|
||||
|
||||
def classify_import_failure(step: ImportStep | None) -> ImportFailure:
|
||||
"""Turn a failure into the one code and the one sentence that will be recorded.
|
||||
|
||||
The step alone decides the outcome — an exception's text carries paths, table names and type
|
||||
names, and belongs in the log, never in the record.
|
||||
|
||||
Args:
|
||||
step: the step that was being attempted, or `None` when none had been entered.
|
||||
|
||||
Returns:
|
||||
ImportFailure: step, code and the reason to store.
|
||||
"""
|
||||
if step is None or step not in IMPORT_FAILURE_CATALOG:
|
||||
code, reason = UNCLASSIFIED_FAILURE
|
||||
return ImportFailure(step=step, code=code, reason=reason)
|
||||
code, reason, responsibility = IMPORT_FAILURE_CATALOG[step]
|
||||
return ImportFailure(step=step, code=code, reason=reason, responsibility=responsibility)
|
||||
|
||||
|
||||
class ImportLogRowNotClaimableError(ApplicationError):
|
||||
"""The record named by `import_run_id` is not this workflow's to write.
|
||||
|
||||
Raised when the claim affects zero rows: no such row, a `TRAIN` row, or a row already running
|
||||
under another workflow. Non-retryable on purpose — a second attempt would read the same rows and
|
||||
reach the same conclusion — and loud everywhere except the row itself, which belongs to someone
|
||||
else.
|
||||
"""
|
||||
|
||||
def __init__(self, message: str, *, import_run_id: Any = None, observed: str = '') -> None:
|
||||
super().__init__(
|
||||
message,
|
||||
type='ImportLogRowNotClaimableError',
|
||||
non_retryable=True,
|
||||
)
|
||||
self.import_run_id = import_run_id
|
||||
self.observed = observed
|
||||
|
||||
|
||||
class ImportStatusRejectedError(ApplicationError):
|
||||
"""The database accepted the connection and refused the statement.
|
||||
|
||||
A `CHECK` violation, an undefined column, a value too long for its column. `V12` not being
|
||||
applied yet lands here, at the terminal write, and it must fail on the first attempt instead of
|
||||
spending the whole retry envelope to say so.
|
||||
"""
|
||||
|
||||
def __init__(self, message: str) -> None:
|
||||
super().__init__(message, type='ImportStatusRejectedError', non_retryable=True)
|
||||
|
||||
|
||||
class ImportMisconfiguredError(ApplicationError):
|
||||
"""A required repository was not injected into this worker. A deploy defect, never transient."""
|
||||
|
||||
def __init__(self, message: str) -> None:
|
||||
super().__init__(message, type='ImportMisconfiguredError', non_retryable=True)
|
||||
|
||||
|
||||
class ImportInputError(ApplicationError):
|
||||
"""The request is incomplete or contradictory. Fails at `received`, on the claimed row."""
|
||||
|
||||
def __init__(self, message: str) -> None:
|
||||
super().__init__(
|
||||
message,
|
||||
{'step': ImportStep.RECEIVED.value, 'gate': None},
|
||||
type='ImportInputError',
|
||||
non_retryable=True,
|
||||
)
|
||||
|
||||
|
||||
class ImportGateRejectedError(ApplicationError):
|
||||
"""A gate rejected the bundle, and the error says which one.
|
||||
|
||||
Gates 2 to 7 all run inside one activity, so the step the workflow was *attempting* is not
|
||||
precise enough to record. An exception loses its Python type crossing an activity boundary but
|
||||
keeps its `details`, so the step and the gate travel there and the terminal write can name them.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
message: str,
|
||||
*,
|
||||
step: ImportStep,
|
||||
gate: int | None = None,
|
||||
cause: str | None = None,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
message,
|
||||
{'step': step.value, 'gate': gate, 'cause': cause},
|
||||
type='ImportGateRejectedError',
|
||||
non_retryable=True,
|
||||
)
|
||||
self.step = step
|
||||
self.gate = gate
|
||||
|
||||
|
||||
_SUMMARY_LIMIT = 400
|
||||
_CHAIN_LIMIT = 10
|
||||
_UNKNOWN_ERROR_TYPE = 'unavailable'
|
||||
|
||||
|
||||
def walk_cause_chain(
|
||||
exception: BaseException, limit: int = _CHAIN_LIMIT
|
||||
) -> Iterator[BaseException]:
|
||||
"""Yield `exception`, then each `__cause__`, up to `limit` links. Bounded against a cycle."""
|
||||
candidate: BaseException | None = exception
|
||||
for _ in range(limit):
|
||||
if candidate is None:
|
||||
return
|
||||
yield candidate
|
||||
candidate = candidate.__cause__
|
||||
|
||||
|
||||
def _root_cause(exception: BaseException) -> BaseException:
|
||||
"""The deepest `__cause__` in the chain: the failure that actually happened.
|
||||
|
||||
A `MlflowException` reaches this channel wrapped — by an `ActivityError`, by an
|
||||
`ApplicationError`, by `requests`' own re-raises — and the outermost type is the wrapper's, which
|
||||
says nothing an operator can act on. The chain is walked to its end, and bounded, because a cycle
|
||||
in `__cause__` would otherwise hang the reporting of a failure.
|
||||
"""
|
||||
return deque(walk_cause_chain(exception), maxlen=1)[-1]
|
||||
|
||||
|
||||
def _type_name(exception: BaseException) -> str:
|
||||
"""The name to report a failure under: `ApplicationError.type` when there is one.
|
||||
|
||||
An activity's own exception type is the one thing worth grouping failures by, and Temporal keeps
|
||||
it in `ApplicationError.type` rather than in the Python class — a `MlflowException` raised in an
|
||||
activity reaches the workflow as an `ApplicationError` whose `type` still says `MlflowException`.
|
||||
Reporting the class name here would say `ApplicationError` for every activity failure there is.
|
||||
"""
|
||||
return str(getattr(exception, 'type', None) or type(exception).__name__)
|
||||
|
||||
|
||||
def _one_line(text: str) -> str:
|
||||
"""The first non-empty line, bounded — a log *field* has to stay one field.
|
||||
|
||||
The traceback still reaches the log in full, in the message. This is the value a dashboard groups
|
||||
by and an alert quotes, so a 40-frame `urllib3` chain would make it useless.
|
||||
"""
|
||||
for line in text.splitlines():
|
||||
stripped = line.strip()
|
||||
if stripped:
|
||||
return stripped[:_SUMMARY_LIMIT]
|
||||
return ''
|
||||
|
||||
|
||||
def summarize_failure(
|
||||
exception: BaseException | None,
|
||||
detail_text: str | None = None,
|
||||
error_type: str | None = None,
|
||||
) -> tuple[str, str]:
|
||||
"""`(error_type, error_summary)` — the two technical fields every failure log carries.
|
||||
|
||||
Args:
|
||||
exception: the failure, when the object itself is to hand.
|
||||
detail_text: the failure as text, for a caller whose exception crossed an activity boundary.
|
||||
error_type: the type name that travelled with such a text, when it is known.
|
||||
|
||||
Returns:
|
||||
tuple: the exception type name and a single-line summary of what it said.
|
||||
"""
|
||||
if exception is not None:
|
||||
cause = _root_cause(exception)
|
||||
return _type_name(cause), _one_line(str(cause))
|
||||
return error_type or _UNKNOWN_ERROR_TYPE, _one_line(detail_text or '')
|
||||
|
||||
|
||||
def describe_failure_chain(exception: BaseException) -> tuple[str, str]:
|
||||
"""`(error_type, detail)` for a failure that is about to lose its exception object.
|
||||
|
||||
The workflow is the one place that holds the whole chain — `ActivityError` wrapping the
|
||||
`ApplicationError` that carries the activity's own type and message — and it cannot hand that
|
||||
object to the channel, because the channel runs in an activity. So the chain is flattened here,
|
||||
**root cause first**, one link per line: the first line is what actually broke, which is what
|
||||
becomes `error_summary`, and the wrappers follow for whoever reads further.
|
||||
|
||||
`ApplicationError.type` is preferred over the Python class name because it is the only place the
|
||||
activity's original exception type survives the boundary: `MlflowException` arrives as an
|
||||
`ApplicationError` whose `type` says so.
|
||||
"""
|
||||
links: list[str] = []
|
||||
for candidate in walk_cause_chain(exception):
|
||||
kind = _type_name(candidate)
|
||||
message = _one_line(str(candidate))
|
||||
# `ApplicationError.__str__` already prefixes its own type, and `MlflowException:
|
||||
# MlflowException: …` reads like two failures rather than one.
|
||||
if not message:
|
||||
links.append(kind)
|
||||
elif message.startswith(f'{kind}:'):
|
||||
links.append(message)
|
||||
else:
|
||||
links.append(f'{kind}: {message}')
|
||||
links.reverse()
|
||||
root = links[0].split(':', 1)[0] if links else _UNKNOWN_ERROR_TYPE
|
||||
return root, '\n'.join(links)
|
||||
|
||||
|
||||
def _log_fields(
|
||||
*,
|
||||
failure: ImportFailure,
|
||||
where: str,
|
||||
import_run_id: Any,
|
||||
workflow_id: str,
|
||||
gate: int | None,
|
||||
error_type: str,
|
||||
error_summary: str,
|
||||
verdict: str | None,
|
||||
metadata: dict[str, Any] | None,
|
||||
) -> dict[str, Any]:
|
||||
"""The workflow metadata plus this failure's own fields, as `logfmt` keys.
|
||||
|
||||
Every key is prefixed `import_` or `error_` so none of them can collide with the metadata the
|
||||
rest of the runtime already logs (`model_id`, `model_name`, `schedule_name`, `workflow_name`) or
|
||||
with the formatter's own (`at`, `msg`, `timestamp`). Values are single-line by construction: the
|
||||
traceback stays in the message, where a newline costs nothing.
|
||||
"""
|
||||
fields: dict[str, Any] = {
|
||||
**(metadata or {}),
|
||||
'import_run_id': import_run_id,
|
||||
'import_workflow_id': workflow_id,
|
||||
'import_step': where,
|
||||
'import_code': failure.code.value,
|
||||
'error_type': error_type,
|
||||
}
|
||||
if error_summary:
|
||||
fields['error_summary'] = error_summary
|
||||
if gate is not None:
|
||||
fields['import_gate'] = gate
|
||||
if verdict:
|
||||
fields['import_verdict_not_recorded'] = verdict
|
||||
return fields
|
||||
|
||||
|
||||
def failure_location(exception: BaseException | None) -> tuple[str | None, int | None]:
|
||||
"""Read the step and gate an activity's failure carries, if it carried any.
|
||||
|
||||
Walks the `__cause__` chain because Temporal wraps an activity failure in an `ActivityError`.
|
||||
Returns `(None, None)` when the failure did not say where it happened, and the caller then falls
|
||||
back to the step it was attempting.
|
||||
"""
|
||||
if exception is None:
|
||||
return None, None
|
||||
for candidate in walk_cause_chain(exception):
|
||||
details = getattr(candidate, 'details', None)
|
||||
if details and isinstance(details[0], dict) and 'step' in details[0]:
|
||||
location = details[0]
|
||||
gate = location.get('gate')
|
||||
return location.get('step'), int(gate) if isinstance(gate, int) else None
|
||||
return None, None
|
||||
|
||||
|
||||
class FailureSinks(Protocol):
|
||||
"""What the channel needs from its caller: the repository's existing observability surface."""
|
||||
|
||||
def error(self, message: str, metadata: dict | None = None) -> None: ...
|
||||
|
||||
def get_core_labels(self, metadata: dict | None = None, operation_type: str = '-') -> dict: ...
|
||||
|
||||
async def emit_metric(
|
||||
self,
|
||||
metric_object: Any,
|
||||
tags: dict[str, Any],
|
||||
method: str = 'inc',
|
||||
value: float | None = None,
|
||||
) -> None: ...
|
||||
|
||||
async def send_notification_async(
|
||||
self,
|
||||
metadata: dict,
|
||||
notification_id: str,
|
||||
message: str,
|
||||
block: str,
|
||||
level: str = 'INFO',
|
||||
attachment_content: str | None = None,
|
||||
) -> None: ...
|
||||
|
||||
|
||||
async def report_import_failure(
|
||||
sinks: FailureSinks,
|
||||
step: ImportStep | None,
|
||||
exception: BaseException | None,
|
||||
*,
|
||||
import_run_id: Any,
|
||||
workflow_id: str,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
code: ImportErrorCode | None = None,
|
||||
marker: str | None = None,
|
||||
verdict: str | None = None,
|
||||
detail_text: str | None = None,
|
||||
error_type: str | None = None,
|
||||
gate: int | None = None,
|
||||
) -> ImportFailure:
|
||||
"""Report one failure through every sink, exactly once.
|
||||
|
||||
Sinks, in order: a structured log carrying the ids, the step, the code **and** the exception
|
||||
with its traceback (the only sink that gets the exception text); the import error metric, with
|
||||
the code in `operation_type`; a notification through the existing handler. The exception itself
|
||||
reaches Temporal by being raised, and the record is the fifth, optional sink — the terminal
|
||||
write, when the row can be written at all.
|
||||
|
||||
Both log lines this produces — the technical one here and the user-facing sentence the
|
||||
notification logs — carry the same `logfmt` fields, built once by `_log_fields`. That is what
|
||||
makes the two greppable and joinable: `import_run_id=29` finds the pair, `import_code=` groups
|
||||
them, and `error_type=` says what actually broke without reading a traceback. The message text is
|
||||
kept as it was, ids and all, so a query written against it still matches.
|
||||
|
||||
Args:
|
||||
sinks: the activity providing logger, metrics and notifications.
|
||||
step: the pipeline step, or `None` for a failure that is not one (see `marker`).
|
||||
exception: the failure.
|
||||
import_run_id: the record's id, for correlation.
|
||||
workflow_id: the identifier a user quotes to support.
|
||||
metadata: workflow metadata for labels and the notification.
|
||||
code: an explicit code, for the two failures that are not pipeline steps.
|
||||
marker: `not_claimable` or `status_write` — what a step name would have been.
|
||||
verdict: for a failed terminal write, the verdict it could not record, so the outcome
|
||||
survives in text even when it cannot survive in the row.
|
||||
detail_text: the failure as text, for the callers whose exception object is already gone —
|
||||
anything that crossed an activity boundary. Used in place of a traceback.
|
||||
error_type: the exception type name that travelled with such a text, when it is known.
|
||||
gate: the gate that rejected the bundle, when one did — a field, never a sentence.
|
||||
|
||||
Returns:
|
||||
ImportFailure: the classification that was reported.
|
||||
"""
|
||||
failure = classify_import_failure(step)
|
||||
if code is not None:
|
||||
failure = ImportFailure(
|
||||
step=step, code=code, reason=_OVERRIDE_REASONS.get(code, failure.reason)
|
||||
)
|
||||
|
||||
where = marker or (step.value if step else 'unknown')
|
||||
if exception is not None:
|
||||
detail = ''.join(traceback.format_exception(exception)).strip()
|
||||
else:
|
||||
detail = detail_text or 'no exception object available'
|
||||
kind, summary = summarize_failure(exception, detail_text, error_type)
|
||||
fields = _log_fields(
|
||||
failure=failure,
|
||||
where=where,
|
||||
import_run_id=import_run_id,
|
||||
workflow_id=workflow_id,
|
||||
gate=gate,
|
||||
error_type=kind,
|
||||
error_summary=summary,
|
||||
verdict=verdict,
|
||||
metadata=metadata,
|
||||
)
|
||||
verdict_line = f' verdict_not_recorded={verdict}' if verdict else ''
|
||||
sinks.error(
|
||||
f'model import failed: step={where} code={failure.code.value} '
|
||||
f'import_run_id={import_run_id} workflow_id={workflow_id}{verdict_line}\n{detail}',
|
||||
fields,
|
||||
)
|
||||
|
||||
labels = sinks.get_core_labels(metadata or {}, operation_type=failure.code.value)
|
||||
await sinks.emit_metric(metric_object=metrics.MODEL_IMPORT_ERROR_COUNT, tags=labels)
|
||||
await sinks.send_notification_async(
|
||||
metadata=fields,
|
||||
notification_id=failure.code.value,
|
||||
message=failure.reason,
|
||||
block=where,
|
||||
level='ERROR',
|
||||
)
|
||||
return failure
|
||||
@@ -21,7 +21,7 @@ from typing import Any, Literal
|
||||
|
||||
from pandas import DataFrame, read_parquet
|
||||
from sientia_do.observability.logger import Logger
|
||||
from sientia_do.repository.minio_repository import MinioRepository
|
||||
from sientia_do.repository.minio_repository_sync import MinioRepository
|
||||
from sientia_do.temporal.constants import DATETIME_FORMAT_FILENAME, DATETIME_FORMAT_WITH_TZ, now
|
||||
|
||||
# Keys that are part of the serialized wire format (not arbitrary metadata).
|
||||
@@ -89,7 +89,7 @@ class MinioDataFramePayload:
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Emit debug logs only when logger is provided
|
||||
Emit a debug message only when a logger instance is available.
|
||||
|
||||
Args:
|
||||
- logger (Logger | None): Logger instance used for debug messages
|
||||
@@ -182,7 +182,14 @@ class MinioDataFramePayload:
|
||||
|
||||
def cleanup_prefix(self) -> str | None:
|
||||
"""
|
||||
Return True if cleanup is enabled for this payload.
|
||||
Return the MinIO prefix eligible for retention cleanup.
|
||||
|
||||
Cleanup is only applicable when payload data was offloaded to MinIO
|
||||
(``object_key`` present and inline ``data`` absent). Inline-only payloads
|
||||
return ``None`` because there is no object tree to prune.
|
||||
|
||||
Return:
|
||||
str | None: Prefix used by cleanup listing, or ``None`` when cleanup does not apply.
|
||||
"""
|
||||
if self.object_key is not None and self.data is None:
|
||||
return self.object_prefix
|
||||
@@ -190,12 +197,19 @@ class MinioDataFramePayload:
|
||||
|
||||
def has_data(self) -> bool:
|
||||
"""
|
||||
Return True if the payload has some data internally or in MinIO.
|
||||
Indicate whether the payload contains retrievable tabular content.
|
||||
|
||||
A payload is considered non-empty when either inline ``data`` exists
|
||||
(and is not an empty dict) or an ``object_key`` is available for MinIO
|
||||
download.
|
||||
|
||||
Return:
|
||||
bool: ``True`` when data can be retrieved, ``False`` otherwise.
|
||||
"""
|
||||
return (self.data is not None and self.data != {}) or self.object_key is not None
|
||||
|
||||
@classmethod
|
||||
async def from_dataframe(
|
||||
def from_dataframe(
|
||||
cls,
|
||||
dataframe: DataFrame | None,
|
||||
minio_repo: MinioRepository,
|
||||
@@ -274,7 +288,7 @@ class MinioDataFramePayload:
|
||||
dataframe.to_parquet(parquet_buffer, engine='pyarrow', index=True)
|
||||
file_bytes = parquet_buffer.getvalue()
|
||||
|
||||
upload_result = await minio_repo.upload_file(
|
||||
upload_result = minio_repo.upload_file(
|
||||
file_bytes=file_bytes,
|
||||
relative_key=object_key,
|
||||
metadata=workflow_metadata,
|
||||
@@ -299,7 +313,7 @@ class MinioDataFramePayload:
|
||||
status=status,
|
||||
)
|
||||
|
||||
async def retrieve(
|
||||
def retrieve(
|
||||
self,
|
||||
minio_repo: MinioRepository,
|
||||
workflow_metadata: dict[str, Any] | None = None,
|
||||
@@ -336,7 +350,7 @@ class MinioDataFramePayload:
|
||||
f'MinioDataFramePayload.retrieve downloading object from MinIO: {self.object_key}',
|
||||
workflow_metadata,
|
||||
)
|
||||
file_bytes = await minio_repo.download_file(
|
||||
file_bytes = minio_repo.download_file(
|
||||
object_name=self.object_key, metadata=workflow_metadata
|
||||
)
|
||||
df = read_parquet(BytesIO(file_bytes))
|
||||
|
||||
@@ -1,32 +0,0 @@
|
||||
from sientia_do.notifications.handlers import NotificationHandler
|
||||
from sientia_do.observability.logger import Logger
|
||||
from sientia_do.observability.metrics_controller import MetricsController
|
||||
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
|
||||
from sientia_do.repository.minio_repository import MinioRepository
|
||||
|
||||
|
||||
class MinioManager(SientiaMonitoring):
|
||||
minio_repository: MinioRepository | None = None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
minio_repository: MinioRepository | None = None,
|
||||
logger: Logger | None = None,
|
||||
notification_handler: NotificationHandler | None = None,
|
||||
metrics_controller: MetricsController | None = None,
|
||||
):
|
||||
if self.minio_repository is None:
|
||||
self.minio_repository = minio_repository
|
||||
SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
|
||||
|
||||
def close(self) -> None:
|
||||
"""
|
||||
Close the MinioManager and clean up resources.
|
||||
"""
|
||||
if self.minio_repository is not None:
|
||||
try:
|
||||
self.minio_repository.close()
|
||||
finally:
|
||||
self.minio_repository = None
|
||||
|
||||
SientiaMonitoring.shutdown(self)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,14 +1,23 @@
|
||||
import asyncio
|
||||
"""
|
||||
Synchronous OPC UA client repository using asyncua ``sync`` API.
|
||||
|
||||
``asyncua.sync.Client`` runs the asyncio stack on a background thread so Temporal
|
||||
activities and other callers stay blocking while preserving the same session
|
||||
lifecycle, security policy, reconnect semantics, and write error classification
|
||||
as the async ``main`` implementation at ``fcc8920a8be4`` (async → sync/thread conversion).
|
||||
"""
|
||||
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from asyncua import Client
|
||||
from asyncua.crypto.security_policies import SecurityPolicyBasic256
|
||||
from asyncua.ua import DataValue, Variant, VariantType
|
||||
from asyncua import ua
|
||||
from asyncua.crypto import security_policies
|
||||
from asyncua.sync import Client
|
||||
from asyncua.ua.uaerrors import UaStatusCodeError
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
@@ -17,7 +26,6 @@ from sientia_do.observability.metrics_controller import MetricsController
|
||||
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
|
||||
|
||||
from laborious import metrics
|
||||
from laborious.utils.models.import_errors import walk_cause_chain
|
||||
|
||||
# Requested session and secure channel lifetime (ms) before server revision; 10 minutes.
|
||||
OPC_UA_SESSION_AND_CHANNEL_TIMEOUT_MS = 10 * 60 * 1000
|
||||
@@ -72,7 +80,7 @@ def _opc_authentication_token_str(client: Client | None) -> str:
|
||||
if client is None:
|
||||
return 'unknown'
|
||||
try:
|
||||
proto = client.uaclient.protocol
|
||||
proto = client.aio_obj.uaclient.protocol
|
||||
if proto is None:
|
||||
return 'unknown'
|
||||
tok = getattr(proto, 'authentication_token', None)
|
||||
@@ -93,9 +101,11 @@ def _opc_status_from_exception(exc: BaseException) -> str:
|
||||
Return:
|
||||
str: Status class name or generic Python exception name.
|
||||
"""
|
||||
for candidate in walk_cause_chain(exc):
|
||||
if isinstance(candidate, UaStatusCodeError):
|
||||
return type(candidate).__name__
|
||||
current: BaseException | None = exc
|
||||
while current is not None:
|
||||
if isinstance(current, UaStatusCodeError):
|
||||
return type(current).__name__
|
||||
current = current.__cause__
|
||||
return type(exc).__name__
|
||||
|
||||
|
||||
@@ -133,28 +143,35 @@ def _model_labels_from_write_metadata(metadata: dict[str, Any] | None) -> dict[s
|
||||
data_type_map = {
|
||||
'float': {
|
||||
'converter': float,
|
||||
'opc_type': VariantType.Float,
|
||||
'opc_type': ua.VariantType.Float,
|
||||
},
|
||||
'double': {
|
||||
'converter': float,
|
||||
'opc_type': VariantType.Double,
|
||||
'opc_type': ua.VariantType.Double,
|
||||
},
|
||||
'int': {
|
||||
'converter': int,
|
||||
'opc_type': VariantType.Int32,
|
||||
'opc_type': ua.VariantType.Int32,
|
||||
},
|
||||
'bool': {
|
||||
'converter': bool,
|
||||
'opc_type': VariantType.Boolean,
|
||||
'opc_type': ua.VariantType.Boolean,
|
||||
},
|
||||
'str': {
|
||||
'converter': str,
|
||||
'opc_type': VariantType.String,
|
||||
'opc_type': ua.VariantType.String,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class OpcRepository(SientiaMonitoring):
|
||||
"""
|
||||
Synchronous OPC UA repository for connect/disconnect and typed writes.
|
||||
|
||||
Uses ``asyncua.sync.Client`` with the same session metrics, Tier-1 Bad* reconnect,
|
||||
and structured write error payloads as the async repository on ``main``.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
opc_id: str,
|
||||
@@ -191,9 +208,9 @@ class OpcRepository(SientiaMonitoring):
|
||||
'schedule_name': '-',
|
||||
}
|
||||
self._last_write_mono: float | None = None
|
||||
self._connection_lock = asyncio.Lock()
|
||||
self._session_ready = asyncio.Event()
|
||||
self._reconnect_task: asyncio.Task[None] | None = None
|
||||
self._connection_lock = threading.Lock()
|
||||
self._session_ready = threading.Event()
|
||||
self._reconnect_thread: threading.Thread | None = None
|
||||
self._allow_reconnect = True
|
||||
|
||||
def _opc_debug_tags(self, session_id: str) -> dict[str, str]:
|
||||
@@ -224,7 +241,7 @@ class OpcRepository(SientiaMonitoring):
|
||||
if self.client is None:
|
||||
return False
|
||||
try:
|
||||
proto = self.client.uaclient.protocol
|
||||
proto = self.client.aio_obj.uaclient.protocol
|
||||
return proto is not None and proto.state != 'closed'
|
||||
except Exception:
|
||||
return False
|
||||
@@ -256,9 +273,9 @@ class OpcRepository(SientiaMonitoring):
|
||||
'level': NotificationLevel.WARNING,
|
||||
}
|
||||
|
||||
async def set_security(self) -> None:
|
||||
def set_security(self) -> None:
|
||||
"""
|
||||
Configure certificates and timeouts on the asyncua client.
|
||||
Configure certificates and timeouts on the sync asyncua client.
|
||||
|
||||
Raises:
|
||||
ValueError: If cert paths, server URI or client are missing.
|
||||
@@ -279,18 +296,20 @@ class OpcRepository(SientiaMonitoring):
|
||||
|
||||
self.client.application_uri = self.server_uri
|
||||
self.info('Setting security...', self.metadata)
|
||||
await self.client.set_security(
|
||||
SecurityPolicyBasic256,
|
||||
certificate=str(cert),
|
||||
private_key=str(private_key),
|
||||
server_certificate=str(server_cert) if server_cert else None,
|
||||
self.client.set_security(
|
||||
security_policies.SecurityPolicyBasic256,
|
||||
str(cert),
|
||||
str(private_key),
|
||||
None,
|
||||
str(server_cert) if server_cert else None,
|
||||
)
|
||||
self.client.secure_channel_timeout = OPC_UA_SESSION_AND_CHANNEL_TIMEOUT_MS
|
||||
self.client.session_timeout = OPC_UA_SESSION_AND_CHANNEL_TIMEOUT_MS
|
||||
aio = self.client.aio_obj
|
||||
aio.secure_channel_timeout = OPC_UA_SESSION_AND_CHANNEL_TIMEOUT_MS
|
||||
aio.session_timeout = OPC_UA_SESSION_AND_CHANNEL_TIMEOUT_MS
|
||||
|
||||
async def _create_client(self) -> None:
|
||||
def _create_client(self) -> None:
|
||||
"""
|
||||
Instantiate the asyncua Client and apply security when configured.
|
||||
Instantiate the sync Client and apply security when configured.
|
||||
|
||||
Caller must hold _connection_lock. Does not open a UA session.
|
||||
|
||||
@@ -303,15 +322,19 @@ class OpcRepository(SientiaMonitoring):
|
||||
'call disconnect() before creating a new client'
|
||||
)
|
||||
|
||||
self.client = Client(self.url, timeout=10, watchdog_intervall=50) # type: ignore[attr-defined]
|
||||
self.client.name = self.pod_id
|
||||
self.client = Client(self.url, timeout=10)
|
||||
aio = self.client.aio_obj
|
||||
if hasattr(aio, 'watchdog_intervall'):
|
||||
aio.watchdog_intervall = 50
|
||||
aio.name = self.pod_id
|
||||
aio.description = self.pod_id
|
||||
pod_uri = self.pod_id.replace('-', ':')
|
||||
self.client.application_uri = pod_uri
|
||||
self.client.product_uri = pod_uri
|
||||
aio.product_uri = pod_uri
|
||||
if self.cert_path:
|
||||
await self.set_security()
|
||||
self.set_security()
|
||||
|
||||
async def _open_session(self) -> tuple[bool, dict[str, Any]]:
|
||||
def _open_session(self) -> tuple[bool, dict[str, Any]]:
|
||||
"""
|
||||
Open the OPC UA session on the existing client.
|
||||
|
||||
@@ -339,30 +362,31 @@ class OpcRepository(SientiaMonitoring):
|
||||
'pod_id': self.pod_id,
|
||||
'server_name': self.server_name,
|
||||
}
|
||||
await self.emit_metric(metrics.OPC_CONNECTIONS_TOTAL, tags)
|
||||
self.emit_metric_sync(metrics.OPC_CONNECTIONS_TOTAL, tags)
|
||||
|
||||
try:
|
||||
await self.client.connect()
|
||||
self.client.connect()
|
||||
|
||||
aio = self.client.aio_obj
|
||||
session_id = _opc_authentication_token_str(self.client)
|
||||
revised_session_timeout_ms = int(self.client.session_timeout)
|
||||
revised_secure_channel_timeout_ms = int(self.client.secure_channel_timeout)
|
||||
revised_session_timeout_ms = int(aio.session_timeout)
|
||||
revised_secure_channel_timeout_ms = int(aio.secure_channel_timeout)
|
||||
self.info(
|
||||
f'OPC new session connected opc_server_id={self.id} session_id={session_id} '
|
||||
f'revised_session_timeout_ms={revised_session_timeout_ms} '
|
||||
f'revised_secure_channel_timeout_ms={revised_secure_channel_timeout_ms}',
|
||||
self.metadata,
|
||||
)
|
||||
await self.emit_metric(
|
||||
self.emit_metric_sync(
|
||||
metrics.OPC_SESSION_CREATED_TOTAL, self._opc_debug_tags(session_id)
|
||||
)
|
||||
await self.emit_metric(
|
||||
self.emit_metric_sync(
|
||||
metric_object=metrics.OPC_SESSION_REVISED_TIMEOUT_MS,
|
||||
method='set',
|
||||
tags=self._opc_debug_tags(session_id),
|
||||
value=revised_session_timeout_ms,
|
||||
)
|
||||
await self.emit_metric(
|
||||
self.emit_metric_sync(
|
||||
metric_object=metrics.OPC_CONNECTION_STATUS,
|
||||
method='set',
|
||||
tags={**tags, 'server_url': self.url},
|
||||
@@ -374,10 +398,10 @@ class OpcRepository(SientiaMonitoring):
|
||||
return True, {}
|
||||
|
||||
except Exception as e:
|
||||
await self._disconnect_locked()
|
||||
self._disconnect_locked()
|
||||
trace = traceback.format_exc()
|
||||
self.error(trace, self.metadata)
|
||||
await self.emit_metric(metrics.OPC_CONNECTIONS_FAILED, tags)
|
||||
self.emit_metric_sync(metrics.OPC_CONNECTIONS_FAILED, tags)
|
||||
return False, {
|
||||
'notification_id': f'OPC_CONNECTION_ERROR_{self.id}',
|
||||
'message': f'Failed to connect to OPC server: {e}',
|
||||
@@ -386,7 +410,7 @@ class OpcRepository(SientiaMonitoring):
|
||||
'attachment_content': trace,
|
||||
}
|
||||
|
||||
async def _connect_locked(self) -> tuple[bool, dict[str, Any]]:
|
||||
def _connect_locked(self) -> tuple[bool, dict[str, Any]]:
|
||||
"""
|
||||
Create the client when absent, then open a UA session.
|
||||
|
||||
@@ -404,10 +428,10 @@ class OpcRepository(SientiaMonitoring):
|
||||
'call disconnect() before connecting again'
|
||||
)
|
||||
if self.client is None:
|
||||
await self._create_client()
|
||||
return await self._open_session()
|
||||
self._create_client()
|
||||
return self._open_session()
|
||||
|
||||
async def _disconnection_fallback(self) -> list[dict[str, Any]]:
|
||||
def _disconnection_fallback(self) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Try up to five times to disconnect from the OPC UA server.
|
||||
"""
|
||||
@@ -419,7 +443,7 @@ class OpcRepository(SientiaMonitoring):
|
||||
f'Disconnecting from OPC UA server, attempt {i + 1} of 5',
|
||||
self.metadata,
|
||||
)
|
||||
await self.client.disconnect()
|
||||
self.client.disconnect()
|
||||
return []
|
||||
except Exception as e:
|
||||
self.error(
|
||||
@@ -433,10 +457,10 @@ class OpcRepository(SientiaMonitoring):
|
||||
'traceback': traceback.format_exc(),
|
||||
}
|
||||
)
|
||||
await asyncio.sleep(self.disconnection_interval * i)
|
||||
time.sleep(self.disconnection_interval * i)
|
||||
return error_stack
|
||||
|
||||
async def _disconnect_locked(self) -> None:
|
||||
def _disconnect_locked(self) -> None:
|
||||
"""
|
||||
Tear down the current session and client.
|
||||
|
||||
@@ -453,11 +477,11 @@ class OpcRepository(SientiaMonitoring):
|
||||
f'OPC disconnecting opc_server_id={self.id} session_id={session_id}',
|
||||
self.metadata,
|
||||
)
|
||||
await self.emit_metric(metrics.OPC_SESSION_CLOSED_TOTAL, self._opc_debug_tags(session_id))
|
||||
self.emit_metric_sync(metrics.OPC_SESSION_CLOSED_TOTAL, self._opc_debug_tags(session_id))
|
||||
|
||||
errors = await self._disconnection_fallback()
|
||||
errors = self._disconnection_fallback()
|
||||
if errors:
|
||||
await self.send_notification_async(
|
||||
self.send_notification(
|
||||
metadata=self.metadata,
|
||||
notification_id=f'OPC_DISCONNECTION_ERROR_{self.id}',
|
||||
message='Failed to disconnect from OPC server in 5 attempts.',
|
||||
@@ -468,7 +492,7 @@ class OpcRepository(SientiaMonitoring):
|
||||
else:
|
||||
self.warning(f'Disconnected from OPC server {self.id} successfully', self.metadata)
|
||||
|
||||
await self.emit_metric(
|
||||
self.emit_metric_sync(
|
||||
metric_object=metrics.OPC_CONNECTION_STATUS,
|
||||
method='set',
|
||||
tags={
|
||||
@@ -480,7 +504,7 @@ class OpcRepository(SientiaMonitoring):
|
||||
)
|
||||
self.client = None
|
||||
|
||||
async def _reconnect_locked(self) -> tuple[bool, dict[str, Any]]:
|
||||
def _reconnect_locked(self) -> tuple[bool, dict[str, Any]]:
|
||||
"""
|
||||
Close the current session and open a new one.
|
||||
|
||||
@@ -490,32 +514,32 @@ class OpcRepository(SientiaMonitoring):
|
||||
tuple[bool, dict[str, Any]]: Result from _connect_locked after teardown.
|
||||
"""
|
||||
self.last_reconnection_time = datetime.now()
|
||||
await self._disconnect_locked()
|
||||
return await self._connect_locked()
|
||||
self._disconnect_locked()
|
||||
return self._connect_locked()
|
||||
|
||||
async def connect(self) -> tuple[bool, dict[str, Any]]:
|
||||
def connect(self) -> tuple[bool, dict[str, Any]]:
|
||||
"""
|
||||
Open an OPC UA session under the connection lock (worker initialization).
|
||||
"""
|
||||
async with self._connection_lock:
|
||||
with self._connection_lock:
|
||||
self.info(
|
||||
f'Starting connection to OPC server {self.id}:{self.server_name}...',
|
||||
self.metadata,
|
||||
)
|
||||
return await self._connect_locked()
|
||||
return self._connect_locked()
|
||||
|
||||
async def disconnect(self) -> None:
|
||||
def disconnect(self) -> None:
|
||||
"""
|
||||
Gracefully disconnect from the OPC server under the connection lock.
|
||||
|
||||
Disables background reconnect so late writes during worker shutdown do not
|
||||
respawn sessions.
|
||||
"""
|
||||
async with self._connection_lock:
|
||||
with self._connection_lock:
|
||||
self._allow_reconnect = False
|
||||
await self._disconnect_locked()
|
||||
self._disconnect_locked()
|
||||
|
||||
async def validate_connection(self) -> tuple[bool, dict[str, Any]]:
|
||||
def validate_connection(self) -> tuple[bool, dict[str, Any]]:
|
||||
"""
|
||||
Read-only check that the asyncua protocol is open.
|
||||
|
||||
@@ -529,20 +553,20 @@ class OpcRepository(SientiaMonitoring):
|
||||
self.error(f'OPC server {self.id} is not connected', self.metadata)
|
||||
return False, self._not_connected_error()
|
||||
|
||||
def _reconnect_task_in_progress(self) -> bool:
|
||||
def _reconnect_thread_in_progress(self) -> bool:
|
||||
"""
|
||||
Return whether a background reconnect task is currently running.
|
||||
Return whether a background reconnect thread is currently running.
|
||||
|
||||
Return:
|
||||
bool: True when a reconnect task exists and has not finished.
|
||||
bool: True when a reconnect thread exists and is alive.
|
||||
"""
|
||||
return self._reconnect_task is not None and not self._reconnect_task.done()
|
||||
return self._reconnect_thread is not None and self._reconnect_thread.is_alive()
|
||||
|
||||
async def _start_reconnect(self, reason: str, session_id: str) -> None:
|
||||
def _start_reconnect(self, reason: str, session_id: str) -> None:
|
||||
"""
|
||||
Schedule a background reconnect when allowed by interval and task state.
|
||||
Schedule a background reconnect when allowed by interval and thread state.
|
||||
|
||||
Clears _session_ready before starting the task. No-op when _allow_reconnect is
|
||||
Clears _session_ready before starting the thread. No-op when _allow_reconnect is
|
||||
False, the reconnection window has not elapsed, or a reconnect is already running.
|
||||
|
||||
Args:
|
||||
@@ -558,7 +582,7 @@ class OpcRepository(SientiaMonitoring):
|
||||
self.metadata,
|
||||
)
|
||||
return
|
||||
if self._reconnect_task_in_progress():
|
||||
if self._reconnect_thread_in_progress():
|
||||
self.warning(
|
||||
f'OPC reconnect skipped reason=in_progress opc_server_id={self.id} '
|
||||
f'reconnect_reason={reason}',
|
||||
@@ -572,24 +596,29 @@ class OpcRepository(SientiaMonitoring):
|
||||
f'old_session_id={session_id}',
|
||||
self.metadata,
|
||||
)
|
||||
self._reconnect_task = asyncio.create_task(self._run_reconnect(reason, session_id))
|
||||
self._reconnect_thread = threading.Thread(
|
||||
target=self._run_reconnect,
|
||||
args=(reason, session_id),
|
||||
daemon=True,
|
||||
)
|
||||
self._reconnect_thread.start()
|
||||
|
||||
async def _run_reconnect(self, reason: str, session_id: str) -> None:
|
||||
def _run_reconnect(self, reason: str, session_id: str) -> None:
|
||||
"""
|
||||
Background task that tears down and re-establishes the OPC UA session.
|
||||
Tear down and re-establish the OPC UA session under the connection lock.
|
||||
|
||||
Args:
|
||||
reason (str): Trigger for reconnect (OPC status or ProtocolClosed).
|
||||
session_id (str): Previous session token string for logging.
|
||||
session_id (str): Previous session token string.
|
||||
"""
|
||||
try:
|
||||
async with self._connection_lock:
|
||||
with self._connection_lock:
|
||||
self.info(
|
||||
f'OPC reconnect started reconnect_reason={reason} opc_server_id={self.id} '
|
||||
f'old_session_id={session_id}',
|
||||
self.metadata,
|
||||
)
|
||||
success, error = await self._reconnect_locked()
|
||||
success, error = self._reconnect_locked()
|
||||
if not success:
|
||||
self.error(
|
||||
f'OPC reconnect failed reconnect_reason={reason} opc_server_id={self.id}',
|
||||
@@ -604,7 +633,7 @@ class OpcRepository(SientiaMonitoring):
|
||||
)
|
||||
self.error(traceback.format_exc(), self.metadata)
|
||||
|
||||
async def _log_write_inter_arrival(self, session_id: str, node: str) -> None:
|
||||
def _log_write_inter_arrival(self, session_id: str, node: str) -> None:
|
||||
"""
|
||||
Log elapsed wall time since the previous successful OPC write on this repository.
|
||||
|
||||
@@ -621,26 +650,18 @@ class OpcRepository(SientiaMonitoring):
|
||||
self.metadata,
|
||||
)
|
||||
if self.client is not None:
|
||||
session_timeout_ms = float(self.client.session_timeout)
|
||||
session_timeout_ms = float(self.client.aio_obj.session_timeout)
|
||||
if session_timeout_ms > 0 and delta_s > (session_timeout_ms / 1000.0):
|
||||
await self.emit_metric(
|
||||
self.emit_metric_sync(
|
||||
metrics.OPC_WRITE_INTER_ARRIVAL_OVER_SESSION_TIMEOUT_TOTAL,
|
||||
self._opc_debug_tags(session_id),
|
||||
)
|
||||
self._last_write_mono = now
|
||||
|
||||
async def _emit_opc_write_metric(
|
||||
def _emit_opc_write_metric(
|
||||
self, session_id: str, result: str, metadata: dict[str, Any] | None
|
||||
) -> None:
|
||||
"""
|
||||
Emit opc_write_attempts_total for a single write attempt outcome.
|
||||
|
||||
Args:
|
||||
session_id (str): OPC UA session token string, or "unknown".
|
||||
result (str): Outcome label (OK, OPC status name, ProtocolClosed, etc.).
|
||||
metadata (dict[str, Any] | None): Write context for model_id/model_name labels.
|
||||
"""
|
||||
await self.emit_metric(
|
||||
self.emit_metric_sync(
|
||||
metrics.OPC_WRITE_ATTEMPTS_TOTAL,
|
||||
{
|
||||
**self._opc_debug_tags(session_id),
|
||||
@@ -658,20 +679,6 @@ class OpcRepository(SientiaMonitoring):
|
||||
opc_error_kind: str | None = None,
|
||||
opc_status: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Build a structured error dict returned from failed write_data paths.
|
||||
|
||||
Args:
|
||||
notification_id (str): Stable notification identifier.
|
||||
message (str): Human-readable failure message.
|
||||
level (NotificationLevel): Severity for downstream notifications.
|
||||
attachment_content (str | None): Optional traceback or diagnostic text.
|
||||
opc_error_kind (str | None): Classifier (session_bad, connection_lost, etc.).
|
||||
opc_status (str | None): OPC UA status name or synthetic reason.
|
||||
|
||||
Return:
|
||||
dict[str, Any]: Error payload consumed by the OPC activity layer.
|
||||
"""
|
||||
payload: dict[str, Any] = {
|
||||
'notification_id': notification_id,
|
||||
'message': message,
|
||||
@@ -686,7 +693,7 @@ class OpcRepository(SientiaMonitoring):
|
||||
payload['opc_status'] = opc_status
|
||||
return payload
|
||||
|
||||
async def _handle_tier1_bad(
|
||||
def _handle_tier1_bad(
|
||||
self,
|
||||
exc: BaseException,
|
||||
session_id: str,
|
||||
@@ -710,14 +717,14 @@ class OpcRepository(SientiaMonitoring):
|
||||
opc_status = _opc_status_from_exception(exc)
|
||||
trace = traceback.format_exc()
|
||||
self.error(trace, metadata)
|
||||
await self._emit_opc_write_metric(session_id, opc_status, metadata)
|
||||
self._emit_opc_write_metric(session_id, opc_status, metadata)
|
||||
self.error(
|
||||
f'OPC write failed opc_status={opc_status} opc_server_id={self.id} '
|
||||
f'session_id={session_id} model_id={metadata.get("model_id", "unknown")} '
|
||||
f'model_name={metadata.get("model_name", "unknown")} node={node} phase={phase}',
|
||||
metadata,
|
||||
)
|
||||
await self._start_reconnect(opc_status, session_id)
|
||||
self._start_reconnect(opc_status, session_id)
|
||||
return False, self._write_failure_payload(
|
||||
notification_id=f'OPC_WRITE_DATA_ERROR_{self.id}',
|
||||
message=f'Failed to {phase} on OPC server: {exc} | metadata: {metadata}',
|
||||
@@ -726,11 +733,33 @@ class OpcRepository(SientiaMonitoring):
|
||||
opc_status=opc_status,
|
||||
)
|
||||
|
||||
async def _write_reconnect_in_progress(
|
||||
self, metadata: dict[str, Any]
|
||||
) -> tuple[bool, dict[str, Any]]:
|
||||
@staticmethod
|
||||
def _write_node_value(
|
||||
node_obj: Any,
|
||||
ua_data: ua.DataValue,
|
||||
data: Any,
|
||||
variant_type: ua.VariantType,
|
||||
) -> None:
|
||||
"""
|
||||
Fail a write because a background reconnect task is already running.
|
||||
Write a DataValue to a node, falling back to set_value when write_value is unavailable.
|
||||
|
||||
Args:
|
||||
node_obj: Sync or async node wrapper from asyncua.
|
||||
ua_data (ua.DataValue): Encoded value for write_value.
|
||||
data: Scalar converted value for set_value fallback.
|
||||
variant_type (ua.VariantType): OPC UA type for set_value fallback.
|
||||
"""
|
||||
if hasattr(node_obj, 'write_value'):
|
||||
try:
|
||||
node_obj.write_value(ua_data)
|
||||
return
|
||||
except (AttributeError, TypeError):
|
||||
pass
|
||||
node_obj.set_value(data, variant_type)
|
||||
|
||||
def _write_reconnect_in_progress(self, metadata: dict[str, Any]) -> tuple[bool, dict[str, Any]]:
|
||||
"""
|
||||
Fail a write because a background reconnect thread is already running.
|
||||
|
||||
Args:
|
||||
metadata (dict[str, Any]): Write context passed through to the activity.
|
||||
@@ -738,7 +767,7 @@ class OpcRepository(SientiaMonitoring):
|
||||
Return:
|
||||
tuple[bool, dict[str, Any]]: (False, error info with opc_error_kind reconnect_in_progress).
|
||||
"""
|
||||
await self._emit_opc_write_metric('unknown', 'ReconnectInProgress', metadata)
|
||||
self._emit_opc_write_metric('unknown', 'ReconnectInProgress', metadata)
|
||||
self.warning(
|
||||
f'OPC write rejected reconnect_in_progress opc_server_id={self.id} '
|
||||
f'model_id={metadata.get("model_id", "unknown")} '
|
||||
@@ -753,7 +782,7 @@ class OpcRepository(SientiaMonitoring):
|
||||
'opc_error_kind': 'reconnect_in_progress',
|
||||
}
|
||||
|
||||
async def _write_connection_lost(
|
||||
def _write_connection_lost(
|
||||
self, metadata: dict[str, Any], opc_status: str
|
||||
) -> tuple[bool, dict[str, Any]]:
|
||||
"""
|
||||
@@ -766,7 +795,7 @@ class OpcRepository(SientiaMonitoring):
|
||||
Return:
|
||||
tuple[bool, dict[str, Any]]: (False, error info with opc_error_kind connection_lost).
|
||||
"""
|
||||
await self._emit_opc_write_metric('unknown', opc_status, metadata)
|
||||
self._emit_opc_write_metric('unknown', opc_status, metadata)
|
||||
return False, self._write_failure_payload(
|
||||
notification_id=f'OPC_WRITE_CONNECTION_LOST_{self.id}',
|
||||
message=f'OPC write skipped: connection lost ({opc_status}) | metadata: {metadata}',
|
||||
@@ -775,7 +804,7 @@ class OpcRepository(SientiaMonitoring):
|
||||
opc_status=opc_status,
|
||||
)
|
||||
|
||||
async def write_data(
|
||||
def write_data(
|
||||
self, node: str, value: Any, data_type: str, metadata: dict[str, Any]
|
||||
) -> tuple[bool, dict[str, Any]]:
|
||||
"""
|
||||
@@ -794,35 +823,37 @@ class OpcRepository(SientiaMonitoring):
|
||||
tuple[bool, dict[str, Any]]: (True, {response_time}) on success, or
|
||||
(False, structured error info) on failure.
|
||||
"""
|
||||
if self._reconnect_task_in_progress():
|
||||
return await self._write_reconnect_in_progress(metadata)
|
||||
if self._reconnect_thread_in_progress():
|
||||
return self._write_reconnect_in_progress(metadata)
|
||||
|
||||
if not self._session_ready.is_set():
|
||||
session_id = _opc_authentication_token_str(self.client)
|
||||
await self._start_reconnect('SessionNotReady', session_id)
|
||||
if self._reconnect_task_in_progress():
|
||||
return await self._write_reconnect_in_progress(metadata)
|
||||
return await self._write_connection_lost(metadata, 'SessionNotReady')
|
||||
self._start_reconnect('SessionNotReady', session_id)
|
||||
if self._reconnect_thread_in_progress():
|
||||
return self._write_reconnect_in_progress(metadata)
|
||||
return self._write_connection_lost(metadata, 'SessionNotReady')
|
||||
|
||||
is_connected, _error = await self.validate_connection()
|
||||
is_connected, _error = self.validate_connection()
|
||||
if not is_connected:
|
||||
session_id = _opc_authentication_token_str(self.client)
|
||||
await self._start_reconnect('ProtocolClosed', session_id)
|
||||
return await self._write_connection_lost(metadata, 'ProtocolClosed')
|
||||
self._start_reconnect('ProtocolClosed', session_id)
|
||||
return self._write_connection_lost(metadata, 'ProtocolClosed')
|
||||
|
||||
start_time = time.time()
|
||||
session_id = _opc_authentication_token_str(self.client)
|
||||
|
||||
try:
|
||||
node_obj = self.client.get_node(node) # type: ignore[union-attr]
|
||||
if self.client is None:
|
||||
raise OpcClientNotInitializedError(
|
||||
f'OPC client is not initialized for server {self.id}; cannot write {node}'
|
||||
)
|
||||
node_obj = self.client.get_node(node)
|
||||
except Exception as e:
|
||||
if is_reconnectable_opcua_bad(e):
|
||||
return await self._handle_tier1_bad(e, session_id, node, metadata, 'get_node')
|
||||
return self._handle_tier1_bad(e, session_id, node, metadata, 'get_node')
|
||||
trace = traceback.format_exc()
|
||||
self.error(trace, metadata)
|
||||
await self._emit_opc_write_metric(
|
||||
session_id, f'GetNodeError:{type(e).__name__}', metadata
|
||||
)
|
||||
self._emit_opc_write_metric(session_id, f'GetNodeError:{type(e).__name__}', metadata)
|
||||
return False, self._write_failure_payload(
|
||||
notification_id=f'OPC_WRITE_GET_NODE_ERROR_{self.id}',
|
||||
message=f'Failed to get node from OPC server: {e} | metadata: {metadata}',
|
||||
@@ -830,7 +861,7 @@ class OpcRepository(SientiaMonitoring):
|
||||
)
|
||||
|
||||
if data_type not in data_type_map:
|
||||
await self._emit_opc_write_metric(session_id, 'UnsupportedDataType', metadata)
|
||||
self._emit_opc_write_metric(session_id, 'UnsupportedDataType', metadata)
|
||||
return False, self._write_failure_payload(
|
||||
notification_id=f'OPC_WRITE_DATA_TYPE_ERROR_{self.id}',
|
||||
message=f'Unsupported data type: {data_type} | metadata: {metadata}',
|
||||
@@ -838,28 +869,29 @@ class OpcRepository(SientiaMonitoring):
|
||||
|
||||
data = data_type_map[data_type]['converter'](value)
|
||||
self.info(f'Writing {data} - {type(data)} to {node}', metadata)
|
||||
ua_data = DataValue(
|
||||
Variant(data, data_type_map[data_type]['opc_type']),
|
||||
variant_type = data_type_map[data_type]['opc_type']
|
||||
ua_data = ua.DataValue(
|
||||
ua.Variant(data, variant_type),
|
||||
)
|
||||
|
||||
try:
|
||||
await node_obj.write_value(ua_data)
|
||||
self._write_node_value(node_obj, ua_data, data, variant_type)
|
||||
end_time = time.time()
|
||||
response_time = end_time - start_time
|
||||
except Exception as e:
|
||||
if is_reconnectable_opcua_bad(e):
|
||||
return await self._handle_tier1_bad(e, session_id, node, metadata, 'write_value')
|
||||
return self._handle_tier1_bad(e, session_id, node, metadata, 'write_value')
|
||||
trace = traceback.format_exc()
|
||||
self.error(trace, metadata)
|
||||
await self._emit_opc_write_metric(session_id, type(e).__name__, metadata)
|
||||
self._emit_opc_write_metric(session_id, type(e).__name__, metadata)
|
||||
return False, self._write_failure_payload(
|
||||
notification_id=f'OPC_WRITE_DATA_ERROR_{self.id}',
|
||||
message=f'Failed to write data to OPC server: {e} | metadata: {metadata}',
|
||||
attachment_content=trace,
|
||||
)
|
||||
|
||||
await self._emit_opc_write_metric(session_id, 'OK', metadata)
|
||||
await self._log_write_inter_arrival(session_id, node)
|
||||
self._emit_opc_write_metric(session_id, 'OK', metadata)
|
||||
self._log_write_inter_arrival(session_id, node)
|
||||
|
||||
return True, {
|
||||
'response_time': response_time,
|
||||
|
||||
@@ -1,40 +1,32 @@
|
||||
"""
|
||||
Laborious Worker Module
|
||||
|
||||
This module provides the main worker implementation for the Sientia DataOps Laborious system.
|
||||
It orchestrates Temporal workers, manages task queues, and handles the lifecycle of
|
||||
prediction and retraining workflows.
|
||||
Entry process that connects to Temporal, registers Laborious activities, and runs four workers in
|
||||
parallel. Each worker shares the same ``Activities`` instance (single Postgres pool, single MLflow
|
||||
repository, single PluginStore handle) but polls a different task queue.
|
||||
|
||||
The worker supports multiple runtime-scoped task queues (via ``sientia_do.temporal.worker.prepare_worker``):
|
||||
- predictions_batch-{runtime}-queue: Batch prediction workflows (heavy workload)
|
||||
- minimal_retrain-{runtime}-queue: Model retraining workflows
|
||||
- drift-{runtime}-queue: Drift detection workflows
|
||||
- simple_metrics-{runtime}-queue: Simple metrics workflows
|
||||
- import_model-{runtime}-queue: `.sientia` model import — registered **only** when
|
||||
``IMPORT_MODEL_ENABLED`` is truthy, so the containment boundary stays a deployment decision
|
||||
Task queues (see ``sientia_do.temporal.worker.prepare_worker``):
|
||||
- ``predictions_batch-{runtime}-queue`` + sub-workflows on the same queue (ML-heavy path).
|
||||
- ``minimal_retrain-{runtime}-queue`` (retrain + promote + export).
|
||||
- ``drift-core-queue`` and ``simple_metrics-core-queue``: pinned to the ``core`` runtime because
|
||||
drift and metrics do not load model code, so every runtime shares the same queues.
|
||||
|
||||
``RUNTIME`` must be set; it is passed to every ``prepare_worker`` call. Schedulers must use the
|
||||
same queue names (breaking change vs legacy ``drift-queue`` / ``simple_metrics-queue``).
|
||||
Bootstrap order:
|
||||
1. Prometheus app metrics and Mongo-backed notification handler.
|
||||
2. ``RUNTIME`` validation and ``PluginStore.install_runtime`` so ``SientiaModel`` code is importable.
|
||||
3. ``Activities`` construction (builds ``SientiaMLflowRepository`` internally from env).
|
||||
4. OPC client initialization inside activities.
|
||||
5. Temporal ``Runtime`` with SDK Prometheus bind, client connect, then ``prepare_worker`` per workflow.
|
||||
|
||||
Key Features:
|
||||
- Resource-based scaling with WorkerTuner (CPU and memory aware)
|
||||
- Automatic polling scaling with PollerBehaviorAutoscaling
|
||||
- Prometheus metrics integration
|
||||
- Comprehensive error handling and logging
|
||||
- Graceful shutdown with cleanup
|
||||
- Multiple worker instances for different workflow types
|
||||
Shutdown closes workers, notifications, activities (pools + OPC), and clears ``app_up``.
|
||||
|
||||
Environment Variables:
|
||||
- RUNTIME: Required non-empty string; suffix for all task queue names
|
||||
- TEMPORAL_HOST: Temporal server address (default: localhost:7233)
|
||||
- TEMPORAL_NAMESPACE: Temporal namespace (default: laborious)
|
||||
- POD_ID: Kubernetes pod identifier for metrics
|
||||
- HTTP_METRICS_PORT: Prometheus metrics server port (default: 9090)
|
||||
- HTTP_SDK_METRICS_PORT: Temporal SDK metrics port (default: 9091)
|
||||
- PROJECT_NAME: Project name for notifications (default: laborious)
|
||||
- IMPORT_MODEL_ENABLED: register the fifth (import) worker; unset by default
|
||||
- IMPORT_PASSWORD_KEY / IMPORT_BUNDLE_* / IMPORT_STATUS_DB_NAME: import configuration, see
|
||||
``laborious.utils.connectors_config`` (the import log's server is POSTGRES_*)
|
||||
- RUNTIME: Required non-empty string passed to ``install_runtime``.
|
||||
- STORE_* / PYPI_*: Plugin store and private index (see ``build_plugin_store_config``).
|
||||
- TEMPORAL_HOST, TEMPORAL_NAMESPACE: Cluster connection.
|
||||
- HOSTNAME (pod id), HTTP_METRICS_PORT, HTTP_SDK_METRICS_PORT: Observability.
|
||||
- PROJECT_NAME, MONGODB_*: Notifications (via ``build_mongodb_config`` in handler).
|
||||
- POSTGRES_*, MINIO_*, OPC_*, PI_WEB_API_*, MLFLOW_*: Passed through ``Activities`` helpers.
|
||||
"""
|
||||
|
||||
from temporalio import client, workflow
|
||||
@@ -48,24 +40,23 @@ with workflow.unsafe.imports_passed_through():
|
||||
from prometheus_client import start_http_server
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||
from sientia_do.observability.logger import get_logger
|
||||
from sientia_do.observability.metrics_controller import MetricsController
|
||||
from sientia_do.temporal.worker.prepare_worker import prepare_worker
|
||||
from sientia_do.utils.connectors_config import (
|
||||
build_api_config,
|
||||
build_mongodb_config,
|
||||
build_postgres_config,
|
||||
)
|
||||
from sientia_model.model_repository.plugin_store import PluginStore
|
||||
|
||||
from laborious import metrics
|
||||
from laborious.activities.activities import Activities
|
||||
from laborious.utils.connectors_config import (
|
||||
build_import_config,
|
||||
build_import_status_config,
|
||||
build_minio_config,
|
||||
build_mlflow_config,
|
||||
build_opc_config,
|
||||
build_plugin_store_config,
|
||||
)
|
||||
from laborious.workflows.drift import Drift
|
||||
from laborious.workflows.import_model import ImportModel
|
||||
from laborious.workflows.minimal_retrain import MinimalRetrain
|
||||
from laborious.workflows.predictions_batch import PredictionsBatch
|
||||
from laborious.workflows.simple_metrics import SimpleMetrics
|
||||
@@ -77,39 +68,21 @@ with workflow.unsafe.imports_passed_through():
|
||||
POD_ID = os.getenv('HOSTNAME')
|
||||
SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', '9091'))
|
||||
|
||||
# The values that turn the import worker on. Anything else — unset, empty, `0`, `false`, a typo —
|
||||
# leaves the four default workers alone and polls no import queue at all.
|
||||
IMPORT_ENABLED_VALUES = ('1', 'true', 'yes', 'on')
|
||||
|
||||
|
||||
def import_worker_enabled() -> bool:
|
||||
"""Whether this runtime registers the fifth (import) worker.
|
||||
|
||||
Read at worker start, not per workflow: containment (QTZPOC-19) wants the importer in its own
|
||||
runtime with its own credentials, and this flag is what makes that a deployment decision.
|
||||
"""
|
||||
return os.getenv('IMPORT_MODEL_ENABLED', '').strip().lower() in IMPORT_ENABLED_VALUES
|
||||
|
||||
|
||||
async def main():
|
||||
"""
|
||||
Main entry point for the Laborious worker application.
|
||||
Run the full worker lifecycle: metrics, notifications, runtime install, workers, gather.
|
||||
|
||||
This function initializes and starts all components of the worker:
|
||||
1. Sets up logging and metadata
|
||||
2. Starts Prometheus metrics server
|
||||
3. Initializes notification handler
|
||||
4. Creates and configures activities
|
||||
5. Initializes OPC connections
|
||||
6. Starts Temporal client and workers
|
||||
7. Manages worker lifecycle and graceful shutdown
|
||||
|
||||
The function runs indefinitely until interrupted or an error occurs.
|
||||
On error, it performs cleanup and exits with a non-zero status code.
|
||||
Exits the process with code 0 on normal completion of all worker tasks, or 1 after logging
|
||||
if any worker raises. ``finally`` always shuts down notifications and activities and sets
|
||||
``app_up`` to 0; ``sys.exit`` runs after it.
|
||||
|
||||
Raises:
|
||||
Exception: Any unhandled exception during worker execution
|
||||
SystemExit: On graceful shutdown or error conditions
|
||||
Exception: Propagated from ``asyncio.gather`` only before ``finally`` handling; typically
|
||||
workers run until cancelled.
|
||||
|
||||
Return:
|
||||
None (process terminates via ``sys.exit`` after the ``finally`` block).
|
||||
"""
|
||||
host = os.getenv('TEMPORAL_HOST', 'localhost:7233')
|
||||
logger = get_logger(__name__)
|
||||
@@ -124,6 +97,21 @@ async def main():
|
||||
|
||||
logger.custom_info(f'Starting Worker with POD_ID: {POD_ID}', metadata)
|
||||
|
||||
logger.custom_info('Starting prometheus client...', metadata)
|
||||
start_prometheus_server(logger, metadata)
|
||||
|
||||
logger.custom_info('Starting Notification Handler...', metadata)
|
||||
|
||||
mongo_config = build_mongodb_config()
|
||||
notification_handler = NotificationHandler(
|
||||
connection_string=mongo_config['connection_string'],
|
||||
database=mongo_config['database_name'],
|
||||
logger=logger,
|
||||
project_name=os.getenv('PROJECT_NAME', 'laborious'),
|
||||
)
|
||||
|
||||
metrics_controller = MetricsController(logger=logger)
|
||||
|
||||
runtime = os.getenv('RUNTIME', '').strip()
|
||||
if not runtime:
|
||||
logger.custom_critical(
|
||||
@@ -134,37 +122,56 @@ async def main():
|
||||
sys.exit(1)
|
||||
|
||||
metadata_runtime = {**metadata, 'runtime': runtime}
|
||||
logger.custom_info(f'Installing PluginStore runtime: {runtime}', metadata_runtime)
|
||||
|
||||
logger.custom_info('Starting prometheus client...', metadata_runtime)
|
||||
start_prometheus_server(logger, metadata_runtime)
|
||||
|
||||
logger.custom_info('Starting Notification Handler...', metadata_runtime)
|
||||
|
||||
mongo_config = build_mongodb_config()
|
||||
notification_handler = NotificationHandler(
|
||||
connection_string=mongo_config['connection_string'],
|
||||
database=mongo_config['database_name'],
|
||||
ps_cfg = build_plugin_store_config()
|
||||
plugin_store = PluginStore(
|
||||
base_url=ps_cfg['base_url'],
|
||||
owner=ps_cfg['owner'],
|
||||
repo=ps_cfg['repo'],
|
||||
username=ps_cfg['username'],
|
||||
password=ps_cfg['password'],
|
||||
branch=ps_cfg['branch'],
|
||||
cache_ttl_seconds=ps_cfg['cache_ttl_seconds'],
|
||||
pypi_index_url=ps_cfg['pypi_index_url'],
|
||||
pypi_username=ps_cfg['pypi_username'],
|
||||
pypi_password=ps_cfg['pypi_password'],
|
||||
logger=logger,
|
||||
project_name=os.getenv('PROJECT_NAME', 'laborious'),
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
if runtime == 'legacy':
|
||||
to_install_runtime = 'single'
|
||||
else:
|
||||
to_install_runtime = runtime
|
||||
|
||||
try:
|
||||
await plugin_store.install_runtime(
|
||||
runtime_name=to_install_runtime, metadata=metadata_runtime
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.custom_critical(
|
||||
f'Failed to install runtime {to_install_runtime}: {exc}', metadata_runtime
|
||||
)
|
||||
metrics.APP_UP.labels(pod_id=POD_ID).set(0)
|
||||
sys.exit(1)
|
||||
|
||||
logger.custom_info('Starting Activities...', metadata_runtime)
|
||||
|
||||
activities = Activities(
|
||||
postgres_config=build_postgres_config(),
|
||||
mlflow_config=build_mlflow_config(),
|
||||
plugin_store=plugin_store,
|
||||
minio_config=build_minio_config(),
|
||||
opc_config=build_opc_config(),
|
||||
pi_web_api_config=build_api_config(),
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
mongo_config=mongo_config,
|
||||
import_config=build_import_config(),
|
||||
import_status_config=build_import_status_config(),
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
logger.custom_info('Initializing OPC...', metadata_runtime)
|
||||
await activities.init_opc()
|
||||
activities.init_opc()
|
||||
|
||||
logger.custom_info(
|
||||
f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...',
|
||||
@@ -212,7 +219,7 @@ async def main():
|
||||
activities.export_data_to_postgres,
|
||||
],
|
||||
logger=logger,
|
||||
runtime=runtime,
|
||||
runtime='core',
|
||||
),
|
||||
prepare_worker(
|
||||
temporal_client=temporal_client,
|
||||
@@ -225,7 +232,7 @@ async def main():
|
||||
activities.export_data_to_postgres,
|
||||
],
|
||||
logger=logger,
|
||||
runtime=runtime,
|
||||
runtime='core',
|
||||
),
|
||||
prepare_worker(
|
||||
temporal_client=temporal_client,
|
||||
@@ -242,6 +249,8 @@ async def main():
|
||||
activities.format_transformed_data,
|
||||
activities.format_prediction,
|
||||
activities.format_default_prediction,
|
||||
# Prediction metadata
|
||||
activities.write_prediction_metadata,
|
||||
# OPC
|
||||
activities.write_opc_data,
|
||||
# Postgres / MinIO offload
|
||||
@@ -259,41 +268,6 @@ async def main():
|
||||
),
|
||||
]
|
||||
|
||||
# The fifth worker is opt-in: containment (QTZPOC-19) wants the importer in its own runtime with
|
||||
# its own credentials, and a flag lets that be a deployment decision instead of forcing every
|
||||
# Laborious runtime to poll the import queue today. Unset, the four workers above are unaffected
|
||||
# and no import queue is polled at all.
|
||||
if import_worker_enabled():
|
||||
logger.custom_info('Import model worker enabled', metadata_runtime)
|
||||
workers.append(
|
||||
prepare_worker(
|
||||
temporal_client=temporal_client,
|
||||
main_workflow=ImportModel,
|
||||
other_workflows=[],
|
||||
activities=[
|
||||
# Import log row (the BFF database, update-only)
|
||||
activities.claim_import_status,
|
||||
activities.record_import_status,
|
||||
activities.record_import_names,
|
||||
activities.record_import_terminal_status,
|
||||
activities.report_import_status_write_failure,
|
||||
activities.report_import_pipeline_failure,
|
||||
# Bundle download/open, MLflow experiment creation and artifact upload — one
|
||||
# fused activity, so they never split across worker processes.
|
||||
activities.prepare_and_upload_import_bundle,
|
||||
# MLflow provisioning
|
||||
activities.register_import_model_version,
|
||||
# The model listing
|
||||
activities.write_import_model_document,
|
||||
# Cleanup and retention
|
||||
activities.cleanup_import_files,
|
||||
activities.ensure_import_bundle_retention,
|
||||
],
|
||||
logger=logger,
|
||||
runtime=runtime,
|
||||
)
|
||||
)
|
||||
|
||||
handlers = []
|
||||
for w in workers:
|
||||
handlers.append(w.run())
|
||||
@@ -302,8 +276,6 @@ async def main():
|
||||
|
||||
exit_code = 0
|
||||
try:
|
||||
# This will run the workers and wait for them to complete.
|
||||
# If an exception occurs in any of the worker handlers, it will be propagated here.
|
||||
await asyncio.gather(*handlers)
|
||||
except Exception as e:
|
||||
logger.custom_error(f'An unhandled exception occurred: {e}', metadata_runtime)
|
||||
@@ -315,10 +287,8 @@ async def main():
|
||||
# let us relabel its exit code before `sys.exit(exit_code)` below discards it anyway.
|
||||
# Letting it propagate on its own keeps the original code (e.g. a handler's sys.exit(N))
|
||||
# intact instead of collapsing it to 0.
|
||||
if notification_handler:
|
||||
notification_handler.shutdown()
|
||||
if activities:
|
||||
await activities.shutdown()
|
||||
notification_handler.shutdown()
|
||||
activities.shutdown()
|
||||
metrics.APP_UP.labels(pod_id=POD_ID).set(0) # Mark app as DOWN
|
||||
|
||||
sys.exit(exit_code)
|
||||
@@ -336,7 +306,7 @@ def start_prometheus_server(logger, metadata: dict):
|
||||
|
||||
Environment Variables:
|
||||
HTTP_METRICS_PORT: Port for the metrics server (default: 9090)
|
||||
POD_ID: Pod identifier for metrics labeling
|
||||
HOSTNAME: Pod identifier for metrics labeling
|
||||
|
||||
Raises:
|
||||
SystemExit: If the metrics server fails to start
|
||||
|
||||
@@ -37,6 +37,7 @@ class Drift:
|
||||
model_config = input_data['model_config']
|
||||
target_name = model_config['target']
|
||||
|
||||
# Built from internal Temporal workflow config, not user input (bandit B608).
|
||||
gathering_query = f"""
|
||||
SELECT *
|
||||
FROM "{input_data['schema']}"."{input_data['source_table_name']}"
|
||||
@@ -44,7 +45,7 @@ class Drift:
|
||||
model_id = '{input_data['model_id']}' AND
|
||||
timestamp > NOW() - INTERVAL '{input_data['interval']} minutes'
|
||||
ORDER BY timestamp ASC
|
||||
""" # nosec B608 - values come from internal Temporal workflow config, not user input
|
||||
""" # nosec
|
||||
|
||||
target_data_handler = workflow.start_activity_method(
|
||||
Activities.load_custom_query,
|
||||
|
||||
@@ -1,448 +0,0 @@
|
||||
"""The `import_model` workflow: a `.sientia` object becomes a registered model, or a durable ERROR.
|
||||
|
||||
Shape of this workflow, and the reasons it has that shape:
|
||||
|
||||
- **The claim comes first.** The record already exists — the frontend inserted it through the BFF —
|
||||
so the workflow's first act is to take ownership of it. Everything after that failure is recorded
|
||||
*on* the record; a failed claim is recorded nowhere, because the workflow owns no row.
|
||||
- **Fail-fast, and nothing else.** No rollback, no compensation, no resume, no reconciliation, no
|
||||
`continue_as_new`. A failed import stays failed; trying again is a new import with a new upload.
|
||||
- **Only the status writes retry**, and only because `UPDATE ... SET <fixed values> WHERE id` is
|
||||
idempotent. Every provisioning call keeps `maximum_attempts=1`, because a repeated registry call or
|
||||
Mongo insert can write twice.
|
||||
- **One `finally`, two actions**: cleanup, then the single terminal write that classifies the
|
||||
outcome once. No `except` block writes a verdict of its own.
|
||||
- **The workflow holds ids, never contents.** The decrypted bundle never crosses an activity
|
||||
boundary, and the password reaches this module only as an envelope it cannot open.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from datetime import timedelta
|
||||
|
||||
from temporalio import workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any
|
||||
|
||||
from sientia_do.temporal.policies import retry_policy
|
||||
from temporalio.exceptions import is_cancelled_exception
|
||||
|
||||
from laborious.activities.activities import Activities
|
||||
from laborious.activities.model_import import (
|
||||
IMPORT_STATUS_HINT_RETRY_POLICY,
|
||||
IMPORT_STATUS_RETRY_POLICY,
|
||||
)
|
||||
from laborious.utils.bundle.steps import ImportStep
|
||||
from laborious.utils.models.import_errors import (
|
||||
ImportInputError,
|
||||
describe_failure_chain,
|
||||
failure_location,
|
||||
)
|
||||
|
||||
# `start_to_close_timeout` is per attempt. The terminal write additionally gets a
|
||||
# `schedule_to_close_timeout` larger than its retry envelope (10 attempts, ≈151 s of backoff, ≈451 s
|
||||
# if every attempt also burns its timeout) — set it smaller and Temporal cuts the retries short,
|
||||
# which is the one configuration mistake that would silently undo design D20.
|
||||
STATUS_WRITE_TIMEOUT = timedelta(seconds=30)
|
||||
STATUS_WRITE_ENVELOPE = timedelta(seconds=900)
|
||||
# Download, decryption/extraction, experiment creation and artifact upload now run inside one
|
||||
# activity attempt (bug 8): the budget is the sum of what the four separate activities used to get,
|
||||
# since a retry of the fused activity redoes all four phases.
|
||||
PREPARE_AND_UPLOAD_TIMEOUT = (
|
||||
timedelta(seconds=1800)
|
||||
+ timedelta(seconds=1800)
|
||||
+ timedelta(seconds=300)
|
||||
+ timedelta(seconds=1800)
|
||||
)
|
||||
REGISTRATION_TIMEOUT = timedelta(seconds=600)
|
||||
DOCUMENT_TIMEOUT = timedelta(seconds=300)
|
||||
CLEANUP_TIMEOUT = timedelta(seconds=600)
|
||||
RETENTION_TIMEOUT = timedelta(seconds=300)
|
||||
|
||||
# Fields no import may carry. The project is a frontend concern: the model listing is the MongoDB
|
||||
# `models` document, no migrated database holds a project or model table, and nothing in this
|
||||
# repository reads a project-to-model link (design D22). An input naming one is rejected rather than
|
||||
# accepted and ignored, so no caller believes a project link was recorded. The same goes for a
|
||||
# relational target: there is none, and for a plaintext password, which does not belong in an event
|
||||
# history whatever else is true.
|
||||
REJECTED_INPUT_FIELDS = (
|
||||
'project_id',
|
||||
'project_name',
|
||||
'provisioning_target',
|
||||
'schema',
|
||||
'table_name',
|
||||
'status_table',
|
||||
'password',
|
||||
)
|
||||
|
||||
DIGEST_LENGTH = 64
|
||||
_HEX_DIGITS = frozenset('0123456789abcdef')
|
||||
|
||||
|
||||
@workflow.defn(name='import_model')
|
||||
class ImportModel:
|
||||
"""Turn an encrypted `.sientia` object in MinIO into a registered model, or into an ERROR."""
|
||||
|
||||
@workflow.run
|
||||
async def run(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Execute the import.
|
||||
|
||||
Args:
|
||||
input_data: the import request.
|
||||
Required keys:
|
||||
- import_run_id (int): primary key of the import log row the frontend inserted
|
||||
- object_key (str): the uploaded object's key
|
||||
- expected_digest (str): the SHA-256 the uploader computed for those bytes
|
||||
- password_envelope (str): the encrypted bundle password, never plaintext
|
||||
Optional keys:
|
||||
- bucket (str): defaults to the configured import bucket
|
||||
|
||||
Returns:
|
||||
dict: `import_run_id`, `run_id`, `model_name`, `version` and the model document's id.
|
||||
|
||||
Raises:
|
||||
ImportLogRowNotClaimableError: the record is not this workflow's — the run ends FAILED
|
||||
with that error type and nothing is written to any row.
|
||||
Exception: whatever failed. The record carries the failing step, its code and its
|
||||
sentence; this exception carries the technical truth.
|
||||
"""
|
||||
workflow_id = workflow.info().workflow_id
|
||||
import_run_id = self._require_import_run_id(input_data.get('import_run_id'))
|
||||
metadata: dict[str, Any] = {
|
||||
'model_name': '-',
|
||||
'model_id': str(import_run_id),
|
||||
'workflow_name': 'import_model',
|
||||
'schedule_name': '-',
|
||||
}
|
||||
base = {
|
||||
'metadata': metadata,
|
||||
'import_run_id': import_run_id,
|
||||
'workflow_id': workflow_id,
|
||||
}
|
||||
|
||||
# Class A. Outside the try/finally on purpose: a claim that matches nothing leaves this
|
||||
# workflow owning no row, so there is nothing to clean up and nothing to record.
|
||||
claim = await workflow.execute_activity_method(
|
||||
Activities.claim_import_status,
|
||||
base,
|
||||
retry_policy=IMPORT_STATUS_RETRY_POLICY,
|
||||
start_to_close_timeout=STATUS_WRITE_TIMEOUT,
|
||||
schedule_to_close_timeout=STATUS_WRITE_ENVELOPE,
|
||||
)
|
||||
|
||||
step: str = ImportStep.RECEIVED.value
|
||||
gate: int | None = None
|
||||
succeeded = False
|
||||
source: dict[str, Any] = {}
|
||||
result: dict[str, Any] = {'import_run_id': import_run_id}
|
||||
|
||||
try:
|
||||
self._validate_input(input_data, claim)
|
||||
|
||||
# Download, decryption/extraction, experiment creation and artifact upload run inside
|
||||
# one activity attempt (bug 8), so no later phase can depend on a file a different,
|
||||
# separately-scheduled activity wrote to a different worker's disk. One hint covers the
|
||||
# whole fused step; an exception it raises with no location of its own falls back to
|
||||
# this step, since that is what the workflow was attempting.
|
||||
step = ImportStep.DOWNLOAD.value
|
||||
await self._record_hint(
|
||||
base, Activities.record_import_status, {'step': step, 'source': None}
|
||||
)
|
||||
# Achado #9: the retention check carries no data this step consumes and already
|
||||
# swallows its own failure (`_ensure_retention`), so it runs alongside the fused
|
||||
# download+decrypt+provision activity instead of blocking it. Achado #8 fused
|
||||
# `download_import_bundle` into `prepare_and_upload_import_bundle` before this change
|
||||
# landed, so this gathers against that fused activity, not the original isolated one.
|
||||
_, prepared = await asyncio.gather(
|
||||
self._ensure_retention(base, input_data),
|
||||
workflow.execute_activity_method(
|
||||
Activities.prepare_and_upload_import_bundle,
|
||||
{
|
||||
**base,
|
||||
'bucket': input_data.get('bucket'),
|
||||
'object_key': input_data['object_key'],
|
||||
'expected_digest': input_data['expected_digest'],
|
||||
'password_envelope': input_data['password_envelope'],
|
||||
},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=PREPARE_AND_UPLOAD_TIMEOUT,
|
||||
),
|
||||
)
|
||||
source = prepared['metadata']
|
||||
parameters = prepared['parameters']
|
||||
model_name = source['model_name']
|
||||
metadata['model_name'] = model_name
|
||||
run_name = prepared['run_name']
|
||||
|
||||
# The names are known only once the fused activity returns; they are still recorded
|
||||
# before the two provisioning steps that follow (the registered version and the model
|
||||
# document), which is the guarantee `model-import-status-record` now makes.
|
||||
await self._record_hint(
|
||||
base,
|
||||
Activities.record_import_names,
|
||||
{
|
||||
'experiment_name': source['experiment_name'],
|
||||
'run_name': run_name,
|
||||
'source': source,
|
||||
'step': ImportStep.EXPERIMENT_CREATION.value,
|
||||
},
|
||||
)
|
||||
|
||||
step = ImportStep.REGISTRATION.value
|
||||
await self._record_hint(
|
||||
base,
|
||||
Activities.record_import_status,
|
||||
{'step': step, 'source': source or None},
|
||||
)
|
||||
version = await workflow.execute_activity_method(
|
||||
Activities.register_import_model_version,
|
||||
{**base, 'model_name': model_name, 'run_id': prepared['run_id']},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=REGISTRATION_TIMEOUT,
|
||||
)
|
||||
|
||||
step = ImportStep.MODEL_DOCUMENT.value
|
||||
await self._record_hint(
|
||||
base,
|
||||
Activities.record_import_status,
|
||||
{'step': step, 'source': source or None},
|
||||
)
|
||||
document = await workflow.execute_activity_method(
|
||||
Activities.write_import_model_document,
|
||||
{
|
||||
**base,
|
||||
'model_name': model_name,
|
||||
'target': parameters['target_variable'],
|
||||
},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=DOCUMENT_TIMEOUT,
|
||||
)
|
||||
|
||||
succeeded = True
|
||||
result = {
|
||||
'import_run_id': import_run_id,
|
||||
'run_id': prepared['run_id'],
|
||||
'model_name': model_name,
|
||||
'version': version['version'],
|
||||
'document_id': document.get('id'),
|
||||
}
|
||||
except BaseException as error: # classified once below, then re-raised untouched
|
||||
located_step, located_gate = failure_location(error)
|
||||
step = located_step or step
|
||||
gate = located_gate
|
||||
await self._report_failure(base, error, step=step, gate=gate)
|
||||
raise
|
||||
finally:
|
||||
cleanup = await self._cleanup(base)
|
||||
await self._record_terminal(
|
||||
base,
|
||||
succeeded=succeeded,
|
||||
step=step,
|
||||
gate=gate,
|
||||
source=source,
|
||||
cleanup_failed=not cleanup.get('cleaned', True),
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
# ------------------------------------------------------------------ validation
|
||||
|
||||
@staticmethod
|
||||
def _require_import_run_id(value: Any) -> int:
|
||||
"""The one check that runs before the claim, because the claim needs its result.
|
||||
|
||||
A missing or unusable id means there is no record to claim and none to write: the run ends
|
||||
FAILED with this error type, and the failure is visible in the workflow's terminal state, its
|
||||
error type, the error metric and a notification — everywhere except a row, which is the
|
||||
point.
|
||||
"""
|
||||
if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
|
||||
raise ImportInputError('import_run_id must be a positive integer')
|
||||
return value
|
||||
|
||||
@staticmethod
|
||||
def _validate_input(input_data: dict[str, Any], claim: dict[str, Any]) -> None:
|
||||
"""Validate the rest of the request, and check it against the record it just claimed.
|
||||
|
||||
This runs **after** the claim so its failure lands on the record rather than only in the
|
||||
logs: by now the row is this workflow's, so a bad request is reported to the person who made
|
||||
it. Everything here fails at step `received`.
|
||||
"""
|
||||
unknown = [name for name in REJECTED_INPUT_FIELDS if name in input_data]
|
||||
if unknown:
|
||||
raise ImportInputError(
|
||||
f'the import request carries fields this workflow does not accept: '
|
||||
f'{", ".join(sorted(unknown))}'
|
||||
)
|
||||
|
||||
for required in ('object_key', 'expected_digest', 'password_envelope'):
|
||||
if not input_data.get(required):
|
||||
raise ImportInputError(f'the import request is missing {required}')
|
||||
|
||||
digest = str(input_data['expected_digest']).lower()
|
||||
if len(digest) != DIGEST_LENGTH or not set(digest) <= _HEX_DIGITS:
|
||||
raise ImportInputError('expected_digest must be a 64-character hexadecimal SHA-256')
|
||||
|
||||
object_key = input_data['object_key']
|
||||
file_name = claim.get('file_name')
|
||||
if file_name and file_name != object_key:
|
||||
raise ImportInputError('the record names a different object than the request')
|
||||
|
||||
request_digest = claim.get('request_digest')
|
||||
if request_digest and request_digest != digest:
|
||||
raise ImportInputError('the record names a different digest than the request')
|
||||
|
||||
# --------------------------------------------------------------- status writes
|
||||
|
||||
async def _record_hint(
|
||||
self,
|
||||
base: dict[str, Any],
|
||||
activity: Callable[..., Awaitable[Any]],
|
||||
extra: dict[str, Any],
|
||||
) -> None:
|
||||
"""Record a status hint. A hint: its failure never fails the import."""
|
||||
try:
|
||||
await workflow.execute_activity_method(
|
||||
activity,
|
||||
{**base, **extra},
|
||||
retry_policy=IMPORT_STATUS_HINT_RETRY_POLICY,
|
||||
start_to_close_timeout=STATUS_WRITE_TIMEOUT,
|
||||
)
|
||||
except BaseException as error: # a lost hint must not fail a working import
|
||||
if is_cancelled_exception(error):
|
||||
raise
|
||||
await self._report_status_write_failure(base, error)
|
||||
|
||||
async def _record_terminal(
|
||||
self,
|
||||
base: dict[str, Any],
|
||||
*,
|
||||
succeeded: bool,
|
||||
step: str | None,
|
||||
gate: int | None,
|
||||
source: dict[str, Any],
|
||||
cleanup_failed: bool,
|
||||
) -> None:
|
||||
"""The single authoritative write, and the last thing this workflow does.
|
||||
|
||||
Its own failure is reported through the channel with the verdict it could not write, and it
|
||||
never turns a completed import into a failed one — nor masks the cause of a real failure.
|
||||
"""
|
||||
try:
|
||||
await workflow.execute_activity_method(
|
||||
Activities.record_import_terminal_status,
|
||||
{
|
||||
**base,
|
||||
'succeeded': succeeded,
|
||||
'step': step,
|
||||
'gate': gate,
|
||||
'source': source or None,
|
||||
'cleanup_failed': cleanup_failed,
|
||||
},
|
||||
retry_policy=IMPORT_STATUS_RETRY_POLICY,
|
||||
start_to_close_timeout=STATUS_WRITE_TIMEOUT,
|
||||
schedule_to_close_timeout=STATUS_WRITE_ENVELOPE,
|
||||
)
|
||||
except BaseException as error: # the outcome stands even when it cannot be written
|
||||
if is_cancelled_exception(error):
|
||||
raise
|
||||
await self._report_status_write_failure(
|
||||
base,
|
||||
error,
|
||||
verdict=f'succeeded={succeeded} step={step} gate={gate}',
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def _report_failure(
|
||||
base: dict[str, Any],
|
||||
error: BaseException,
|
||||
*,
|
||||
step: str | None,
|
||||
gate: int | None,
|
||||
) -> None:
|
||||
"""Route the failure that ended this import through the one failure channel.
|
||||
|
||||
Runs from the `except`, before cleanup and before the terminal write, so the log reads in the
|
||||
order things happened: what failed, then what was cleaned, then what was recorded. The step
|
||||
and the gate are the ones already located from the failure, so the log line and the row can
|
||||
never name different ones.
|
||||
|
||||
Its own failure is swallowed. The channel is the thing that logs, so a report that cannot be
|
||||
made has nowhere to say so, and the failure it was reporting is re-raised either way — the
|
||||
event history and the record both still carry it.
|
||||
"""
|
||||
error_type, detail = describe_failure_chain(error)
|
||||
try:
|
||||
await workflow.execute_activity_method(
|
||||
Activities.report_import_pipeline_failure,
|
||||
{**base, 'step': step, 'gate': gate, 'detail': detail, 'error_type': error_type},
|
||||
retry_policy=IMPORT_STATUS_HINT_RETRY_POLICY,
|
||||
start_to_close_timeout=STATUS_WRITE_TIMEOUT,
|
||||
)
|
||||
except BaseException as error: # the report of a report is where this stops
|
||||
if is_cancelled_exception(error):
|
||||
raise
|
||||
|
||||
async def _report_status_write_failure(
|
||||
self,
|
||||
base: dict[str, Any],
|
||||
error: BaseException,
|
||||
verdict: str | None = None,
|
||||
) -> None:
|
||||
"""Route a status write that could not be made through the one failure channel."""
|
||||
try:
|
||||
await workflow.execute_activity_method(
|
||||
Activities.report_import_status_write_failure,
|
||||
{**base, 'detail': str(error), 'verdict': verdict},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=STATUS_WRITE_TIMEOUT,
|
||||
)
|
||||
except BaseException as error: # the report of a report is where this stops
|
||||
if is_cancelled_exception(error):
|
||||
raise
|
||||
# Deliberately silent otherwise: the channel this was trying to reach is the thing that
|
||||
# logs, so there is nowhere left to say it, and an import that worked must not fail
|
||||
# over it.
|
||||
|
||||
# ----------------------------------------------------------- side-effect steps
|
||||
|
||||
@staticmethod
|
||||
async def _ensure_retention(base: dict[str, Any], input_data: dict[str, Any]) -> None:
|
||||
"""Ensure the bundle bucket's expiry rule, without making it a reason to fail an import.
|
||||
|
||||
The rule is applied by the importer rather than assumed from a Helm value, so the acceptance
|
||||
criterion is verifiable from the code that depends on it. But a bucket policy the importer
|
||||
cannot set is an operational problem, not a bad import: failing a valid model import because
|
||||
a lifecycle API refused would be the wrong trade, so the failure is dropped here and left
|
||||
where Temporal already records it — the activity's own failure in the event history.
|
||||
"""
|
||||
try:
|
||||
await workflow.execute_activity_method(
|
||||
Activities.ensure_import_bundle_retention,
|
||||
{**base, 'bucket': input_data.get('bucket')},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=RETENTION_TIMEOUT,
|
||||
)
|
||||
except BaseException as error: # retention is not a precondition of an import
|
||||
if is_cancelled_exception(error):
|
||||
raise
|
||||
# Dropped, not reported otherwise: the activity's own failure is already in the event
|
||||
# history, and a bucket policy the importer cannot set is an operational problem rather
|
||||
# than a bad import. Failing a valid model import over a lifecycle API would be the
|
||||
# wrong trade.
|
||||
|
||||
@staticmethod
|
||||
async def _cleanup(base: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Remove this import's temporary files. Runs on every path, before the terminal write."""
|
||||
try:
|
||||
return await workflow.execute_activity_method(
|
||||
Activities.cleanup_import_files,
|
||||
base,
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=CLEANUP_TIMEOUT,
|
||||
)
|
||||
except BaseException as error: # cleanup never replaces the failure that brought us here
|
||||
if is_cancelled_exception(error):
|
||||
raise
|
||||
return {'cleaned': False}
|
||||
@@ -32,6 +32,7 @@ class SimpleMetrics:
|
||||
model_config = input_data['model_config']
|
||||
target_name = model_config['target']
|
||||
|
||||
# Built from internal Temporal workflow config, not user input (bandit B608).
|
||||
query = f"""
|
||||
select p."timestamp", p.prediction, ld.value as "target"
|
||||
from "{input_data['schema']}"."{input_data['predictions_table_name']}" p
|
||||
@@ -45,7 +46,7 @@ class SimpleMetrics:
|
||||
p."timestamp" >= NOW() - INTERVAL '{interval_minutes} minutes'
|
||||
order by
|
||||
p."timestamp" desc;
|
||||
""" # nosec B608 - values come from internal Temporal workflow config, not user input
|
||||
""" # nosec
|
||||
|
||||
target_data = await workflow.execute_activity_method(
|
||||
Activities.load_custom_query,
|
||||
|
||||
@@ -75,6 +75,11 @@ class FormatAndExportPrediction:
|
||||
Required when path_flag is None
|
||||
- comment (str): Operational comment or error description
|
||||
Required when path_flag is not None
|
||||
- metadata_id (str | None): Id of the `prediction_metadata` Mongo document to
|
||||
reference for this row, already written by the caller; `None` when there's
|
||||
nothing to reference. Only meaningful when path_flag is not None — the
|
||||
normal (path_flag is None) path never has metadata to reference, so this
|
||||
key is irrelevant there.
|
||||
|
||||
Returns:
|
||||
None: The workflow completes successfully when all export operations finish
|
||||
@@ -154,6 +159,7 @@ class FormatAndExportPrediction:
|
||||
'model_id': input_data['model_id'],
|
||||
'prediction_confidence': prediction_confidence,
|
||||
'comment': input_data['comment'],
|
||||
'metadata_id': input_data.get('metadata_id'),
|
||||
},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=60),
|
||||
@@ -161,7 +167,7 @@ class FormatAndExportPrediction:
|
||||
|
||||
write_transformed_handler = None
|
||||
|
||||
opc_metrics = {}
|
||||
opc_metrics: dict[str, dict[str, float | None]] = {}
|
||||
|
||||
# write to pi web api
|
||||
if pi_web_api_output_config:
|
||||
|
||||
@@ -132,7 +132,9 @@ class PredictionProcess:
|
||||
'path_priority': input_data['path_priority'],
|
||||
}
|
||||
|
||||
path_flag, confidence, comment = await workflow.execute_local_activity_method(
|
||||
# `document` describes which filter(s) produced `path_flag` (`None` when nothing was
|
||||
# filtered); `path_flag_handler` is the one that decides whether and how to persist it.
|
||||
path_flag, confidence, comment, document = await workflow.execute_local_activity_method(
|
||||
Activities.input_gate,
|
||||
gate_input,
|
||||
retry_policy=retry_policy,
|
||||
@@ -141,7 +143,7 @@ class PredictionProcess:
|
||||
|
||||
# Handle path decision based on filter results
|
||||
if await self.path_flag_handler(
|
||||
data, path_flag, input_data, confidence, last_timestamp, comment
|
||||
data, path_flag, input_data, confidence, last_timestamp, comment, document
|
||||
):
|
||||
return
|
||||
|
||||
@@ -154,7 +156,7 @@ class PredictionProcess:
|
||||
)
|
||||
|
||||
# Validate MLFlow transform response
|
||||
path_flag, confidence, comment = await workflow.execute_local_activity_method(
|
||||
path_flag, confidence, comment, document = await workflow.execute_local_activity_method(
|
||||
Activities.mlflow_response_gate,
|
||||
{
|
||||
**metadata,
|
||||
@@ -169,11 +171,11 @@ class PredictionProcess:
|
||||
|
||||
# Handle path decision based on transform validation
|
||||
if await self.path_flag_handler(
|
||||
data, path_flag, input_data, confidence, last_timestamp, comment
|
||||
data, path_flag, input_data, confidence, last_timestamp, comment, document
|
||||
):
|
||||
return
|
||||
|
||||
path_flag, confidence, comment = await workflow.execute_local_activity_method(
|
||||
path_flag, confidence, comment, document = await workflow.execute_local_activity_method(
|
||||
Activities.mlflow_content_gate,
|
||||
{
|
||||
**metadata,
|
||||
@@ -187,7 +189,7 @@ class PredictionProcess:
|
||||
)
|
||||
|
||||
if await self.path_flag_handler(
|
||||
data, path_flag, input_data, confidence, last_timestamp, comment
|
||||
data, path_flag, input_data, confidence, last_timestamp, comment, document
|
||||
):
|
||||
return
|
||||
|
||||
@@ -204,7 +206,7 @@ class PredictionProcess:
|
||||
)
|
||||
|
||||
# Validate MLFlow prediction response
|
||||
path_flag, confidence, comment = await workflow.execute_local_activity_method(
|
||||
path_flag, confidence, comment, document = await workflow.execute_local_activity_method(
|
||||
Activities.mlflow_response_gate,
|
||||
{
|
||||
**metadata,
|
||||
@@ -219,7 +221,7 @@ class PredictionProcess:
|
||||
|
||||
# Handle path decision based on prediction validation
|
||||
if await self.path_flag_handler(
|
||||
data, path_flag, input_data, confidence, last_timestamp, comment
|
||||
data, path_flag, input_data, confidence, last_timestamp, comment, document
|
||||
):
|
||||
return
|
||||
|
||||
@@ -247,14 +249,61 @@ class PredictionProcess:
|
||||
},
|
||||
)
|
||||
|
||||
async def _write_prediction_metadata(
|
||||
self, metadata: dict[str, Any], document: dict[str, Any] | None
|
||||
) -> str | None:
|
||||
"""
|
||||
Write a gate's metadata document via the centralized write activity.
|
||||
|
||||
Shared by the CONTINUE and REPEAT branches of `path_flag_handler` — both persist a new
|
||||
`predictions` row (`REPEAT` reuses the previous row's values but still inserts a new
|
||||
one), so both need a `metadata_id` to reference it.
|
||||
|
||||
Args:
|
||||
metadata: Workflow metadata for contextual logging.
|
||||
document: Metadata produced by the gate that resolved this `path_flag`; `None` when
|
||||
nothing was filtered.
|
||||
|
||||
Return:
|
||||
The generated id, or `None` when there's nothing to write, including when the
|
||||
activity call itself fails (see `PredictionMetadata.write_prediction_metadata`).
|
||||
"""
|
||||
if not document:
|
||||
# Skip the remote activity call entirely rather than dispatching it just to have
|
||||
# `write_prediction_metadata` decide there's nothing to write (review finding,
|
||||
# 2026-10-01). All 3 gates now build a `document` whenever they return a non-`None`
|
||||
# `path_flag`, so this branch is a defensive fallback today, not the common case it
|
||||
# was when this was first written — kept in case a future caller passes a `path_flag`
|
||||
# without one.
|
||||
return None
|
||||
|
||||
try:
|
||||
return await workflow.execute_activity_method(
|
||||
Activities.write_prediction_metadata,
|
||||
{**metadata, 'documents': [document]},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(minutes=1),
|
||||
)
|
||||
except Exception as e:
|
||||
# Business rule 1: writing metadata must never block or fail the prediction write.
|
||||
# `write_prediction_metadata` already swallows Mongo-specific failures internally and
|
||||
# returns `None`, but that protection doesn't cover the activity call itself failing
|
||||
# to run at all (its `retry_policy` has `maximum_attempts=1` — no retry — so a
|
||||
# timeout, a momentarily unavailable worker, or any other dispatch failure raises
|
||||
# straight into the workflow). Catching it here, instead of just at the Mongo client
|
||||
# level, is what actually makes the rule hold (review finding, 2026-10-01).
|
||||
workflow.logger.warning(f'Failed to write prediction metadata, continuing anyway: {e}')
|
||||
return None
|
||||
|
||||
async def path_flag_handler(
|
||||
self,
|
||||
data: dict[str, Any],
|
||||
path_flag: str,
|
||||
path_flag: str | None,
|
||||
input_data: dict,
|
||||
confidence: int,
|
||||
last_timestamp: str,
|
||||
comment: str,
|
||||
document: dict[str, Any] | None = None,
|
||||
) -> bool:
|
||||
"""
|
||||
Handle path decisions based on filter results and confidence levels.
|
||||
@@ -280,6 +329,13 @@ class PredictionProcess:
|
||||
confidence: Confidence level from filter validation
|
||||
last_timestamp: Last processed timestamp
|
||||
comment: Additional information about the filter result
|
||||
document: Metadata describing the filter(s) that produced `path_flag`, as returned by
|
||||
the gate (`None` when nothing was filtered, or when the gate that produced this
|
||||
`path_flag` doesn't build one yet). Written via `write_prediction_metadata` for
|
||||
`CONTINUE` and `REPEAT` alike — both persist a `predictions` row (`REPEAT` reuses
|
||||
the previous row's values but still inserts a new one), so both can reference the
|
||||
resulting `metadata_id`; `STOP` never reaches a write at all, so its `document` is
|
||||
simply never used here.
|
||||
|
||||
Returns:
|
||||
bool: True if processing should stop, False to continue
|
||||
@@ -304,7 +360,10 @@ class PredictionProcess:
|
||||
# Stop processing and exit workflow
|
||||
return True
|
||||
elif path_flag == 'REPEAT':
|
||||
# Repeat last prediction if available
|
||||
metadata_id = await self._write_prediction_metadata(metadata, document)
|
||||
# Repeat last prediction, copying the previous row's values forward with a new
|
||||
# timestamp. `repeat_last_prediction` now lives in `Storage` (laborious), not the
|
||||
# shared library — its INSERT...SELECT references `metadata_id` directly.
|
||||
await workflow.execute_activity_method(
|
||||
Activities.repeat_last_prediction,
|
||||
{
|
||||
@@ -313,12 +372,14 @@ class PredictionProcess:
|
||||
'table_name': table_name,
|
||||
'model': model_id,
|
||||
'last_timestamp': last_timestamp,
|
||||
'metadata_id': metadata_id,
|
||||
},
|
||||
retry_policy=retry_policy,
|
||||
start_to_close_timeout=timedelta(minutes=1),
|
||||
)
|
||||
return True
|
||||
elif path_flag == 'CONTINUE':
|
||||
metadata_id = await self._write_prediction_metadata(metadata, document)
|
||||
# call write workflow
|
||||
await workflow.execute_child_workflow(
|
||||
'subworkflow.format_and_export_prediction',
|
||||
@@ -335,6 +396,7 @@ class PredictionProcess:
|
||||
'table_name': table_name,
|
||||
'transform_table_name': transform_table_name,
|
||||
'comment': comment,
|
||||
'metadata_id': metadata_id,
|
||||
'opc_output_config': input_data['opc_output_config'],
|
||||
'pi_web_api_output_config': input_data['pi_web_api_output_config'],
|
||||
'prediction_store_policy': input_data['prediction_store_policy'],
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "laborious"
|
||||
version = "0.9.1"
|
||||
version = "0.9.0-rc4"
|
||||
description = "Sientia DataOps Laborious - ML Model Orchestration System"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
@@ -149,10 +149,6 @@ markers = [
|
||||
"integration: marks tests as integration tests",
|
||||
"unit: marks tests as unit tests",
|
||||
"opc: marks tests that use the in-process OPC UA server (OpcRepository E2E)",
|
||||
# The .sientia import e2e suite. Outside `testpaths`, so `validate --only-tests` and the
|
||||
# coverage gate never collect it; run it on request with `pytest e2e -m import_e2e`.
|
||||
"import_e2e: marks the model-import end-to-end tests (needs Docker and an MLflow server)",
|
||||
"slow: marks tests that run on a real clock and take tens of seconds",
|
||||
]
|
||||
|
||||
[tool.coverage.run]
|
||||
|
||||
@@ -22,4 +22,3 @@ testcontainers[postgres,minio,mongodb]>=4.15.0 # PostgreSQL, MinIO and MongoDB c
|
||||
ipython>=8.12.0 # Enhanced Python shell
|
||||
ipdb>=0.13.13 # IPython debugger
|
||||
ipykernel==6.30.1 # IPython kernel for Jupyter notebooks
|
||||
python-dotenv>=1.0.0 # .env loading for scripts/import_model_manual_run.py
|
||||
@@ -1,21 +0,0 @@
|
||||
temporalio
|
||||
psycopg2-binary
|
||||
sqlalchemy<2.1
|
||||
asyncua==1.0.6
|
||||
redis
|
||||
sientia_do>=1.12.2
|
||||
mlflow==2.10.1
|
||||
prometheus-client
|
||||
botocore
|
||||
boto3
|
||||
s3fs
|
||||
pyarrow
|
||||
kaleido
|
||||
hyperopt
|
||||
shap
|
||||
pycurl
|
||||
scipy<1.14.0
|
||||
scikit-learn==1.4.2
|
||||
pycaret==3.3.2
|
||||
joblib
|
||||
pynacl==1.5.0
|
||||
@@ -1,17 +1,18 @@
|
||||
temporalio
|
||||
psycopg2-binary
|
||||
# SQLAlchemy 2.1 makes psycopg (v3) the default driver for postgresql:// URLs; sientia_do
|
||||
# builds that URL and ships only psycopg2-binary, so 2.1 fails with "No module named psycopg".
|
||||
sqlalchemy<2.1
|
||||
asyncua==1.0.6
|
||||
redis
|
||||
git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.12.2
|
||||
git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.41.1
|
||||
git+ssh://git@github.com/Aignosi/sientia-model-library.git@0.10.0
|
||||
prometheus-client
|
||||
# fsspec/s3fs/gcsfs pinned as a set: nannyml (via the mlops library) pulls both s3fs and
|
||||
# gcsfs, and each pins fsspec exactly, so a free `s3fs` makes the resolver fail against
|
||||
# evidently's fsspec>=2024.2.0. Versions match the library-distribution-controller freeze.
|
||||
# fsspec/s3fs pinned as a pair: s3fs pins fsspec exactly, so a free `s3fs` can make the
|
||||
# resolver fail against evidently's fsspec>=2024.2.0. Versions match the
|
||||
# library-distribution-controller freeze.
|
||||
fsspec==2026.7.0
|
||||
s3fs==2026.7.0
|
||||
gcsfs==2026.8.0
|
||||
# aiobotocore 3.9.0 (pulled by s3fs) requires botocore<1.43.57,>=1.43.3.
|
||||
botocore==1.43.56
|
||||
boto3==1.43.56
|
||||
@@ -21,6 +22,4 @@ hyperopt
|
||||
shap
|
||||
pycurl
|
||||
scipy<1.14.0
|
||||
scikit-learn==1.4.2
|
||||
pycaret==3.3.2
|
||||
joblib
|
||||
scikit-learn==1.5.2
|
||||
|
||||
@@ -1,17 +1,18 @@
|
||||
temporalio
|
||||
psycopg2-binary
|
||||
# SQLAlchemy 2.1 makes psycopg (v3) the default driver for postgresql:// URLs; sientia_do
|
||||
# builds that URL and ships only psycopg2-binary, so 2.1 fails with "No module named psycopg".
|
||||
sqlalchemy<2.1
|
||||
asyncua==1.0.6
|
||||
redis
|
||||
sientia_do>=1.12.2
|
||||
sientia>=0.41.1
|
||||
sientia_model>=0.8.2
|
||||
prometheus-client
|
||||
# fsspec/s3fs/gcsfs pinned as a set: nannyml (via sientia) pulls both s3fs and gcsfs, and
|
||||
# each pins fsspec exactly, so a free `s3fs` makes the resolver fail against evidently's
|
||||
# fsspec>=2024.2.0. Versions match the library-distribution-controller freeze.
|
||||
# fsspec/s3fs pinned as a pair: s3fs pins fsspec exactly, so a free `s3fs` can make the
|
||||
# resolver fail against evidently's fsspec>=2024.2.0. Versions match the
|
||||
# library-distribution-controller freeze.
|
||||
fsspec==2026.7.0
|
||||
s3fs==2026.7.0
|
||||
gcsfs==2026.8.0
|
||||
# aiobotocore 3.9.0 (pulled by s3fs) requires botocore<1.43.57,>=1.43.3.
|
||||
botocore==1.43.56
|
||||
boto3==1.43.56
|
||||
@@ -21,13 +22,4 @@ hyperopt
|
||||
shap
|
||||
pycurl
|
||||
scipy<1.14.0
|
||||
scikit-learn==1.4.2
|
||||
pycaret==3.3.2
|
||||
joblib
|
||||
# .sientia bundle import (Argon2id + XChaCha20-Poly1305 secretstream): same pin as the
|
||||
# producer in sientia-projects-templates, so both sides open the same bytes.
|
||||
pynacl==1.5.0
|
||||
# AES-256-GCM password envelope; already present transitively, declared explicitly.
|
||||
cryptography
|
||||
# MongoDB models collection (the model listing); transitive via sientia_do, declared explicitly.
|
||||
pymongo
|
||||
scikit-learn
|
||||
|
||||
@@ -1,586 +0,0 @@
|
||||
"""Run one `.sientia` model import by hand, standing in for the frontend that does not exist yet.
|
||||
|
||||
**What this is.** The Angular screen plus the Spring Boot BFF will, between them, pick a `.sientia`
|
||||
file, hash it, wrap the bundle password for the worker, upload the object and insert one `PENDING`
|
||||
row into `public.experiment_run` before starting `import_model`. None of that exists yet. This
|
||||
script does exactly those four things, against whatever environment `.env` points at, so the import
|
||||
path can be exercised for real and so whoever writes the frontend has a working reference for the
|
||||
wire shapes (the digest, the envelope, the object key, the row).
|
||||
|
||||
**What it writes.** Two things, and it can undo both: one object under
|
||||
`{IMPORT_BUNDLE_BUCKET}/{IMPORT_BUNDLE_PREFIX}` in MinIO, and one row in `public.experiment_run` in
|
||||
the BFF Postgres database. Everything else it does is a read. If the workflow runs, *it* creates the
|
||||
MLflow experiment, run and registered version and the MongoDB model document — cell 10 can remove
|
||||
those too, behind its own flag.
|
||||
|
||||
**The safety rule.** Every target comes from the environment, and this file contains no host, user,
|
||||
password or key literal. A `.env` inherited from a cluster deployment can perfectly well name a
|
||||
production database and a production bucket, so cell 5 and cell 6 — the only cells that write — are
|
||||
gated on cell 2, where the operator has to restate the Postgres host that cell 1 printed. This is a
|
||||
development tool. It is not for production, and the guard is a mitigation, not a guarantee.
|
||||
|
||||
**Cells.** Run them one at a time in the VS Code interactive window, or run the module top to
|
||||
bottom. Cell order is the contract:
|
||||
|
||||
1 resolve every target and print it 6 insert the PENDING row -> import_run_id
|
||||
2 confirmation guard (first write is after) 7 print the workflow input
|
||||
3 pick the bundle, compute expected_digest 8 optionally start the workflow and wait
|
||||
4 build the password envelope 9 read the outcome back
|
||||
5 upload the object, verify the stored bytes 10 undo
|
||||
|
||||
**Boundaries.** This script imports nothing from `laborious.activities` or `laborious.workflows` —
|
||||
it stands in for the frontend, so it must not borrow the consumer's code. The one exception is
|
||||
`laborious.utils.bundle.format`, which is pure `struct`/`hashlib` over bytes and lets cell 3 print
|
||||
the header of the file it is about to upload. The envelope is built here in four visible lines
|
||||
rather than imported, because those four lines are what the frontend has to transliterate to
|
||||
WebCrypto; `laborious/utils/import_password.py` is the decrypt counterpart.
|
||||
"""
|
||||
|
||||
# %%
|
||||
# Cell 1 — resolve every target from the environment and print what we are pointed at.
|
||||
# Reads only. No secret is printed here or anywhere else in this script.
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import secrets
|
||||
import sys
|
||||
from datetime import UTC, datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import mlflow
|
||||
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
|
||||
from dotenv import load_dotenv
|
||||
from minio import Minio
|
||||
from pymongo import MongoClient
|
||||
from sqlalchemy import create_engine, text
|
||||
from sqlalchemy.engine import URL
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parent.parent
|
||||
if str(REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
load_dotenv(REPO_ROOT / '.env')
|
||||
|
||||
# Below the `sys.path` insert on purpose: `laborious` is only importable once the repository root is
|
||||
# on the path, which is what those three lines above just arranged. Hoisting this to the top of the
|
||||
# cell is a `ModuleNotFoundError` when the script runs from anywhere but the repository root.
|
||||
from laborious.utils.bundle.format import DIGEST_OFFSET, HEADER_SIZE, parse_header # noqa: E402
|
||||
|
||||
|
||||
def env(name: str, default: str = '') -> str:
|
||||
"""The environment variable, stripped, or `default`."""
|
||||
return os.getenv(name, default).strip()
|
||||
|
||||
|
||||
def env_required(name: str) -> str:
|
||||
"""The environment variable, or a `RuntimeError` naming it."""
|
||||
value = env(name)
|
||||
if not value:
|
||||
raise RuntimeError(f'{name} is not set — see .env.example')
|
||||
return value
|
||||
|
||||
|
||||
def env_secret_required(name: str) -> str:
|
||||
"""The environment variable *verbatim*, or a `RuntimeError` naming it.
|
||||
|
||||
Unlike `env_required` this does not strip, and the difference is not cosmetic. A key and a
|
||||
password are byte strings whose surrounding whitespace is part of the secret. The worker reads
|
||||
the same variables with a plain `os.getenv`, so stripping here would seal the envelope under
|
||||
bytes the worker never derives, and the import would fail at a gate with nothing in the message
|
||||
to say why.
|
||||
"""
|
||||
value = os.getenv(name, '')
|
||||
if not value:
|
||||
raise RuntimeError(f'{name} is not set — see .env.example')
|
||||
return value
|
||||
|
||||
|
||||
def _line(label: str, value: object) -> str:
|
||||
"""One aligned `label value` line, for output that stays readable in a terminal."""
|
||||
return f' {label:<24} {value}'
|
||||
|
||||
|
||||
# Every connection target below is *required* rather than defaulted. The worker's own config falls
|
||||
# back to localhost, which is right for a service that must boot; it is wrong for a script that
|
||||
# writes, because a silent fallback is a write to somewhere nobody read out loud.
|
||||
#
|
||||
# The import log database: the POSTGRES_* server from .env, with only the database name of its own,
|
||||
# because the log lives in the Spring Boot BFF's database rather than in `sientia`. The schema and
|
||||
# table are constants of the worker's code, never configuration.
|
||||
#
|
||||
# `build_import_status_config` also honours IMPORT_STATUS_DB_{HOST,PORT,USER,PASSWORD} as overrides,
|
||||
# for a deployment whose log sits on another server. This script does not: it is a development tool
|
||||
# pointed at a real environment, and the override exists for the e2e suite's two containers.
|
||||
STATUS_DB_HOST = env_required('POSTGRES_HOST')
|
||||
STATUS_DB_PORT = int(env_required('POSTGRES_PORT'))
|
||||
STATUS_DB_NAME = env_required('IMPORT_STATUS_DB_NAME')
|
||||
STATUS_DB_USER = env_required('POSTGRES_USER')
|
||||
STATUS_DB_PASSWORD = os.getenv('POSTGRES_PASSWORD', '')
|
||||
STATUS_TABLE = 'public.experiment_run'
|
||||
|
||||
# Where the frontend uploads the bundle. The prefix keeps the worker's own default, which is a
|
||||
# format decision rather than a target.
|
||||
MINIO_ENDPOINT_URL = env_required('MINIO_ENDPOINT_URL')
|
||||
MINIO_ACCESS_KEY = env_required('MINIO_ACCESS_KEY')
|
||||
MINIO_SECRET_KEY = env_required('MINIO_SECRET_KEY')
|
||||
BUNDLE_BUCKET = env_required('IMPORT_BUNDLE_BUCKET')
|
||||
_raw_prefix = env('IMPORT_BUNDLE_PREFIX', 'imported_models/')
|
||||
BUNDLE_PREFIX = _raw_prefix if _raw_prefix.endswith('/') else f'{_raw_prefix}/'
|
||||
|
||||
# Read-back targets. This script never writes to these; the workflow does.
|
||||
MONGODB_URL = env_required('MONGODB_URL')
|
||||
MONGODB_USERNAME = env('MONGODB_USERNAME')
|
||||
MONGODB_PASSWORD = os.getenv('MONGODB_PASSWORD', '')
|
||||
MONGODB_DATABASE = env_required('MONGODB_DATABASE')
|
||||
MODELS_COLLECTION = env('IMPORT_MODELS_COLLECTION', 'models')
|
||||
|
||||
MLFLOW_HOST = env_required('MLFLOW_HOST')
|
||||
MLFLOW_PORT = env_required('MLFLOW_PORT')
|
||||
MLFLOW_TRACKING_URI = f'{MLFLOW_HOST}:{MLFLOW_PORT}'
|
||||
|
||||
TEMPORAL_HOST = env_required('TEMPORAL_HOST')
|
||||
TEMPORAL_NAMESPACE = env_required('TEMPORAL_NAMESPACE')
|
||||
RUNTIME = env_required('RUNTIME')
|
||||
IMPORT_TASK_QUEUE = 'import_model-legacy-queue'
|
||||
|
||||
# This script's own inputs.
|
||||
BUNDLE_PATH = Path(
|
||||
env('IMPORT_SCRIPT_BUNDLE_PATH', 'tests/fixtures/bundle/golden_model_v3.sientia')
|
||||
)
|
||||
if not BUNDLE_PATH.is_absolute():
|
||||
BUNDLE_PATH = REPO_ROOT / BUNDLE_PATH
|
||||
IMPORT_USERNAME = env_required('IMPORT_SCRIPT_USERNAME')
|
||||
|
||||
print('resolved targets — nothing has been written yet')
|
||||
print(_line('status postgres', f'{STATUS_DB_HOST}:{STATUS_DB_PORT}/{STATUS_DB_NAME}'))
|
||||
print(_line('status user', STATUS_DB_USER))
|
||||
print(_line('status table', STATUS_TABLE))
|
||||
print(_line('minio endpoint', MINIO_ENDPOINT_URL))
|
||||
print(_line('minio bucket/prefix', f'{BUNDLE_BUCKET}/{BUNDLE_PREFIX}'))
|
||||
print(_line('mongo', f'{MONGODB_URL}/{MONGODB_DATABASE}.{MODELS_COLLECTION}'))
|
||||
print(_line('mlflow', MLFLOW_TRACKING_URI))
|
||||
print(_line('temporal', f'{TEMPORAL_HOST} ns={TEMPORAL_NAMESPACE}'))
|
||||
print(_line('task queue', IMPORT_TASK_QUEUE))
|
||||
print(_line('bundle', BUNDLE_PATH))
|
||||
print(_line('username', IMPORT_USERNAME))
|
||||
print()
|
||||
print(f'to continue, set the cell 2 confirmation to: {STATUS_DB_HOST}')
|
||||
|
||||
# %%
|
||||
# Cell 2 — the confirmation guard. Cells 5 and 6 write; this is the last cell before them.
|
||||
#
|
||||
# Restate the Postgres host cell 1 printed. Either set IMPORT_SCRIPT_CONFIRM_TARGET in .env, or
|
||||
# replace the right-hand side below with the value cell 1 told you. It is a read-then-restate on
|
||||
# purpose: a boolean flag can be flipped without ever looking at what it is pointed at.
|
||||
CONFIRM_TARGET = 'localhost'
|
||||
|
||||
if CONFIRM_TARGET != STATUS_DB_HOST:
|
||||
raise RuntimeError(
|
||||
'refusing to write: the confirmation does not match the resolved status database host. '
|
||||
f'confirmation={CONFIRM_TARGET!r}, resolved POSTGRES_HOST={STATUS_DB_HOST!r}. '
|
||||
'Read the targets cell 1 printed, then restate the host here.'
|
||||
)
|
||||
|
||||
print(f'target confirmed: {STATUS_DB_HOST} — the write cells may run')
|
||||
|
||||
# %%
|
||||
# Cell 3 — pick the bundle, compute the digest the frontend must send, print its header.
|
||||
#
|
||||
# `expected_digest` is the SHA-256 of the exact bytes that get uploaded, computed the way the
|
||||
# frontend computes it: over the file, not over anything derived from it. The header fields printed
|
||||
# here are plaintext at the front of the file — reading them decrypts nothing.
|
||||
|
||||
|
||||
bundle_bytes = BUNDLE_PATH.read_bytes()
|
||||
expected_digest = hashlib.sha256(bundle_bytes).hexdigest()
|
||||
object_name = BUNDLE_PATH.name
|
||||
object_key = f'{BUNDLE_PREFIX}{object_name}'
|
||||
|
||||
header = parse_header(bundle_bytes[:HEADER_SIZE])
|
||||
|
||||
print(_line('file', BUNDLE_PATH))
|
||||
print(_line('size', f'{len(bundle_bytes)} bytes'))
|
||||
print(_line('expected_digest', expected_digest))
|
||||
print(_line('object key', f'{BUNDLE_BUCKET}/{object_key}'))
|
||||
print(_line('format version', header.format_version))
|
||||
print(_line('aead / kdf id', f'{header.aead_id} / {header.kdf_id}'))
|
||||
print(_line('kdf memlimit', f'{header.kdf_memlimit_bytes} bytes'))
|
||||
print(_line('kdf opslimit', header.kdf_opslimit))
|
||||
print(_line('kdf parallelism', header.kdf_parallelism))
|
||||
print(_line('salt', header.salt.hex()))
|
||||
print(_line('ciphertext digest', header.ciphertext_digest.hex()))
|
||||
print(_line('chunk stream from', f'offset {DIGEST_OFFSET}'))
|
||||
|
||||
# %%
|
||||
# Cell 4 — build the password envelope, exactly as the frontend must build it.
|
||||
#
|
||||
# Wire shape: base64( nonce(12) || ciphertext || GCM tag(16) ) under AES-256-GCM with the 32-byte
|
||||
# key that `IMPORT_PASSWORD_KEY` carries base64-encoded. The four lines below are the whole of it,
|
||||
# written out rather than imported so the Angular side can transliterate them to WebCrypto:
|
||||
#
|
||||
# const k = await crypto.subtle.importKey('raw', keyBytes, 'AES-GCM', false, ['encrypt']);
|
||||
# const nonce = crypto.getRandomValues(new Uint8Array(12));
|
||||
# const sealed = await crypto.subtle.encrypt({name: 'AES-GCM', iv: nonce}, k, passwordBytes);
|
||||
# const envelope = base64(concat(nonce, new Uint8Array(sealed)));
|
||||
#
|
||||
# WebCrypto appends the GCM tag to the ciphertext, and so does `cryptography`'s AESGCM, so the two
|
||||
# produce byte-identical payloads. The decrypt counterpart the worker runs is
|
||||
# `laborious/utils/import_password.py` (`decrypt_password_envelope`).
|
||||
#
|
||||
# `IMPORT_PASSWORD_KEY` must be base64 of exactly 32 bytes — `openssl rand -base64 32`, or the
|
||||
# value the deployment's Kubernetes Secret carries. Anything else fails here, which is the point:
|
||||
# the worker resolves the same variable the same way, so a value this cell cannot decode is a value
|
||||
# the worker could not have decoded either.
|
||||
#
|
||||
# The bundle password is read verbatim — `env_secret_required`, not `env_required` — because a
|
||||
# trailing space is part of the password, and stripping it here would seal an envelope whose
|
||||
# plaintext the bundle reader rejects at gate 3.
|
||||
#
|
||||
# Neither the key nor the password nor the plaintext is printed — only the envelope's length.
|
||||
|
||||
|
||||
envelope_key = base64.b64decode(env_required('IMPORT_PASSWORD_KEY'), validate=True)
|
||||
nonce = secrets.token_bytes(12)
|
||||
sealed = AESGCM(envelope_key).encrypt(
|
||||
nonce, env_secret_required('IMPORT_SCRIPT_BUNDLE_PASSWORD').encode('utf-8'), None
|
||||
)
|
||||
password_envelope = base64.b64encode(nonce + sealed).decode('ascii')
|
||||
|
||||
print(_line('envelope key', f'{len(envelope_key)} bytes (not printed)'))
|
||||
print(_line('envelope', f'{len(password_envelope)} base64 characters (not printed)'))
|
||||
|
||||
# %%
|
||||
# Cell 5 — FIRST WRITE. Upload the object, then prove the stored bytes are the ones we hashed.
|
||||
#
|
||||
# The re-read is the point: `expected_digest` is a promise about what is in the bucket, and the
|
||||
# worker's gate 1 holds the upload to it. If the digest of the stored object differs, the frontend
|
||||
# contract is already broken and there is no reason to insert a row.
|
||||
|
||||
|
||||
_endpoint = urlparse(MINIO_ENDPOINT_URL)
|
||||
minio_client = Minio(
|
||||
_endpoint.netloc or MINIO_ENDPOINT_URL,
|
||||
access_key=MINIO_ACCESS_KEY,
|
||||
secret_key=MINIO_SECRET_KEY,
|
||||
secure=_endpoint.scheme == 'https',
|
||||
)
|
||||
|
||||
minio_client.put_object(
|
||||
BUNDLE_BUCKET,
|
||||
object_key,
|
||||
io.BytesIO(bundle_bytes),
|
||||
length=len(bundle_bytes),
|
||||
content_type='application/octet-stream',
|
||||
)
|
||||
|
||||
stat = minio_client.stat_object(BUNDLE_BUCKET, object_key)
|
||||
_response = minio_client.get_object(BUNDLE_BUCKET, object_key)
|
||||
try:
|
||||
stored_bytes = _response.read()
|
||||
finally:
|
||||
_response.close()
|
||||
_response.release_conn()
|
||||
|
||||
stored_digest = hashlib.sha256(stored_bytes).hexdigest()
|
||||
if stored_digest != expected_digest:
|
||||
raise RuntimeError(
|
||||
f'the stored object does not hash to the digest we computed: stored {stored_digest}, '
|
||||
f'expected {expected_digest} — do not insert a row for this upload'
|
||||
)
|
||||
|
||||
print(_line('uploaded', f'{BUNDLE_BUCKET}/{object_key}'))
|
||||
print(_line('reported size', f'{stat.size} bytes'))
|
||||
print(_line('etag', stat.etag))
|
||||
print(_line('re-read digest', f'{stored_digest} (matches expected_digest)'))
|
||||
|
||||
# %%
|
||||
# Cell 6 — SECOND WRITE. Insert the import log row: the BFF's half of the contract.
|
||||
#
|
||||
# `URL.create` rather than a URL string, so a password containing `@ : /` survives. `experiment_name`
|
||||
# is a placeholder: the column is NOT NULL and the frontend cannot see inside an encrypted bundle,
|
||||
# so the worker overwrites it with the bundle's own experiment name as soon as the bundle is open
|
||||
# (`record_import_names`). `file_name` is the object key, because the worker compares the two and
|
||||
# refuses the import at the `received` step when they disagree.
|
||||
|
||||
|
||||
status_engine = create_engine(
|
||||
URL.create(
|
||||
'postgresql',
|
||||
username=STATUS_DB_USER,
|
||||
password=STATUS_DB_PASSWORD,
|
||||
host=STATUS_DB_HOST,
|
||||
port=STATUS_DB_PORT,
|
||||
database=STATUS_DB_NAME,
|
||||
),
|
||||
pool_pre_ping=True,
|
||||
)
|
||||
|
||||
request_data = {
|
||||
'expected_digest': expected_digest,
|
||||
'bucket': BUNDLE_BUCKET,
|
||||
'object_key': object_key,
|
||||
'file_name': object_name,
|
||||
'source': 'scripts/import_model_manual_run.py',
|
||||
}
|
||||
|
||||
_now = datetime.now(UTC).replace(tzinfo=None)
|
||||
_placeholder_experiment = f'import-{Path(object_name).stem}'[:50]
|
||||
|
||||
INSERT_ROW_SQL = text(
|
||||
'INSERT INTO public.experiment_run '
|
||||
'(experiment_name, username, status, created_at, updated_at, '
|
||||
' bucket_name, file_name, request_data, run_type) '
|
||||
'VALUES (:experiment_name, :username, :status, :created_at, :updated_at, '
|
||||
' :bucket_name, :file_name, CAST(:request_data AS JSONB), :run_type) '
|
||||
'RETURNING id'
|
||||
)
|
||||
|
||||
with status_engine.begin() as connection:
|
||||
import_run_id = connection.execute(
|
||||
INSERT_ROW_SQL,
|
||||
{
|
||||
'experiment_name': _placeholder_experiment,
|
||||
'username': IMPORT_USERNAME,
|
||||
'status': 'PENDING',
|
||||
'created_at': _now,
|
||||
'updated_at': _now,
|
||||
'bucket_name': BUNDLE_BUCKET,
|
||||
'file_name': object_key,
|
||||
'request_data': json.dumps(request_data),
|
||||
'run_type': 'IMPORT',
|
||||
},
|
||||
).scalar_one()
|
||||
|
||||
print(_line('import_run_id', import_run_id))
|
||||
print(_line('status / run_type', 'PENDING / IMPORT'))
|
||||
print(_line('experiment_name', f'{_placeholder_experiment} (the worker overwrites this)'))
|
||||
print(_line('file_name', object_key))
|
||||
|
||||
# %%
|
||||
# Cell 7 — the workflow input, ready to paste into the Temporal UI.
|
||||
#
|
||||
# These five keys are the whole contract. `bucket` is optional — the worker falls back to
|
||||
# IMPORT_BUNDLE_BUCKET — and is sent explicitly so the printed input is self-describing. No model
|
||||
# name is sent: the worker takes it from the bundle's own metadata, verbatim.
|
||||
workflow_input = {
|
||||
'import_run_id': import_run_id,
|
||||
'bucket': BUNDLE_BUCKET,
|
||||
'object_key': object_key,
|
||||
'expected_digest': expected_digest,
|
||||
'password_envelope': password_envelope,
|
||||
}
|
||||
|
||||
print(f'workflow: import_model task queue: {IMPORT_TASK_QUEUE}')
|
||||
print(json.dumps(workflow_input, indent=2))
|
||||
|
||||
# %%
|
||||
# Cell 8 — optionally start the workflow and wait for it.
|
||||
#
|
||||
# Skipping this cell is normal: the operator may prefer to start the run from the Temporal UI with
|
||||
# the JSON cell 7 printed. Cell 9 works either way, because it looks the outcome up by row id.
|
||||
#
|
||||
# The coroutine runs on a loop of its own in a worker thread rather than through `asyncio.run`, so
|
||||
# this cell behaves the same in the interactive window (which already has a running loop) as it
|
||||
# does when the module is executed top to bottom (which does not).
|
||||
START_THE_WORKFLOW = True
|
||||
|
||||
if START_THE_WORKFLOW:
|
||||
import asyncio
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
from temporalio.client import Client
|
||||
|
||||
async def start_and_wait() -> Any:
|
||||
"""Start `import_model` on the import queue and return its result."""
|
||||
client = await Client.connect(TEMPORAL_HOST, namespace=TEMPORAL_NAMESPACE)
|
||||
handle = await client.start_workflow(
|
||||
'import_model',
|
||||
workflow_input,
|
||||
id=f'manual-import-{import_run_id}',
|
||||
task_queue=IMPORT_TASK_QUEUE,
|
||||
)
|
||||
print(_line('workflow id', handle.id))
|
||||
print(_line('run id', handle.result_run_id))
|
||||
return await handle.result()
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as _pool:
|
||||
workflow_result = _pool.submit(asyncio.run, start_and_wait()).result()
|
||||
print(json.dumps(workflow_result, indent=2, default=str))
|
||||
else:
|
||||
print('not started — paste the input from cell 7 into the Temporal UI, or set')
|
||||
print(f'START_THE_WORKFLOW = True and re-run this cell. Task queue: {IMPORT_TASK_QUEUE}')
|
||||
|
||||
# %%
|
||||
# Cell 9 — read the outcome back from all three places it lands.
|
||||
#
|
||||
# The row is authoritative for the verdict; MLflow and MongoDB are what the verdict is about. Each
|
||||
# block is labelled, so the output can be pasted into a ticket as it is.
|
||||
print('--- import log row ---')
|
||||
READ_ROW_SQL = text(
|
||||
'SELECT id, experiment_name, run_name, status, error_message, run_type, '
|
||||
' bucket_name, file_name, created_at, updated_at, '
|
||||
' request_data, orchestrator_response_data '
|
||||
'FROM public.experiment_run WHERE id = :id'
|
||||
)
|
||||
|
||||
with status_engine.begin() as connection:
|
||||
_result = connection.execute(READ_ROW_SQL, {'id': import_run_id})
|
||||
row = _result.mappings().one_or_none()
|
||||
|
||||
if row is None:
|
||||
print(f' no row with id {import_run_id}')
|
||||
else:
|
||||
for _column, _value in row.items():
|
||||
print(_line(_column, _value))
|
||||
|
||||
print()
|
||||
print('--- mlflow ---')
|
||||
|
||||
os.environ['MLFLOW_TRACKING_USERNAME'] = env('MLFLOW_USERNAME')
|
||||
os.environ['MLFLOW_TRACKING_PASSWORD'] = os.getenv('MLFLOW_PASSWORD', '')
|
||||
mlflow.set_tracking_uri(MLFLOW_TRACKING_URI)
|
||||
mlflow_client = mlflow.tracking.MlflowClient()
|
||||
|
||||
_experiment_name = row['experiment_name'] if row is not None else ''
|
||||
_experiment = mlflow_client.get_experiment_by_name(_experiment_name) if _experiment_name else None
|
||||
# The experiment is get-or-create, so it can perfectly well be the model's own experiment with
|
||||
# months of training runs in it. `_run_ids` is everything in it — printed, and used to decide
|
||||
# whether the experiment is this import's alone — while `_import_run_ids` is only the run this
|
||||
# import created, matched on the run name the record carries. The undo acts on the second list.
|
||||
_row_run_name = (row['run_name'] if row is not None else '') or ''
|
||||
_run_ids: list[str] = []
|
||||
_import_run_ids: list[str] = []
|
||||
if _experiment is None:
|
||||
print(f' no experiment named {_experiment_name!r}')
|
||||
else:
|
||||
print(_line('experiment', f'{_experiment.name} ({_experiment.experiment_id})'))
|
||||
for _run in mlflow_client.search_runs([_experiment.experiment_id]):
|
||||
_mine = bool(_row_run_name) and _run.info.run_name == _row_run_name
|
||||
print(
|
||||
_line(
|
||||
'run',
|
||||
f'{_run.info.run_name} ({_run.info.run_id}) {_run.info.status}'
|
||||
f'{" <- this import" if _mine else ""}',
|
||||
)
|
||||
)
|
||||
_run_ids.append(_run.info.run_id)
|
||||
if _mine:
|
||||
_import_run_ids.append(_run.info.run_id)
|
||||
for _artifact in mlflow_client.list_artifacts(_run.info.run_id):
|
||||
print(_line(' artifact', _artifact.path))
|
||||
|
||||
# The registered model is looked up by `run_id`, not by a name reconstructed from `run_name`. The
|
||||
# run name is `import-{model_name}-v{model_version}` truncated to 50 characters (`derive_run_name`),
|
||||
# and a 32-character `model_version` leaves ten characters for the name — so `Quartz-demo` comes
|
||||
# back as `Quartz-dem` and the lookup finds nothing. The version's `run_id` is exact, and it is what
|
||||
# ties the registered model to the run this import created whatever the name was truncated to.
|
||||
_registered_names: set[str] = set()
|
||||
_registered_versions: list[tuple[str, str]] = []
|
||||
for _run_id in _import_run_ids:
|
||||
try:
|
||||
for _version in mlflow_client.search_model_versions(f"run_id='{_run_id}'"):
|
||||
_registered_names.add(_version.name)
|
||||
_registered_versions.append((_version.name, str(_version.version)))
|
||||
print(
|
||||
_line(
|
||||
'version',
|
||||
f'{_version.name} v{_version.version} stage={_version.current_stage} '
|
||||
f'run={_version.run_id}',
|
||||
)
|
||||
)
|
||||
print(_line(' source', _version.source))
|
||||
except Exception as error:
|
||||
# A read-back cell reports what it could not read; it does not fail the inspection.
|
||||
print(f' registry read failed: {type(error).__name__}: {error}')
|
||||
if not _import_run_ids:
|
||||
print(' no run of this import in the experiment yet, so no registered model to look up')
|
||||
elif not _registered_names:
|
||||
print(" no registered version points at this import's run")
|
||||
|
||||
print()
|
||||
print('--- mongodb model document ---')
|
||||
|
||||
_credentials = f'{MONGODB_USERNAME}:{MONGODB_PASSWORD}@' if MONGODB_USERNAME else ''
|
||||
mongo_client: MongoClient = MongoClient(f'mongodb://{_credentials}{MONGODB_URL}')
|
||||
_collection = mongo_client[MONGODB_DATABASE][MODELS_COLLECTION]
|
||||
for _document in _collection.find({}, {'_id': False}).sort('id', 1):
|
||||
print(_line('document', _document))
|
||||
|
||||
# %%
|
||||
# Cell 10 — undo.
|
||||
#
|
||||
# The object and the row always: they are what this script created, and leaving them behind leaves a
|
||||
# PENDING import nobody will ever run. What MLflow and MongoDB hold is left alone unless
|
||||
# DELETE_PROVISIONING is set, because after a successful run that is the thing you wanted to look at.
|
||||
#
|
||||
# When it *is* set, the undo removes only what this import added. An import of a model the platform
|
||||
# already had lands as a new *version* of that registered model and writes no document at all, so
|
||||
# deleting the registered model by name would take versions trained months ago with it. The version
|
||||
# this import created is deleted by number; the registered model and the document go only when
|
||||
# nothing else is left pointing at that name — which is exactly the case where this import was the
|
||||
# first one.
|
||||
DELETE_PROVISIONING = False
|
||||
|
||||
try:
|
||||
minio_client.remove_object(BUNDLE_BUCKET, object_key)
|
||||
print(_line('deleted object', f'{BUNDLE_BUCKET}/{object_key}'))
|
||||
except Exception as error:
|
||||
# Undo reports what it could not remove and carries on to the row.
|
||||
print(f' could not delete the object: {type(error).__name__}: {error}')
|
||||
|
||||
with status_engine.begin() as connection:
|
||||
_deleted = connection.execute(
|
||||
text('DELETE FROM public.experiment_run WHERE id = :id'),
|
||||
{'id': import_run_id},
|
||||
).rowcount
|
||||
print(_line('deleted rows', _deleted))
|
||||
|
||||
if DELETE_PROVISIONING:
|
||||
for _name, _number in _registered_versions:
|
||||
try:
|
||||
mlflow_client.delete_model_version(_name, _number)
|
||||
print(_line('deleted version', f'{_name} v{_number}'))
|
||||
except Exception as error:
|
||||
# Same rule as above: report and carry on to the next one.
|
||||
print(f' could not delete the model version: {type(error).__name__}: {error}')
|
||||
|
||||
for _registered in sorted(_registered_names):
|
||||
try:
|
||||
_remaining = list(mlflow_client.search_model_versions(f"name='{_registered}'"))
|
||||
except Exception as error:
|
||||
print(f' could not count the remaining versions: {type(error).__name__}: {error}')
|
||||
continue
|
||||
if _remaining:
|
||||
# Versions this import did not create: the model was on the platform before it ran, so
|
||||
# the registered model and its document are somebody else's and stay exactly as they are.
|
||||
print(_line('kept model', f'{_registered} ({len(_remaining)} other versions)'))
|
||||
continue
|
||||
try:
|
||||
mlflow_client.delete_registered_model(_registered)
|
||||
print(_line('deleted model', _registered))
|
||||
except Exception as error:
|
||||
print(f' could not delete the registered model: {type(error).__name__}: {error}')
|
||||
_removed = _collection.delete_many({'name': _registered}).deleted_count
|
||||
print(_line('deleted documents', f'{_removed} named {_registered}'))
|
||||
|
||||
# Last, and only when the experiment holds nothing but this import's own run: the experiment is
|
||||
# get-or-create, so an import into a model the platform already had reuses the experiment its
|
||||
# training runs live in, and deleting it would take them all.
|
||||
if _experiment is not None:
|
||||
_others = len(_run_ids) - len(_import_run_ids)
|
||||
if _others:
|
||||
print(_line('kept experiment', f'{_experiment.name} ({_others} other runs)'))
|
||||
else:
|
||||
mlflow_client.delete_experiment(_experiment.experiment_id)
|
||||
print(_line('deleted experiment', _experiment.experiment_id))
|
||||
else:
|
||||
print(' provisioning kept — set DELETE_PROVISIONING = True to remove it too')
|
||||
|
||||
status_engine.dispose()
|
||||
mongo_client.close()
|
||||
@@ -1,6 +1,17 @@
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from sientia_do.temporal.activities.postgres_sync import Postgres
|
||||
|
||||
|
||||
def _noop_postgres_del(_self):
|
||||
"""
|
||||
Unit tests use MagicMock metrics controllers; postgres_sync.Postgres.__del__ calls
|
||||
close() during GC and triggers async shutdown. Explicit ``close()`` is covered in tests.
|
||||
"""
|
||||
return None
|
||||
|
||||
|
||||
Postgres.__del__ = _noop_postgres_del # type: ignore[method-assign]
|
||||
|
||||
# The production code converts SIENTIA_MINIO_OFFLOAD_THRESHOLD_MEGABYTES to int at import-time.
|
||||
# Tests must set it to a valid integer string to avoid import errors.
|
||||
@@ -47,13 +58,9 @@ class DummyMinioDataFramePayload:
|
||||
"""
|
||||
Pytest configuration file with global mocks for external dependencies.
|
||||
|
||||
This module mocks the 'sientia' module to avoid requiring its installation
|
||||
during unit tests. The mock is registered in sys.modules before any test
|
||||
imports are executed.
|
||||
The historical ``sientia`` package is no longer imported by the codebase;
|
||||
drift analysis lives in ``sientia_model.analytics.drift_analysis`` and is
|
||||
imported lazily inside Temporal activities. No global module-level mock is
|
||||
required here — unit tests that need to control ``DriftAnalysis`` outputs
|
||||
should patch ``laborious.activities.model_metrics.DriftAnalysis`` directly.
|
||||
"""
|
||||
|
||||
# Mock sientia module
|
||||
sientia_mock = MagicMock()
|
||||
sientia_mock.ModelAnalysis = MagicMock
|
||||
sys.modules['sientia'] = sientia_mock
|
||||
sys.modules['sientia.ModelAnalysis'] = MagicMock()
|
||||
|
||||
86
tests/fixtures/bundle/README.md
vendored
86
tests/fixtures/bundle/README.md
vendored
@@ -1,86 +0,0 @@
|
||||
# Golden `.sientia` fixtures
|
||||
|
||||
Binary conformance vectors for `laborious/utils/bundle/`. They are **bytes the producer's own code
|
||||
emitted**, not bytes this repository knows how to write: there is no production writer here (change
|
||||
`import-model-workflow`, design D1). If the producer's layout moves, the reader's tests fail against
|
||||
these files instead of failing in production.
|
||||
|
||||
| File | What it is |
|
||||
|---|---|
|
||||
| `golden_model_v3.sientia` | `format_version` 1 bundle, 110767 bytes, plaintext zip 110569 bytes — **two chunks**, so `MESSAGE`/`FINAL` framing is exercised by the golden fixture itself |
|
||||
| `golden_model_v3.metadata.json` | the `metadata.json` that went into that bundle, in plaintext, so a test can assert the recovered value without decrypting first |
|
||||
| `regenerate.py` | the script that produced them, importing the producer's `bundle.py`/`crypto.py` by path |
|
||||
|
||||
## Password
|
||||
|
||||
```
|
||||
Import1Model!Fixture
|
||||
```
|
||||
|
||||
Recorded deliberately: the fixture protects nothing. It satisfies the producer's export-side
|
||||
strength rule (≥ 12 characters, upper, lower, digit, symbol), which is checked by
|
||||
`crypto.validate_password_strength` at generation time.
|
||||
|
||||
## Provenance
|
||||
|
||||
- Producer repository: `sientia-projects-templates`
|
||||
- Producer commit: `8d3c1fd117c697e8b352c34f11c732809a4a129b`
|
||||
(`8d3c1fd QTZPOC-7: adapt to lint rules picked up from merging main`)
|
||||
- Producer modules used: `app/src/operations/model_export/bundle.py`,
|
||||
`app/src/operations/model_export/crypto.py` — loaded by path, so every constant, the `struct`
|
||||
layout, the zip layout, the KDF parameters and the chunk framing come from the producer and none of
|
||||
them is re-implemented here
|
||||
- Producer pin honoured: `pynacl==1.5.0` (this repository pins the same version)
|
||||
|
||||
One deviation from task 0.3, recorded rather than hidden: the bytes were generated with the
|
||||
producer's **code** under *this* repository's venv, not inside the producer's own venv, because that
|
||||
venv has neither `pynacl` nor `mlflow` installed. The dependency that matters for the bytes is
|
||||
`pynacl`, pinned identically on both sides, and `mlflow` is not reached — `regenerate.py` builds the
|
||||
artifact tree itself instead of calling `client.download_artifacts`.
|
||||
|
||||
## Contents of the bundle
|
||||
|
||||
```
|
||||
metadata.json (written first, no indent, by bundle.build_zip)
|
||||
artifacts/prediction_model/MLmodel
|
||||
artifacts/prediction_model/model.pkl (90 KiB of incompressible filler)
|
||||
artifacts/prediction_model/conda.yaml
|
||||
artifacts/prediction_model/python_env.yaml
|
||||
artifacts/prediction_model/requirements.txt
|
||||
artifacts/data_model/MLmodel
|
||||
artifacts/data_model/model.pkl (8 KiB)
|
||||
artifacts/data_model/conda.yaml
|
||||
artifacts/data_model/python_env.yaml
|
||||
artifacts/data_model/requirements.txt
|
||||
artifacts/model_card/model_card.json
|
||||
artifacts/train_data.csv
|
||||
artifacts/test_data.csv
|
||||
artifacts/raw_data.csv
|
||||
```
|
||||
|
||||
No directory entries are stored — the producer only adds paths for which `is_file()` is true.
|
||||
|
||||
`parameters` carries `target_variable: "SE"`, which gate 6 requires and which becomes
|
||||
`model_config.target` in the MongoDB model document.
|
||||
|
||||
## Digest boundary — the reason this fixture exists
|
||||
|
||||
The header's `ciphertext_digest` is `sha256(file[164:])`: the producer computes
|
||||
`hashlib.sha256(ciphertext)` over the joined chunks and only then writes
|
||||
`header + stream_header + ciphertext`, so the 24-byte `secretstream` header is **not** covered.
|
||||
Verified against this fixture and against a freshly produced bundle. Earlier prose (this change's own
|
||||
first draft, and `aig-docs/03-sientia-bundle-format.md`) said "everything after the plaintext
|
||||
header", i.e. from offset 140 — that is wrong, and `tests/laborious/utils/bundle/test_format.py`
|
||||
asserts both directions so it cannot come back.
|
||||
|
||||
## Regenerating
|
||||
|
||||
```bash
|
||||
source ./venv/bin/activate
|
||||
PRODUCER_REPO=../sientia-projects-templates python tests/fixtures/bundle/regenerate.py
|
||||
```
|
||||
|
||||
Two fields change on every run and are expected to: the 16-byte salt (`os.urandom`) and
|
||||
`metadata.metadata.export_timestamp` (`datetime.now(UTC)`). Everything else is deterministic —
|
||||
the filler bytes come from a seeded PRNG. After regenerating, update the producer commit above and
|
||||
the byte counts in the table.
|
||||
@@ -1,24 +0,0 @@
|
||||
{
|
||||
"parameters": {
|
||||
"model_name": "fixture_model",
|
||||
"target_variable": "SE",
|
||||
"input_variables": "['feature_a', 'feature_b']",
|
||||
"models_params": "{'alpha': 0.1}",
|
||||
"nan_treatment": "drop",
|
||||
"scaler_name": "StandardScaler",
|
||||
"train_size": "0.8",
|
||||
"test_size": "0.2",
|
||||
"include_ar": "False",
|
||||
"retrain": "False",
|
||||
"start_date": "2026-01-01 00:00:00",
|
||||
"end_date": "2026-08-01 00:00:00"
|
||||
},
|
||||
"metadata": {
|
||||
"model_name": "fixture_model",
|
||||
"experiment_name": "fixture_experiment",
|
||||
"run_id": "8f14e45fceea167a5a36dedd4bea2543",
|
||||
"model_version": "3",
|
||||
"model_project": "fixture_project",
|
||||
"export_timestamp": "2026-08-25T12:08:59.533441+00:00"
|
||||
}
|
||||
}
|
||||
BIN
tests/fixtures/bundle/golden_model_v3.sientia
vendored
BIN
tests/fixtures/bundle/golden_model_v3.sientia
vendored
Binary file not shown.
185
tests/fixtures/bundle/regenerate.py
vendored
185
tests/fixtures/bundle/regenerate.py
vendored
@@ -1,185 +0,0 @@
|
||||
"""Regenerate the golden `.sientia` fixtures using the producer's own source.
|
||||
|
||||
This script is the auditable half of the anti-drift contract described in the change
|
||||
`import-model-workflow` (design D1): the committed fixtures are bytes produced by
|
||||
`sientia-projects-templates/app/src/operations/model_export/{bundle,crypto}.py`, loaded here by
|
||||
path so no constant is re-implemented. Regenerating is a deliberate act — see the README next to
|
||||
this file for the recorded producer commit and the fixture password.
|
||||
|
||||
Usage (from the repository root, inside the venv):
|
||||
|
||||
PRODUCER_REPO=../sientia-projects-templates ./venv/bin/python tests/fixtures/bundle/regenerate.py
|
||||
|
||||
The producer's own venv is not used: it does not have `pynacl` installed, while this repository
|
||||
pins the producer's exact version (`pynacl==1.5.0`). The bytes are therefore produced by the
|
||||
producer's code on the producer's dependency pin, which is what the fixture is for.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
FIXTURE_DIR = Path(__file__).resolve().parent
|
||||
DEFAULT_PRODUCER_REPO = FIXTURE_DIR.parents[2].parent / 'sientia-projects-templates'
|
||||
|
||||
# Recorded in the README alongside the fixtures; the password satisfies the producer's own
|
||||
# export-side strength rule (>= 12 chars, upper, lower, digit, symbol).
|
||||
FIXTURE_PASSWORD = 'Import1Model!Fixture'
|
||||
|
||||
# Deterministic filler so a regeneration differs only in the salt and the export timestamp.
|
||||
FILLER_SEED = 20260825
|
||||
|
||||
|
||||
def _load_producer_module(repo: Path, name: str) -> types.ModuleType:
|
||||
"""Import one producer module by path, without importing the Streamlit app."""
|
||||
path = repo / 'app' / 'src' / 'operations' / 'model_export' / f'{name}.py'
|
||||
if not path.is_file():
|
||||
raise SystemExit(f'producer module not found: {path}')
|
||||
spec = importlib.util.spec_from_file_location(f'producer_{name}', path)
|
||||
if spec is None or spec.loader is None:
|
||||
raise SystemExit(f'cannot load producer module: {path}')
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def _producer_commit(repo: Path) -> str:
|
||||
result = subprocess.run( # noqa: S603 - fixed argv, no shell
|
||||
['git', '-C', str(repo), 'rev-parse', 'HEAD'], # noqa: S607 - git from PATH is intended
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=False,
|
||||
)
|
||||
return result.stdout.strip() or 'unknown'
|
||||
|
||||
|
||||
def _write_model_dir(root: Path, flavor: str, pickle_bytes: bytes) -> None:
|
||||
"""Write the files an MLflow 2.x model directory carries for this platform."""
|
||||
root.mkdir(parents=True, exist_ok=True)
|
||||
(root / 'MLmodel').write_text(
|
||||
f'artifact_path: {flavor}\n'
|
||||
'flavors:\n'
|
||||
' python_function:\n'
|
||||
' loader_module: mlflow.sklearn\n'
|
||||
' model_path: model.pkl\n'
|
||||
' sklearn:\n'
|
||||
' pickled_model: model.pkl\n'
|
||||
' serialization_format: cloudpickle\n'
|
||||
' sklearn_version: 1.4.2\n'
|
||||
'mlflow_version: 2.10.1\n'
|
||||
'model_uuid: 0f9c1c1f9c1f4c0f9c1c1f9c1f4c0f9c\n'
|
||||
'utc_time_created: "2026-08-25 12:00:00.000000"\n',
|
||||
encoding='utf-8',
|
||||
)
|
||||
(root / 'model.pkl').write_bytes(pickle_bytes)
|
||||
(root / 'conda.yaml').write_text(
|
||||
'channels:\n- conda-forge\ndependencies:\n- python=3.11\n- pip\n- pip:\n'
|
||||
' - mlflow==2.10.1\n - scikit-learn==1.4.2\nname: mlflow-env\n',
|
||||
encoding='utf-8',
|
||||
)
|
||||
(root / 'python_env.yaml').write_text(
|
||||
'python: 3.11.15\nbuild_dependencies:\n- pip==24.0\ndependencies:\n- -r requirements.txt\n',
|
||||
encoding='utf-8',
|
||||
)
|
||||
(root / 'requirements.txt').write_text(
|
||||
'mlflow==2.10.1\nscikit-learn==1.4.2\ncloudpickle==3.0.0\n', encoding='utf-8'
|
||||
)
|
||||
|
||||
|
||||
def _write_csv(path: Path, rows: int, rng: random.Random) -> None:
|
||||
lines = ['datetime,feature_a,feature_b,SE']
|
||||
for i in range(rows):
|
||||
lines.append(
|
||||
f'2026-08-{(i % 27) + 1:02d}T00:00:00,'
|
||||
f'{rng.uniform(0, 100):.6f},{rng.uniform(0, 100):.6f},{rng.uniform(0, 100):.6f}'
|
||||
)
|
||||
path.write_text('\n'.join(lines) + '\n', encoding='utf-8')
|
||||
|
||||
|
||||
def _build_artifact_tree(root: Path) -> None:
|
||||
rng = random.Random(FILLER_SEED) # noqa: S311 - test filler, not key material
|
||||
# Incompressible filler keeps the zip above the 64 KiB chunk size, so the golden fixture
|
||||
# exercises multi-chunk `MESSAGE`/`FINAL` framing with real producer bytes.
|
||||
big_pickle = bytes(rng.getrandbits(8) for _ in range(90 * 1024))
|
||||
small_pickle = bytes(rng.getrandbits(8) for _ in range(8 * 1024))
|
||||
_write_model_dir(root / 'prediction_model', 'prediction_model', big_pickle)
|
||||
_write_model_dir(root / 'data_model', 'data_model', small_pickle)
|
||||
(root / 'model_card').mkdir(parents=True, exist_ok=True)
|
||||
(root / 'model_card' / 'model_card.json').write_text(
|
||||
json.dumps(
|
||||
{
|
||||
'model_name': 'fixture_model',
|
||||
'target': 'SE',
|
||||
'created_by': 'fixture',
|
||||
'creation_method': 'time_series_template',
|
||||
}
|
||||
),
|
||||
encoding='utf-8',
|
||||
)
|
||||
_write_csv(root / 'train_data.csv', 120, rng)
|
||||
_write_csv(root / 'test_data.csv', 40, rng)
|
||||
_write_csv(root / 'raw_data.csv', 200, rng)
|
||||
|
||||
|
||||
RUN_PARAMETERS = {
|
||||
'model_name': 'fixture_model',
|
||||
'target_variable': 'SE',
|
||||
'input_variables': "['feature_a', 'feature_b']",
|
||||
'models_params': "{'alpha': 0.1}",
|
||||
'nan_treatment': 'drop',
|
||||
'scaler_name': 'StandardScaler',
|
||||
'train_size': '0.8',
|
||||
'test_size': '0.2',
|
||||
'include_ar': 'False',
|
||||
'retrain': 'False',
|
||||
'start_date': '2026-01-01 00:00:00',
|
||||
'end_date': '2026-08-01 00:00:00',
|
||||
}
|
||||
|
||||
METADATA_FIELDS = {
|
||||
'model_name': 'fixture_model',
|
||||
'experiment_name': 'fixture_experiment',
|
||||
'run_id': '8f14e45fceea167a5a36dedd4bea2543',
|
||||
'model_version': '3',
|
||||
'model_project': 'fixture_project',
|
||||
}
|
||||
|
||||
|
||||
def main() -> int:
|
||||
repo = Path(os.getenv('PRODUCER_REPO', str(DEFAULT_PRODUCER_REPO))).resolve()
|
||||
bundle = _load_producer_module(repo, 'bundle')
|
||||
crypto = _load_producer_module(repo, 'crypto')
|
||||
|
||||
crypto.validate_password_strength(FIXTURE_PASSWORD)
|
||||
|
||||
staging = Path(tempfile.mkdtemp(prefix='sientia_fixture_'))
|
||||
try:
|
||||
_build_artifact_tree(staging)
|
||||
metadata = bundle.build_metadata_json(run_parameters=RUN_PARAMETERS, **METADATA_FIELDS)
|
||||
zip_bytes = bundle.build_zip(staging, metadata)
|
||||
sientia_bytes = crypto.encrypt_archive(zip_bytes, FIXTURE_PASSWORD)
|
||||
finally:
|
||||
shutil.rmtree(staging, ignore_errors=True)
|
||||
|
||||
target = FIXTURE_DIR / 'golden_model_v3.sientia'
|
||||
target.write_bytes(sientia_bytes)
|
||||
(FIXTURE_DIR / 'golden_model_v3.metadata.json').write_text(
|
||||
json.dumps(metadata, indent=2) + '\n', encoding='utf-8'
|
||||
)
|
||||
print(f'wrote {target} ({len(sientia_bytes)} bytes, zip {len(zip_bytes)} bytes)')
|
||||
print(f'producer commit: {_producer_commit(repo)}')
|
||||
print(f'password: {FIXTURE_PASSWORD}')
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
sys.exit(main())
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user