feat: add multi hooks support (#251)

* support multi-hooks

Signed-off-by: kerthcet <kerthcet@gmail.com>

* add more tests

Signed-off-by: kerthcet <kerthcet@gmail.com>

* fix lint

Signed-off-by: kerthcet <kerthcet@gmail.com>

* fix test

Signed-off-by: kerthcet <kerthcet@gmail.com>

* fix test

Signed-off-by: kerthcet <kerthcet@gmail.com>

---------

Signed-off-by: kerthcet <kerthcet@gmail.com>

Author: kerthcet

alphatrion/run/hooks.py

@@ -1,9 +1,13 @@
 """Built-in post-run hooks for metadata enrichment."""
 
+import logging
 import uuid
 from typing import Any
 
 from alphatrion.runtime.runtime import global_runtime
+from alphatrion.storage.sql_models import Status
+
+logger = logging.getLogger(__name__)
 
 
 class PostRunHookFn:
@@ -14,16 +18,17 @@ def sync_metadata(run_id: uuid.UUID, result: Any) -> None:
         """
         Sync function result to run metadata.
 
-        If the function returns a dict, it will be merged into the run's metadata.
-        This is useful for automatically capturing metrics, model info, etc.
+        Looks for 'metadata' key in result dict and syncs it to run metadata.
 
         Example:
             async def train_model():
-                # ... training code ...
                 return {
-                    "accuracy": 0.95,
-                    "loss": 0.05,
-                    "num_epochs": 10,
+                    "metadata": {
+                        "accuracy": 0.95,
+                        "loss": 0.05,
+                        "num_epochs": 10,
+                    },
+                    "status": "COMPLETED"
                 }
 
             run = exp.run(train_model, post_run_hooks=[PostRunHookFn.sync_metadata])
@@ -32,6 +37,70 @@ async def train_model():
         :param run_id: UUID of the run
         :param result: Return value from the run function
         """
-        if isinstance(result, dict):
+        if result is None:
+            return
+
+        if isinstance(result, dict) and "metadata" in result:
+            metadata = result["metadata"]
+            if isinstance(metadata, dict):
+                metadb = global_runtime().metadb
+                metadb.update_run(run_id=run_id, meta=metadata)
+            else:
+                logger.warning(
+                    f"PostRunHookFn.sync_metadata: 'metadata' key in result for run {run_id} is not a dict. Skipping metadata sync."
+                )
+        else:
+            logger.warning(
+                f"PostRunHookFn.sync_metadata: Result for run {run_id} does not contain 'metadata' key or is not a dict. Skipping metadata sync."
+            )
+
+    @staticmethod
+    def sync_status(run_id: uuid.UUID, result: Any) -> None:
+        """
+        Sync function result to run status.
+
+        Looks for 'status' key in result dict. Status can be a string representation,
+        or integer value.
+
+        Example:
+            async def train_model():
+                return {
+                    "status": "COMPLETED"  # or 9
+                }
+
+            run = exp.run(train_model, post_run_hooks=[
+                PostRunHookFn.sync_status
+            ])
+
+        :param run_id: UUID of the run
+        :param result: Return value from the run function
+        """
+        if result is None:
+            return
+
+        status = None
+
+        # Extract status from dict
+        if isinstance(result, dict) and "status" in result:
+            status_value = result["status"]
+
+            if isinstance(status_value, str):
+                try:
+                    status = Status[status_value.upper()]
+                except (KeyError, AttributeError):
+                    logger.warning(
+                        f"PostRunHookFn.sync_status: Invalid status value '{status_value}' for run {run_id}. Skipping status sync."
+                    )
+                    return
+            elif isinstance(status_value, int):
+                try:
+                    status = Status(status_value)
+                except ValueError:
+                    logger.warning(
+                        f"PostRunHookFn.sync_status: Invalid status value '{status_value}' for run {run_id}. Skipping status sync."
+                    )
+                    return
+
+        if status is not None:
             metadb = global_runtime().metadb
-            metadb.update_run(run_id=run_id, meta=result)
+            metadb.update_run(run_id=run_id, status=status)

tests/integration/test_run_hooks.py

@@ -9,6 +9,7 @@
 from alphatrion.experiment import CraftExperiment, ExperimentConfig
 from alphatrion.run import PostRunHookFn
 from alphatrion.runtime.runtime import global_runtime
+from alphatrion.storage.sql_models import Status
 
 
 @pytest.fixture
@@ -32,13 +33,15 @@ async def test_run_hook_sync_metadata(test_org_id, test_user_id, test_team_id):
     alpha.init(org_id=test_org_id, team_id=test_team_id, user_id=test_user_id)
 
     async def train_model():
-        """Function that returns metrics as dict"""
+        """Function that returns metrics in nested metadata structure"""
         await asyncio.sleep(0.1)
         return {
-            "accuracy": 0.95,
-            "loss": 0.05,
-            "learning_rate": 0.001,
-            "num_epochs": 10,
+            "metadata": {
+                "accuracy": 0.95,
+                "loss": 0.05,
+                "learning_rate": 0.001,
+                "num_epochs": 10,
+            }
         }
 
     async with CraftExperiment.start("test_hook_experiment") as exp:
@@ -84,18 +87,44 @@ async def task_with_string_result():
         assert run_obj.meta is None or run_obj.meta == {}
 
 
+@pytest.mark.asyncio
+async def test_run_hook_with_dict_without_metadata_key(
+    test_org_id, test_user_id, test_team_id
+):
+    """Test sync_metadata hook with dict result but no 'metadata' key"""
+    alpha.init(org_id=test_org_id, team_id=test_team_id, user_id=test_user_id)
+
+    async def task_without_metadata_key():
+        """Function that returns dict without 'metadata' key"""
+        await asyncio.sleep(0.1)
+        return {"accuracy": 0.95}
+
+    async with CraftExperiment.start("test_hook_no_metadata_key") as exp:
+        run = exp.run(
+            task_without_metadata_key, post_run_hooks=[PostRunHookFn.sync_metadata]
+        )
+        await exp.wait()
+
+        # Verify metadata was not updated
+        metadb = global_runtime().metadb
+        run_obj = metadb.get_run(run_id=run.id)
+
+        # Metadata should be None or empty (hook didn't update it)
+        assert run_obj.meta is None or run_obj.meta == {}
+
+
 @pytest.mark.asyncio
 async def test_experiment_level_hooks(test_org_id, test_user_id, test_team_id):
     """Test hooks configured at experiment level apply to all runs"""
     alpha.init(org_id=test_org_id, team_id=test_team_id, user_id=test_user_id)
 
     async def task1():
         await asyncio.sleep(0.1)
-        return {"task": "task1", "accuracy": 0.92}
+        return {"metadata": {"task": "task1", "accuracy": 0.92}}
 
     async def task2():
         await asyncio.sleep(0.1)
-        return {"task": "task2", "accuracy": 0.94}
+        return {"metadata": {"task": "task2", "accuracy": 0.94}}
 
     # Configure experiment with sync_metadata hook
     config = ExperimentConfig(post_run_hooks=[PostRunHookFn.sync_metadata])
@@ -138,7 +167,7 @@ def add_custom_info(run_id, result):
 
     async def train_model():
         await asyncio.sleep(0.1)
-        return {"accuracy": 0.95}
+        return {"metadata": {"accuracy": 0.95}}
 
     async with CraftExperiment.start("test_custom_hook") as exp:
         # Use both built-in and custom hooks
@@ -169,7 +198,7 @@ async def test_hook_merges_with_existing_metadata(
 
     async def train_model():
         await asyncio.sleep(0.1)
-        return {"accuracy": 0.96, "loss": 0.04}
+        return {"metadata": {"accuracy": 0.96, "loss": 0.04}}
 
     async with CraftExperiment.start("test_merge_metadata") as exp:
         run = exp.run(train_model, post_run_hooks=[PostRunHookFn.sync_metadata])
@@ -205,7 +234,7 @@ def buggy_hook(run_id, result):
 
     async def train_model():
         await asyncio.sleep(0.1)
-        return {"accuracy": 0.95}
+        return {"metadata": {"accuracy": 0.95}}
 
     async with CraftExperiment.start("test_hook_failure") as exp:
         run = exp.run(
@@ -220,3 +249,60 @@ async def train_model():
         metadb = global_runtime().metadb
         run_obj = metadb.get_run(run_id=run.id)
         assert run_obj.meta["accuracy"] == 0.95
+
+
+@pytest.mark.asyncio
+async def test_both_hooks_together(test_org_id, test_user_id, test_team_id):
+    """Test sync_metadata and sync_status hooks working together"""
+    alpha.init(org_id=test_org_id, team_id=test_team_id, user_id=test_user_id)
+
+    async def train_model():
+        await asyncio.sleep(0.1)
+        return {
+            "metadata": {"accuracy": 0.95, "loss": 0.05, "num_epochs": 10},
+            "status": "failed",
+        }
+
+    async with CraftExperiment.start("test_both_hooks") as exp:
+        # Use both hooks
+        run = exp.run(
+            train_model,
+            post_run_hooks=[PostRunHookFn.sync_metadata, PostRunHookFn.sync_status],
+        )
+        await exp.wait()
+
+        # Verify both hooks ran
+        metadb = global_runtime().metadb
+        run_obj = metadb.get_run(run_id=run.id)
+
+        # From sync_metadata hook
+        assert run_obj.meta["accuracy"] == 0.95
+        assert run_obj.meta["loss"] == 0.05
+        assert run_obj.meta["num_epochs"] == 10
+
+        assert run_obj.status == Status.FAILED
+
+
+@pytest.mark.asyncio
+async def test_sync_metadata_with_none(test_org_id, test_user_id, test_team_id):
+    """Test that sync_metadata with None result doesn't update metadata"""
+    alpha.init(org_id=test_org_id, team_id=test_team_id, user_id=test_user_id)
+
+    async def task_with_none_result():
+        """Function that returns None"""
+        await asyncio.sleep(0.1)
+
+    async with CraftExperiment.start("test_hook_none_result") as exp:
+        run = exp.run(
+            task_with_none_result,
+            post_run_hooks=[PostRunHookFn.sync_metadata, PostRunHookFn.sync_status],
+        )
+        await exp.wait()
+
+        # Verify metadata was not updated
+        metadb = global_runtime().metadb
+        run_obj = metadb.get_run(run_id=run.id)
+
+        # Metadata should be None or empty (hook didn't update it)
+        assert run_obj.meta is None or run_obj.meta == {}
+        assert run_obj.status == Status.COMPLETED