Deprecate KernelRow.image column and add image_id column
- Dominant language
- Python
- Stars
- 670
- Forks
- 183
- Avg merge
- 17h 7m
- Merged PRs (30d)
- 358
Description
## Motivation
The `KernelRow` model currently uses the `image` column (VARCHAR 512) to store the canonical image name string (e.g.,
`//:`). This column should be deprecated and replaced with an `image_id` column that references the
image by ID instead of by name.
## Required Features
1. **Investigation Task**: Determine if the `image` column truly needs to be nullable
- Current definition: `image: Mapped[str | None] = mapped_column("image", sa.String(length=512), nullable=True)`
- The `from_kernel_info()` method sets `image=None` when `info.image.identifier` is None
- Historical migration (`eb9441fcf90a`) filters with `kernels.c.image.is_not(None)`, suggesting null values existed
- **Action**: Query production database to check if any `kernels.image IS NULL` records exist and understand the business context
1. **Schema Migration**:
- Add `image_id` column
- Determine nullability based on investigation results
- Create data migration to populate `image_id` from existing `image` values
1. **Code Updates**:
- Update `KernelRow` model to add `image_id` column
- Mark `image` column as deprecated
- Update `image_row` relationship to use `image_id` instead of name-based join
- Update `from_kernel_info()` and `to_kernel_info()` methods
- Update all query locations that reference `KernelRow.image`
## Impact
### Affected Components
- `src/ai/backend/manager/models/kernel/row.py` - KernelRow model definition
- `src/ai/backend/manager/registry.py` - Session/kernel creation and queries
- `src/ai/backend/manager/repositories/scheduler/db_source/db_source.py` - Scheduler queries
- `src/ai/backend/manager/repositories/schedule/` - Schedule-related queries
- `src/ai/backend/manager/repositories/group/repository.py` - Group queries
- `src/ai/backend/manager/api/gql_legacy/kernel.py` - GraphQL kernel API
- `src/ai/backend/manager/models/resource_usage.py` - Resource usage queries
- `src/ai/backend/manager/sokovan/scheduling_controller/` - Scheduling controller
### Database
- Requires Alembic migration for schema change
- Data migration needed to populate `image_id` from existing records
## Testing Scenarios
1. **Unit Tests**:
- Verify `KernelRow` creation with valid `image_id`
- Verify `image_row` relationship loads correctly via new column
- Test `from_kernel_info()` and `to_kernel_info()` with the new schema
1. **Integration Tests**:
- Create session with image and verify `image_id` is properly set
- Query kernels and verify `image_row` relationship works
- Test scheduler queries that filter/join on image
1. **Migration Tests**:
- Verify forward migration populates `image_id` correctly
- Test with production-like data including edge cases (null images if they exist)
1. **Backward Compatibility**:
- Ensure existing APIs continue to work during transition period
- Verify GraphQL queries return correct image information
JIRA Issue: BA-4074
Contributor guide
Assessment
This issue has not been assessed yet.