diff --git a/CHANGELOG.md b/CHANGELOG.md index 6d0241f8..ddb8a01e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ ### ✨ Added +- **Container Status / Progress (`mageflow.astatus`)** (#134): Container signatures (swarms and chains) now expose an `astatus()` method returning a structured `ContainerStatus` (total / finished / failed / running / pending, terminal-state percentage, and completion flag). The new `mageflow.astatus(*ids)` loads several containers in a single Redis lookup and returns a `ContainersStatus` with an aggregate `overall_percentage`, raising on missing or non-container ids. - **Signing Hatchet Workflows (`MageWorkflow`)**: Native Hatchet `Workflow` objects can now be tracked by mageflow's signature lifecycle, enabling status callbacks (success/failure) without wrapping tasks in mageflow decorators. ### 🐛 Fixed diff --git a/docs/api/chain.md b/docs/api/chain.md index 5d6c9bc0..dccf4928 100644 --- a/docs/api/chain.md +++ b/docs/api/chain.md @@ -81,3 +81,11 @@ async def interrupt() ``` Interrupts all tasks in the chain and sets the status to `INTERRUPTED`. + +#### `astatus()` + +Return a `ContainerStatus` describing the chain's progress. Child tasks are classified by their `task_status` (done / failed / running / pending) and the percentage is terminal-state based. See [`mageflow.astatus`](functions.md#mageflowastatussignature_ids) for the model fields and the batch helper that covers several containers at once. + +```python +async def astatus() -> ContainerStatus +``` diff --git a/docs/api/client.md b/docs/api/client.md index b61ed1d5..94937195 100644 --- a/docs/api/client.md +++ b/docs/api/client.md @@ -67,6 +67,15 @@ Create a task swarm. swarm = await client.aswarm(tasks=[task1, task2], task_name="my-swarm") ``` +#### `astatus(*signature_ids)` + +Report the progress of one or more container signatures (swarms / chains). See [`mageflow.astatus`](functions.md#mageflowastatussignature_ids) for the returned `ContainersStatus` / `ContainerStatus` models. + +```python +status = await client.astatus(swarm.key, chain.key) +print(status.overall_percentage) +``` + #### `with_ctx` Override the default parameter configuration to enable context for a specific task. diff --git a/docs/api/functions.md b/docs/api/functions.md index 64a35da8..59d81d60 100644 --- a/docs/api/functions.md +++ b/docs/api/functions.md @@ -130,6 +130,66 @@ async def load_signature(key: RapyerKey) -> Optional[Signature] signature = await mageflow.load_signature(task_key) ``` +## Status + +### `mageflow.astatus(*signature_ids)` + +Report the progress of one or more **container** signatures (swarms / chains) in a single Redis lookup. + +```python +async def astatus(*signature_ids: RapyerKey) -> ContainersStatus +``` + +**Parameters:** + +- `*signature_ids` (RapyerKey): Keys of the container signatures to inspect. + +**Returns:** a `ContainersStatus` model describing every requested container. Calling it with no ids returns an empty `ContainersStatus` without touching Redis. + +**Raises:** + +- `MissingSignatureError`: one of the keys does not exist. +- `NotAContainerError`: one of the keys points to a non-container signature (e.g. a plain `TaskSignature`). + +```python +status = await mageflow.astatus(swarm.key, chain.key) +for container in status.containers: + print(container.task_name, container.percentage, container.is_done) +print("overall", status.overall_percentage) +``` + +The percentage is terminal-state based — `(finished + failed) / total * 100` — so it reflects how many child tasks have reached a final state. + +The result models are defined in `thirdmagic` and imported from there: + +```python +from thirdmagic import ContainerStatus, ContainersStatus +``` + +#### `ContainerStatus` + +Per-container breakdown returned inside `ContainersStatus.containers`. + +| Field | Type | Description | +| --- | --- | --- | +| `signature_id` | RapyerKey | Key of the container signature | +| `task_name` | str | The container's task name | +| `status` | SignatureStatus | The container's own lifecycle status | +| `total` | int | Total child tasks | +| `finished` | int | Children that completed successfully | +| `failed` | int | Children that errored | +| `running` | int | Children currently running | +| `pending` | int | Children not yet started | +| `percentage` | float | `(finished + failed) / total * 100`, `0.0` when `total == 0` | +| `is_done` | bool | Whether the container itself has completed | + +#### `ContainersStatus` + +| Member | Type | Description | +| --- | --- | --- | +| `containers` | list[ContainerStatus] | One entry per requested container | +| `overall_percentage` | float (property) | Terminal-state percentage aggregated across all containers | + ## Atomic Operations ### `mageflow.abounded_field(ignore_redis_error=False)` diff --git a/docs/api/swarm.md b/docs/api/swarm.md index b40ebc45..8de82821 100644 --- a/docs/api/swarm.md +++ b/docs/api/swarm.md @@ -135,6 +135,14 @@ Check if swarm has completed all tasks. async def is_swarm_done() -> bool ``` +#### `astatus()` + +Return a `ContainerStatus` describing the swarm's progress (counts + terminal-state percentage), computed from its bookkeeping lists without loading child tasks. See [`mageflow.astatus`](functions.md#mageflowastatussignature_ids) for the model fields and the batch helper that covers several containers at once. + +```python +async def astatus() -> ContainerStatus +``` + ## Error Classes ### TooManyTasksError diff --git a/libs/mageflow/mageflow/__init__.py b/libs/mageflow/mageflow/__init__.py index 95c8000a..88bf4f81 100644 --- a/libs/mageflow/mageflow/__init__.py +++ b/libs/mageflow/mageflow/__init__.py @@ -3,6 +3,7 @@ from thirdmagic import abounded_field from thirdmagic.chain.creator import chain as achain from thirdmagic.signature import Signature +from thirdmagic.status import astatus from thirdmagic.swarm.creator import swarm as aswarm from thirdmagic.task import TaskSignature from thirdmagic.task import sign as asign @@ -39,6 +40,7 @@ async def load_sign(key: RapyerKey) -> Signature: "SignatureTTLConfig", "achain", "aswarm", + "astatus", "start_mageflow", "abounded_field", ] diff --git a/libs/mageflow/mageflow/clients/hatchet/mageflow.py b/libs/mageflow/mageflow/clients/hatchet/mageflow.py index 7f114651..2db91102 100644 --- a/libs/mageflow/mageflow/clients/hatchet/mageflow.py +++ b/libs/mageflow/mageflow/clients/hatchet/mageflow.py @@ -17,10 +17,13 @@ ) from hatchet_sdk.runnables.workflow import BaseWorkflow, Standalone from hatchet_sdk.worker.worker import LifespanFn +from rapyer.fields import RapyerKey from redis.asyncio import Redis from thirdmagic import chain, sign from thirdmagic.chain import ChainTaskSignature from thirdmagic.signature import Signature +from thirdmagic.signature.status import ContainersStatus +from thirdmagic.status import astatus from thirdmagic.swarm import SwarmTaskSignature from thirdmagic.swarm.creator import SignatureOptions, swarm from thirdmagic.task import TaskInputType, TaskSignature, TaskSignatureConvertible @@ -329,6 +332,9 @@ async def aswarm( ): return await swarm(tasks, task_name, **kwargs) + async def astatus(self, *signature_ids: RapyerKey) -> ContainersStatus: + return await astatus(*signature_ids) + def with_ctx(self, func): func.__user_ctx__ = True return func diff --git a/libs/mageflow/tests/unit/workflows/test_astatus.py b/libs/mageflow/tests/unit/workflows/test_astatus.py new file mode 100644 index 00000000..0b286738 --- /dev/null +++ b/libs/mageflow/tests/unit/workflows/test_astatus.py @@ -0,0 +1,118 @@ +import pytest +from thirdmagic import ContainersStatus, ContainerStatus +from thirdmagic.errors import MissingSignatureError, NotAContainerError +from thirdmagic.signature.status import SignatureStatus + +import mageflow +from tests.integration.hatchet.models import ContextMessage +from tests.unit.workflows.conftest import create_swarm_item_test_setup + + +@pytest.mark.asyncio +async def test_astatus_returns_status_for_single_container(mock_adapter): + # Arrange + setup = await create_swarm_item_test_setup( + num_tasks=4, + stop_after_n_failures=None, + current_running=1, + tasks_left_indices=[3], + finished_indices=[0], + failed_indices=[1], + ) + expected = ContainersStatus( + containers=[ + ContainerStatus( + signature_id=setup.swarm_task.key, + task_name="test_swarm", + status=SignatureStatus.PENDING, + total=4, + finished=1, + failed=1, + running=1, + pending=1, + percentage=50.0, + is_done=False, + ) + ] + ) + + # Act + result = await mageflow.astatus(setup.swarm_task.key) + + # Assert + assert result == expected + + +@pytest.mark.asyncio +async def test_astatus_aggregates_multiple_containers(mock_adapter): + # Arrange + first = await create_swarm_item_test_setup( + num_tasks=2, stop_after_n_failures=None, finished_indices=[0, 1] + ) + second = await create_swarm_item_test_setup(num_tasks=2, stop_after_n_failures=None) + expected = ContainersStatus( + containers=[ + ContainerStatus( + signature_id=first.swarm_task.key, + task_name="test_swarm", + status=SignatureStatus.PENDING, + total=2, + finished=2, + failed=0, + running=1, + pending=0, + percentage=100.0, + is_done=False, + ), + ContainerStatus( + signature_id=second.swarm_task.key, + task_name="test_swarm", + status=SignatureStatus.PENDING, + total=2, + finished=0, + failed=0, + running=1, + pending=0, + percentage=0.0, + is_done=False, + ), + ] + ) + + # Act + result = await mageflow.astatus(first.swarm_task.key, second.swarm_task.key) + + # Assert + assert result == expected + # 2 terminal tasks out of 4 total across both swarms + assert result.overall_percentage == 50.0 + + +@pytest.mark.asyncio +async def test_astatus_empty_ids_returns_empty_without_db_scan(mock_adapter): + # Act + result = await mageflow.astatus() + + # Assert + assert isinstance(result, ContainersStatus) + assert result.containers == [] + + +@pytest.mark.asyncio +async def test_astatus_raises_for_non_container(mock_adapter): + # Arrange + task = await mageflow.asign("plain_task", model_validators=ContextMessage) + + # Act / Assert + with pytest.raises(NotAContainerError): + await mageflow.astatus(task.key) + + +@pytest.mark.asyncio +async def test_astatus_raises_for_missing_signature(mock_adapter): + # Arrange + missing_key = "SwarmTaskSignature:00000000-0000-0000-0000-000000000000" + + # Act / Assert + with pytest.raises(MissingSignatureError): + await mageflow.astatus(missing_key) diff --git a/libs/third-magic/tests/unit/test_container_astatus.py b/libs/third-magic/tests/unit/test_container_astatus.py new file mode 100644 index 00000000..526bea0f --- /dev/null +++ b/libs/third-magic/tests/unit/test_container_astatus.py @@ -0,0 +1,126 @@ +import pytest + +import thirdmagic +from thirdmagic import ContainerStatus +from thirdmagic.signature.status import SignatureStatus + + +@pytest.mark.asyncio +async def test_swarm_astatus_terminal_percentage(mock_task_def): + # Arrange + swarm = await thirdmagic.swarm(task_name="test_swarm") + tasks = [await thirdmagic.sign(f"test_task_{i}") for i in range(4)] + await swarm.add_tasks(tasks) + + async with swarm.apipeline(): + swarm.tasks_left_to_run.remove_range(0, len(swarm.tasks_left_to_run)) + swarm.finished_tasks.append(tasks[0].key) + swarm.failed_tasks.append(tasks[1].key) + swarm.current_running_tasks = 1 + swarm.tasks_left_to_run.append(tasks[3].key) + + expected = ContainerStatus( + signature_id=swarm.key, + task_name="test_swarm", + status=SignatureStatus.PENDING, + total=4, + finished=1, + failed=1, + running=1, + pending=1, + percentage=50.0, + is_done=False, + ) + + # Act + status = await swarm.astatus() + + # Assert + assert status == expected + + +@pytest.mark.asyncio +async def test_swarm_astatus_done_is_full(mock_task_def): + # Arrange + swarm = await thirdmagic.swarm(task_name="test_swarm") + tasks = [await thirdmagic.sign(f"test_task_{i}") for i in range(2)] + await swarm.add_tasks(tasks) + + async with swarm.apipeline(): + swarm.tasks_left_to_run.remove_range(0, len(swarm.tasks_left_to_run)) + swarm.finished_tasks.extend([task.key for task in tasks]) + swarm.is_swarm_closed = True + + expected = ContainerStatus( + signature_id=swarm.key, + task_name="test_swarm", + status=SignatureStatus.PENDING, + total=2, + finished=2, + failed=0, + running=0, + pending=0, + percentage=100.0, + is_done=True, + ) + + # Act + status = await swarm.astatus() + + # Assert + assert status == expected + + +@pytest.mark.asyncio +async def test_swarm_astatus_empty_is_zero(mock_task_def): + # Arrange + swarm = await thirdmagic.swarm(task_name="test_swarm") + expected = ContainerStatus( + signature_id=swarm.key, + task_name="test_swarm", + status=SignatureStatus.PENDING, + total=0, + finished=0, + failed=0, + running=0, + pending=0, + percentage=0.0, + is_done=False, + ) + + # Act + status = await swarm.astatus() + + # Assert + assert status == expected + + +@pytest.mark.asyncio +async def test_chain_astatus_classifies_children(mock_task_def): + # Arrange + tasks = [await thirdmagic.sign(f"chain_task_{i}") for i in range(4)] + chain = await thirdmagic.chain([task.key for task in tasks]) + + await tasks[0].change_status(SignatureStatus.DONE) + await tasks[1].change_status(SignatureStatus.FAILED) + await tasks[2].change_status(SignatureStatus.ACTIVE) + # tasks[3] stays PENDING + + expected = ContainerStatus( + signature_id=chain.key, + task_name=chain.task_name, + status=SignatureStatus.PENDING, + total=4, + finished=1, + failed=1, + running=1, + pending=1, + percentage=50.0, + is_done=False, + ) + + # Act + status = await chain.astatus() + + # Assert + assert status == expected diff --git a/libs/third-magic/thirdmagic/__init__.py b/libs/third-magic/thirdmagic/__init__.py index 5daa2b19..3d9eca64 100644 --- a/libs/third-magic/thirdmagic/__init__.py +++ b/libs/third-magic/thirdmagic/__init__.py @@ -1,9 +1,17 @@ import rapyer from thirdmagic.chain.creator import chain +from thirdmagic.signature.status import ContainersStatus, ContainerStatus from thirdmagic.swarm.creator import swarm from thirdmagic.task.creator import sign abounded_field = rapyer.apipeline -__all__ = ["sign", "chain", "swarm", "abounded_field"] +__all__ = [ + "sign", + "chain", + "swarm", + "abounded_field", + "ContainerStatus", + "ContainersStatus", +] diff --git a/libs/third-magic/thirdmagic/chain/model.py b/libs/third-magic/thirdmagic/chain/model.py index 2cd16365..3ac387b9 100644 --- a/libs/third-magic/thirdmagic/chain/model.py +++ b/libs/third-magic/thirdmagic/chain/model.py @@ -7,7 +7,7 @@ from thirdmagic.container import ContainerTaskSignature from thirdmagic.errors import MissingSignatureError -from thirdmagic.signature.status import SignatureStatus +from thirdmagic.signature.status import ContainerStatus, SignatureStatus from thirdmagic.task.model import TaskSignature from thirdmagic.utils import HAS_HATCHET @@ -47,6 +47,31 @@ async def sub_tasks(self) -> list[TaskSignature]: sub_tasks = await rapyer.afind(*self.tasks, skip_missing=True) return cast(list[TaskSignature], sub_tasks) + async def astatus(self) -> ContainerStatus: + sub_tasks = await self.sub_tasks() + finished = failed = running = 0 + for task in sub_tasks: + status = task.task_status.status + if status == SignatureStatus.DONE: + finished += 1 + elif status == SignatureStatus.FAILED: + failed += 1 + elif status == SignatureStatus.ACTIVE: + running += 1 + total = len(self.tasks) + pending = total - finished - failed - running + return ContainerStatus.from_counts( + signature_id=self.key, + task_name=self.task_name, + status=self.task_status.status, + total=total, + finished=finished, + failed=failed, + running=running, + pending=pending, + is_done=self.task_status.is_done(), + ) + async def acall(self, msg: Any, set_return_field: bool = True, **kwargs): first_task = await rapyer.afind_one(self.tasks[0]) if first_task is None: diff --git a/libs/third-magic/thirdmagic/container.py b/libs/third-magic/thirdmagic/container.py index 5d4f158b..0af80352 100644 --- a/libs/third-magic/thirdmagic/container.py +++ b/libs/third-magic/thirdmagic/container.py @@ -6,6 +6,7 @@ from rapyer.fields import RapyerKey from thirdmagic.signature import Signature +from thirdmagic.signature.status import ContainerStatus class ContainerTaskSignature(Signature, ABC): @@ -18,6 +19,10 @@ def task_ids(self) -> list[RapyerKey]: async def sub_tasks(self) -> list[Signature]: pass + @abc.abstractmethod + async def astatus(self) -> ContainerStatus: + pass + async def remove_references(self): sub_tasks = await self.sub_tasks() await asyncio.gather( diff --git a/libs/third-magic/thirdmagic/errors.py b/libs/third-magic/thirdmagic/errors.py index 8fe87817..c09e8ce9 100644 --- a/libs/third-magic/thirdmagic/errors.py +++ b/libs/third-magic/thirdmagic/errors.py @@ -28,3 +28,7 @@ class TaskAndMsgsDontMatchForSwarmError(SwarmError, RuntimeError): class UnrecognizedTaskError(MageflowError): pass + + +class NotAContainerError(MageflowError): + pass diff --git a/libs/third-magic/thirdmagic/signature/__init__.py b/libs/third-magic/thirdmagic/signature/__init__.py index afe301f7..cb5e3bed 100644 --- a/libs/third-magic/thirdmagic/signature/__init__.py +++ b/libs/third-magic/thirdmagic/signature/__init__.py @@ -1,6 +1,12 @@ from thirdmagic.signature.model import Signature, SignatureConfig from thirdmagic.signature.retry_cache import SignatureRetryCache, retry_cache_ctx -from thirdmagic.signature.status import PauseActionTypes, SignatureStatus, TaskStatus +from thirdmagic.signature.status import ( + ContainersStatus, + ContainerStatus, + PauseActionTypes, + SignatureStatus, + TaskStatus, +) __all__ = [ "Signature", @@ -9,5 +15,7 @@ "SignatureStatus", "PauseActionTypes", "TaskStatus", + "ContainerStatus", + "ContainersStatus", "retry_cache_ctx", ] diff --git a/libs/third-magic/thirdmagic/signature/status.py b/libs/third-magic/thirdmagic/signature/status.py index bc123efd..7ff2d570 100644 --- a/libs/third-magic/thirdmagic/signature/status.py +++ b/libs/third-magic/thirdmagic/signature/status.py @@ -1,8 +1,10 @@ from enum import Enum from typing import ClassVar +from pydantic import BaseModel from rapyer import AtomicRedisModel from rapyer.config import RedisConfig +from rapyer.fields import RapyerKey class SignatureStatus(str, Enum): @@ -33,3 +35,57 @@ def should_run(self): def is_done(self): return self.status in [SignatureStatus.DONE, SignatureStatus.FAILED] + + +class ContainerStatus(BaseModel): + signature_id: RapyerKey + task_name: str + status: SignatureStatus + total: int + finished: int + failed: int + running: int + pending: int + percentage: float + is_done: bool + + @classmethod + def from_counts( + cls, + signature_id: RapyerKey, + task_name: str, + status: SignatureStatus, + total: int, + finished: int, + failed: int, + running: int, + pending: int, + is_done: bool, + ) -> "ContainerStatus": + percentage = (finished + failed) / total * 100 if total else 0.0 + return cls( + signature_id=signature_id, + task_name=task_name, + status=status, + total=total, + finished=finished, + failed=failed, + running=running, + pending=pending, + percentage=percentage, + is_done=is_done, + ) + + +class ContainersStatus(BaseModel): + containers: list[ContainerStatus] + + @property + def overall_percentage(self) -> float: + total = sum(container.total for container in self.containers) + if not total: + return 0.0 + terminal = sum( + container.finished + container.failed for container in self.containers + ) + return terminal / total * 100 diff --git a/libs/third-magic/thirdmagic/status.py b/libs/third-magic/thirdmagic/status.py new file mode 100644 index 00000000..a9e73e97 --- /dev/null +++ b/libs/third-magic/thirdmagic/status.py @@ -0,0 +1,24 @@ +import rapyer +from rapyer.fields import RapyerKey + +from thirdmagic.container import ContainerTaskSignature +from thirdmagic.errors import MissingSignatureError, NotAContainerError +from thirdmagic.signature.status import ContainersStatus + + +async def astatus(*signature_ids: RapyerKey) -> ContainersStatus: + if not signature_ids: + return ContainersStatus(containers=[]) + + signatures = await rapyer.afind(*signature_ids, skip_missing=True) + if len(signatures) != len(signature_ids): + raise MissingSignatureError(f"Some signatures were not found: {signature_ids}") + + statuses = [] + for signature in signatures: + if not isinstance(signature, ContainerTaskSignature): + raise NotAContainerError( + f"Signature {signature.key} is not a container signature" + ) + statuses.append(await signature.astatus()) + return ContainersStatus(containers=statuses) diff --git a/libs/third-magic/thirdmagic/swarm/model.py b/libs/third-magic/thirdmagic/swarm/model.py index e2114cb9..6819f0f8 100644 --- a/libs/third-magic/thirdmagic/swarm/model.py +++ b/libs/third-magic/thirdmagic/swarm/model.py @@ -14,7 +14,7 @@ TooManyTasksError, ) from thirdmagic.signature import Signature -from thirdmagic.signature.status import SignatureStatus +from thirdmagic.signature.status import ContainerStatus, SignatureStatus from thirdmagic.swarm.consts import SWARM_MESSAGE_PARAM_NAME from thirdmagic.swarm.state import PublishState from thirdmagic.task.creator import TaskSignatureConvertible, resolve_signatures @@ -187,6 +187,19 @@ async def is_swarm_done(self): finished_all_tasks = set(done_tasks) == set(self.tasks) return self.is_swarm_closed and finished_all_tasks + async def astatus(self) -> ContainerStatus: + return ContainerStatus.from_counts( + signature_id=self.key, + task_name=self.task_name, + status=self.task_status.status, + total=len(self.tasks), + finished=len(self.finished_tasks), + failed=len(self.failed_tasks), + running=self.current_running_tasks, + pending=len(self.tasks_left_to_run), + is_done=await self.is_swarm_done(), + ) + def has_published_callback(self): return self.task_status.status == SignatureStatus.DONE