feat: add interrupted status

alphatrion/experiment/base.py

@@ -13,6 +13,7 @@
 from alphatrion.runtime.contextvars import current_exp_id
 from alphatrion.runtime.runtime import global_runtime
 from alphatrion.storage import runtime as storage_runtime
+from alphatrion.storage import sql_models
 from alphatrion.storage.sql_models import Status
 from alphatrion.types import CallableEntry, PostRunHookFn
 from alphatrion.utils import context
@@ -168,6 +169,7 @@ class Experiment(ABC):
         # The end status, None, Err or Cancelled.
         "_end_status",
         "_stopped",
+        "_received_signal",
         "_signal_task",
     )
 
@@ -187,6 +189,7 @@ def __init__(self, config: ExperimentConfig | None = None):
         self._start_waiting = False
         self._end_status = None
         self._stopped = asyncio.Event()
+        self._received_signal: int | None = None
         self._signal_task: asyncio.Task | None = None
 
     async def __aenter__(self):
@@ -199,7 +202,6 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
             self.done_with_err()
         else:
             self.done()
-        self._end_status = None
 
         if self._signal_task:
             # Already done, will not update the status again.
@@ -225,22 +227,7 @@ def _start(
             name=name, team_id=self._runtime.team_id
         )
 
-        # Just in case of kubernetes pod restarts, we want to make sure the experiment
-        # can be resumed if it is not completed, instead of creating a new experiment
-        # with the same name. If the experiment is already completed, we raise an error
-        # to avoid confusion.
-        if exp_obj and exp_obj.status != Status.COMPLETED:
-            self._id = exp_obj.uuid
-            self._runtime._metadb.update_experiment(
-                experiment_id=self._id,
-                status=Status.RUNNING,
-            )
-        elif exp_obj and exp_obj.status == Status.COMPLETED:
-            raise RuntimeError(
-                f"Experiment with name '{name}' already exists and is completed. \
-                Please choose a different name or delete the existing experiment."
-            )
-        else:
+        if not exp_obj:
             self._id = self._runtime._metadb.create_experiment(
                 name=name,
                 org_id=self._runtime._org_id,
@@ -253,6 +240,21 @@ def _start(
                 params=params,
                 status=Status.RUNNING,
             )
+        # Just in case of kubernetes pod restarts, we want to make sure the experiment
+        # can be resumed if it is not terminated, instead of creating a new experiment
+        # with the same name. If the experiment is already terminated, we raise an error
+        # to avoid confusion.
+        elif exp_obj.status not in sql_models.TERMINAL_STATUS:
+            self._id = exp_obj.uuid
+            self._runtime._metadb.update_experiment(
+                experiment_id=self._id,
+                status=Status.RUNNING,
+            )
+        else:
+            raise RuntimeError(
+                f"Experiment with name '{name}' already exists and is terminated with status '{exp_obj.status}'. \
+                        Please choose a different name or delete the existing experiment."
+            )
 
         self._context = context.Context(
             cancel_func=self._stop,
@@ -397,11 +399,19 @@ def done(self):
         self._cancel()
 
     def done_with_err(self):
-        self._end_status = "Err"
+        self._end_status = Status.FAILED
         self.done()
 
     def done_with_cancel(self):
-        self._end_status = "Cancelled"
+        self._end_status = Status.CANCELLED
+        self.done()
+
+    def done_with_interrupt(self):
+        self._end_status = Status.INTERRUPTED
+        self.done()
+
+    def done_with_abort(self):
+        self._end_status = Status.ABORTED
         self.done()
 
     def _cancel(self):
@@ -422,10 +432,8 @@ def _stop(self):
             ).total_seconds()
 
             status = Status.COMPLETED
-            if self._end_status == "Err":
-                status = Status.FAILED
-            elif self._end_status == "Cancelled":
-                status = Status.CANCELLED
+            if self._end_status is not None:
+                status = self._end_status
 
             self._runtime.metadb.update_experiment(
                 experiment_id=self._id, status=status, duration=duration
@@ -500,15 +508,22 @@ def _start_signal_handlers(self):
         loop = asyncio.get_running_loop()
 
         # Handle SIGINT and SIGTERM to allow graceful shutdown.
-        # Make sure to call done() on receiving the signal.
+        # SIGINT (Ctrl+C) → CANCELLED (user intent to stop)
+        # SIGTERM (kill/K8s) → INTERRUPTED (resumable)
         for sig in (signal.SIGINT, signal.SIGTERM):
-            loop.add_signal_handler(sig, self._on_signal)
+            loop.add_signal_handler(sig, lambda s=sig: self._on_signal(s))
 
         return asyncio.create_task(self._wait_for_stop())
 
-    def _on_signal(self):
+    def _on_signal(self, sig: int):
+        self._received_signal = sig
         self._stopped.set()
 
     async def _wait_for_stop(self):
         await self._stopped.wait()
-        self.done_with_cancel()
+
+        # Distinguish between SIGINT (user Ctrl+C) and SIGTERM (system/K8s)
+        if self._received_signal == signal.SIGINT:
+            self.done_with_cancel()  # User explicitly stopped with Ctrl+C
+        else:
+            self.done_with_interrupt()  # SIGTERM or other, treat as resumable

alphatrion/server/graphql/resolvers.py

@@ -2027,11 +2027,13 @@ def delete_datasets(
         return True
 
     @staticmethod
-    def abort_experiment(
+    def stop_experiment(
         info: Info[GraphQLContext, None], experiment_id: strawberry.ID
     ) -> Experiment:
-        """Abort an experiment by changing its status to ABORTED.
-        Only works if the experiment is in PENDING status."""
+        """Stop an experiment.
+        - PENDING → ABORTED (validation failure before start)
+        - RUNNING → CANCELLED (user-initiated stop)
+        """
 
         user_id = uuid.UUID(info.context.user_id)
         experiment_id_uuid = uuid.UUID(experiment_id)
@@ -2048,23 +2050,27 @@ def abort_experiment(
                 "Not allowed to update experiment in team that user does not belong to"
             )
 
-        # Only abort if experiment is in PENDING status
-        if exp.status != Status.PENDING:
+        # Determine target status based on current status
+        if exp.status == Status.PENDING:
+            new_status = Status.ABORTED
+        elif exp.status == Status.RUNNING:
+            new_status = Status.CANCELLED
+        else:
             raise RuntimeError(
-                f"Cannot abort experiment with status '{StatusMap[Status(exp.status)]}'. "
-                "Only experiments in PENDING status can be aborted."
+                f"Cannot stop experiment with status '{StatusMap[Status(exp.status)]}'. "
+                "Only experiments in PENDING or RUNNING status can be stopped."
             )
 
-        # Update status to ABORTED
+        # Update status
         metadb.update_experiment(
             experiment_id=experiment_id_uuid,
-            status=Status.ABORTED,
+            status=new_status,
         )
 
         # Get the updated experiment
         updated_exp = metadb.get_experiment(experiment_id=experiment_id_uuid)
         if not updated_exp:
-            raise RuntimeError("Failed to retrieve aborted experiment")
+            raise RuntimeError("Failed to retrieve stopped experiment")
 
         return Experiment(
             id=updated_exp.uuid,

alphatrion/server/graphql/schema.py

@@ -302,10 +302,10 @@ def update_experiment(
         return GraphQLMutations.update_experiment(info=info, input=input)
 
     @strawberry.mutation
-    def abort_experiment(
+    def stop_experiment(
         self, experiment_id: strawberry.ID, info: Info[GraphQLContext, None]
     ) -> Experiment:
-        return GraphQLMutations.abort_experiment(info=info, experiment_id=experiment_id)
+        return GraphQLMutations.stop_experiment(info=info, experiment_id=experiment_id)
 
 
 schema = strawberry.Schema(query=Query, mutation=Mutation)

alphatrion/server/graphql/types.py

@@ -148,13 +148,25 @@ def teams(self, info: Info) -> list[Team] | None:
 
 
 class GraphQLStatus(Enum):
+    """Status enum shared between Experiments and Runs.
+
+    Run status is a subset: RUNNING, COMPLETED, CANCELLED, FAILED.
+    Experiment status includes all values.
+
+    Terminal states (cannot resume): COMPLETED, CANCELLED, ABORTED
+    Recoverable states (can resume): FAILED, INTERRUPTED
+    """
+
     UNKNOWN = "UNKNOWN"
     PENDING = "PENDING"
     RUNNING = "RUNNING"
-    CANCELLED = "CANCELLED"
     COMPLETED = "COMPLETED"
-    FAILED = "FAILED"
-    ABORTED = "ABORTED"
+    CANCELLED = "CANCELLED"  # User intentionally stopped, cannot resume
+    FAILED = "FAILED"  # Recoverable: can retry after fixing error
+    ABORTED = "ABORTED"  # Programmatically aborted, cannot resume
+    INTERRUPTED = (
+        "INTERRUPTED"  # Recoverable: system stopped (e.g. preemption), can resume
+    )
 
 
 GraphQLStatusEnum = strawberry.enum(GraphQLStatus)

alphatrion/storage/sql_models.py

@@ -19,14 +19,16 @@
 Base = declarative_base()
 
 
+# Status includes both run and experiment status, run status is a subset of experiment status.
 class Status(enum.IntEnum):
     UNKNOWN = 0
     PENDING = 1
     RUNNING = 2
     COMPLETED = 9
-    CANCELLED = 10
-    FAILED = 11
-    ABORTED = 12
+    CANCELLED = 10  # User intentionally stopped, cannot resume
+    FAILED = 11  # Recoverable: can retry after fixing error
+    ABORTED = 12  # Programmatically aborted, cannot resume
+    INTERRUPTED = 13  # Recoverable: system stopped (e.g. preemption), can resume
 
 
 StatusMap = {
@@ -37,9 +39,20 @@ class Status(enum.IntEnum):
     Status.COMPLETED: "COMPLETED",
     Status.FAILED: "FAILED",
     Status.ABORTED: "ABORTED",
+    Status.INTERRUPTED: "INTERRUPTED",
 }
 
-FINISHED_STATUS = [Status.COMPLETED, Status.FAILED, Status.CANCELLED, Status.ABORTED]
+# Terminal states: cannot transition to any other state
+TERMINAL_STATUS = [Status.COMPLETED, Status.CANCELLED, Status.ABORTED]
+
+# Recoverable states: can transition back to RUNNING
+RECOVERABLE_STATUS = [Status.FAILED, Status.INTERRUPTED]
+
+# All finished states (terminal + recoverable)
+FINISHED_STATUS = TERMINAL_STATUS + RECOVERABLE_STATUS
+
+# Run status is a subset of experiment status to make it simpler.
+RUN_STATUS = [Status.RUNNING, Status.COMPLETED, Status.CANCELLED, Status.FAILED]
 
 
 class Organization(Base):

docs/architecture/exp_state_machine.md

@@ -0,0 +1,98 @@
+# Experiment Status State Machine
+
+## Overview
+
+The `Status` enum is shared between **Experiments** and **Runs**:
+
+```python
+class Status(enum.IntEnum):
+    UNKNOWN = 0
+    PENDING = 1
+    RUNNING = 2
+    COMPLETED = 9
+    CANCELLED = 10    # User intentionally stopped, cannot resume
+    FAILED = 11       # Recoverable: can retry after fixing error
+    ABORTED = 12      # Programmatically aborted, cannot resume
+    INTERRUPTED = 13  # Recoverable: system stopped (e.g. preemption), can resume
+```
+
+**Experiment statuses**: All values above
+**Run statuses**: Subset only (RUNNING, COMPLETED, CANCELLED, FAILED)
+
+## State Diagram
+
+```
+                             ┌─────────────┐
+                             │ Experiment  │
+                             └──────┬──────┘
+                                    │
+                                    ▼
+                             ┌─────────────┐
+                  ┌──────────│   PENDING   │
+                  │          └──────┬──────┘
+                  │                 │
+                  │ abort()         │ start()
+                  │                 │
+                  ▼                 ▼
+            ┌────────────┐   ┌─────────────┐
+            │  ABORTED   │   │   RUNNING   │ ◀───────────────────────────────────────────┐
+            └────────────┘   └──────┬──────┘                                             │
+             (terminal)             │                                                    │
+                                    │                                                    │
+        ┌───────────────────────────┼───────────────────────────────────────┐            │
+        │                           │                 │                     │            │
+        │                           │                 │                     │            │
+cancel()/signal(SIGINT)           done()           exception          signal(SIGTERM)    │
+        │                           │            (unintentional)            │            │
+        ▼                           ▼                 ▼                     ▼            │
+  ┌────────────┐             ┌────────────┐     ┌────────────┐      ┌──────────────┐     │
+  │ CANCELLED  │             │ COMPLETED  │     │   FAILED   │      │ INTERRUPTED  │     │
+  └────────────┘             └────────────┘     └─────┬──────┘      └──────┬───────┘     │
+   (terminal)                 (terminal)              │                    │             │
+                                                      │                    │             │
+                                                      └────────────────────┘             │
+                                                                │              resume    │
+                                                                └────────────────────────┘
+
+
+────────────────────────────────────────────────────────────────────────────────────────
+Terminal States (No Resume):
+  • COMPLETED   - Success, finished normally
+  • CANCELLED   - User stopped permanently (SIGINT/Ctrl+C or stopExperiment API)
+  • ABORTED     - Validation failed before starting (only from PENDING state)
+
+Recoverable States (Can Resume):
+  • FAILED      - Exception/error during execution, can retry after fix
+  • INTERRUPTED - System stopped (SIGTERM/K8s pod termination), can resume after restart
+────────────────────────────────────────────────────────────────────────────────────────
+```
+
+## State Categories
+
+| Category | States | Can Resume? | Description |
+|----------|--------|-------------|-------------|
+| **Active** | PENDING, RUNNING | N/A | Normal execution flow |
+| **Terminal** | COMPLETED, CANCELLED, ABORTED | ❌ No | Final states, immutable |
+| **Recoverable** | FAILED, INTERRUPTED | ✅ Yes | Can resume to RUNNING |
+
+### Semantics
+
+**Resumability is the key distinction:**
+
+- **COMPLETED**: Success, nothing more to do
+- **CANCELLED**: User decided to stop permanently (SIGINT in local dev, or stopExperiment API)
+- **ABORTED**: Validation failed before starting (only from PENDING via stopExperiment API)
+- **FAILED**: Unintentional error, retry after fixing
+- **INTERRUPTED**: System event (SIGTERM from K8s pod termination, preemption), resume when restarted
+
+## Implementation
+
+### Constants
+
+```python
+# In sql_models.py
+TERMINAL_STATUS = [Status.COMPLETED, Status.CANCELLED, Status.ABORTED]
+RECOVERABLE_STATUS = [Status.FAILED, Status.INTERRUPTED]
+FINISHED_STATUS = TERMINAL_STATUS + RECOVERABLE_STATUS
+RUN_STATUS = [Status.RUNNING, Status.COMPLETED, Status.CANCELLED, Status.FAILED]
+```