lablup / lablup/backend.ai

Migrate commit session background task to retriable pattern

Open
#6,624 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
670
Forks
183
Avg merge
17h 7m
Merged PRs (30d)
358

Description

## Overview

Migrated the commit session background task from inline function pattern using `start()` to the retriable pattern using `start_retriable()` with Manifest and Handler classes.

## Changes Made

### 1. New Handler Module

Created `src/ai/backend/manager/bgtask/tasks/commit_session.py` (273 lines):

- `CommitSessionManifest`: Pydantic model with 7 fields (session_id, registry_hostname, registry_project, image_name, image_visibility, image_owner_id, user_email)
- `CommitSessionResult`: Result model with image_id or error_message
- `CommitSessionHandler`: Handler implementing full business logic previously in inline function

### 2. Enum Update

Added `COMMIT_SESSION = "commit_session"` to `ManagerBgtaskName` enum in `src/ai/backend/manager/bgtask/types.py`

### 3. Handler Registration

Updated `src/ai/backend/manager/server_bgtask_ctx.py` to register `CommitSessionHandler` with dependencies:

- session_repository
- agent_registry
- event_hub
- event_fetcher

### 4. Call Site Refactoring

Modified `src/ai/backend/manager/services/session/service.py`:

- Removed 170+ line inline `_commit_and_upload()` function
- Replaced with manifest creation and `start_retriable()` call
- Simplified logic with early validation

### 5. Test Coverage

Created `tests/manager/bgtask/tasks/test_commit_session.py` with tests for:

- Manifest creation with required fields
- Manifest serialization/deserialization
- Handler metadata (name, manifest_type)

## Benefits

- **Retriability**: Task state is now serializable and retriable via Valkey
- **Separation of Concerns**: Dedicated handler class instead of inline closure
- **Consistency**: Follows same pattern as other background tasks (RescanGPUAllocMaps, RescanImages, PurgeImages)
- **Testability**: Isolated handler easier to unit test
- **Maintainability**: Clearer code structure with explicit manifest fields

## Technical Details

The handler executes the following workflow:

1. Validate session exists
1. Get registry configuration
1. Resolve base image and build new canonical name
1. Check for existing customized image and reuse ID if exists
1. Commit session with image labels
1. Push image to registry (if not local)
1. Rescan updated image
1. Return result with image_id or error_message

Event propagation for waiting on agent background tasks is handled internally by the handler using `WithCachePropagator` and `EventHub`.

## Verification

✅ Type check passed: `pants check src/ai/backend/manager::`
✅ Lint passed: `pants lint src/ai/backend/manager::`
✅ Tests passed: `pants test tests/manager/bgtask/tasks/test_commit_session.py`
✅ Pre-commit hooks passed

## Commit

Commit hash: 656b4e324

JIRA Issue: BA-2953

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.