lablup / lablup/backend.ai

Refactor agent's kernel creation routines using a state machine

Open
#759 5 comments 1 reaction 0 assignees View on GitHub
Dominant language
Python
Stars
670
Forks
183
Avg merge
17h 7m
Merged PRs (30d)
358

Description

It is hard to debug with `print()` statements when there are many external resourecs invoked with `await` statements chained as a long function body.

If we could model the steps to create a kernel using a state machine and implement debugging features on the base state machine, it will make our life much easier. For example, just dumping out all transitions as logs will allow pinpointing where the process got stuck.

So, let's write a simple `ai.backend.common.automata.AsyncStateMachine` (or maybe a separate `aiostate` package) which executes per-state asyncio tasks by the state transition. The state transition may be triggered upon the result of tasks like success, exception, cancelled, timeout, and max-retry-limit-exceeded. Each state may have optional timeout or retry conditions. State transition callbacks may include production of `KernelStartedEvent` or `KernelCancelledEvent` depending on the target state. The state machine itself could serve as a supervisor-scope (currently known as `PersistentTaskGroup`) with a planned feature in `aiotools` that implements a context-scoped `asyncio.create_task()` override so that termination of the state machine guarantees termination of all involved, nested asyncio tasks.

Some idea sketch:

```python
from tenacity import AsyncRetrying, ...

state_machine = AsyncStateMachine(
ctx={"event_dispatcher": event_dispatcher, ...}, # state_machine is injected when invoking states/transition callbacks
states=[
State("check-image", coro_factory=lambda ctx: check_image(ctx), retry=None),
State("pull-image", coro_factory=lambda ctx: pull_image(ctx),
retry=AsyncRetrying(stop=stop_after_attempts(2))),
State("prepare-network", coro_factory=lambda ctx: prepare_network(ctx),
retry=AsyncRetrying(stop=stop_after_timeout(30.0))),
State("create-container", coro_factory=lambda ctx: create_container(ctx), retry=None),
State("start-container", coro_factory=lambda ctx: start_container(ctx), retry=None),
State("wait-kernel-runner", coro_factory=lambda ctx: wait_check_status(ctx),
retry=AsyncRetrying(stop=stop_after_attempts(10), wait=wait_exponential(...)),
State("cleanup", coro_factory=lambda ctx: remove_container_if_exists(ctx), retry=None),
State("done", coro_factory=lambda ctx: None), # may be declared as an intrinsic state
State("failure", coro_factory=lambda ctx: None), # may be declared as an intrinsic state
],
transitions=[
Transition("check-image", "pull-image", trigger="requires-pull", # optionally the task may set a custom result
callback=lambda ctx: send_kernel_pulling_event(ctx)),
Transition("pull-image", "prepare-network", trigger=SUCCESS,
callback=lambda ctx: send_kernel_preparing_event(ctx)),
Transition("check-image", "prepare-network", trigger=SUCCESS,
callback=lambda ctx: send_kernel_preparing_event(ctx)),
Transition("prepare-network", "create-container", trigger=SUCCESS),
Transition("create-container", "start-container", trigger=SUCCESS),
Transition("create-container", "cleanup", trigger=FAILURE),
Transition("start-container", "wait-kernel-runner", trigger=SUCCESS),
Transition("start-container", "cleanup", trigger=FAILURE),
Transition("wait-kernel-runner", "done", trigger=SUCCESS,
callback=lambda ctx: send_kernel_started_event(ctx)),
Transition("wait-kernel-runner", "cleanup", trigger=(ERROR, TIEMOUT)),
Transition("cleanup", "failure", trigger=SUCCESS,
callback=lambda ctx: send_kernel_cancelled_event(ctx)),
],
initial_state="check-image",
success_state="done",
failure_state="failure",
fallback_exception_handler=lambda ctx: log_error_and_send_kernel_cancelled_event(ctx),
)
await state_machine.run()
```

This approach will also allow validation of the state machine if there are any missing error handlers. For example, it could raise an error before actually running it when a state has `retry` argument but the transitions from it does not include `TIMEOUT` transition.

JIRA Issue: BA-243

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.