Use InstanceTerminationReason instead of strings

Use the `InstanceTerminationReason` enum as the
type of `InstanceModel.termination_reason`. To
handle old instances, automatically convert their
legacy termination reason strings to relevant
`InstanceTerminationReason` enum members on reads.

Author: jvstme

src/dstack/_internal/core/models/instances.py

@@ -15,6 +15,9 @@
 from dstack._internal.core.models.health import HealthStatus
 from dstack._internal.core.models.volumes import Volume
 from dstack._internal.utils.common import pretty_resources
+from dstack._internal.utils.logging import get_logger
+
+logger = get_logger(__name__)
 
 
 class Gpu(CoreModel):
@@ -256,14 +259,68 @@ def finished_statuses(cls) -> List["InstanceStatus"]:
 
 class InstanceTerminationReason(str, Enum):
     IDLE_TIMEOUT = "idle_timeout"
-    PROOVISIONING_TIMEOUT = "provisioning_timeout"
+    PROVISIONING_TIMEOUT = "provisioning_timeout"
     ERROR = "error"
     JOB_FINISHED = "job_finished"
     TERMINATION_TIMEOUT = "termination_timeout"
     STARTING_TIMEOUT = "starting_timeout"
     NO_OFFERS = "no_offers"
     MASTER_FAILED = "master_failed"
-    NO_BALANCE = "no_balance"
+    MAX_INSTANCES_LIMIT = "max_instances_limit"
+    NO_BALANCE = "no_balance"  # used in dstack Sky
+
+    @classmethod
+    def from_legacy_str(cls, v: str) -> "InstanceTerminationReason":
+        """
+        Convert legacy termination reason string to relevant termination reason enum.
+
+        dstack versions prior to 0.20.1 represented instance termination reasons as raw
+        strings. Such strings may still be stored in the database.
+        """
+
+        if v == "Idle timeout":
+            return cls.IDLE_TIMEOUT
+        if v in (
+            "Provisioning timeout expired",
+            "Proivisioning timeout expired",  # typo is intentional
+            "The proivisioning timeout expired",  # typo is intentional
+        ):
+            return cls.PROVISIONING_TIMEOUT
+        if v in (
+            "Unsupported private SSH key type",
+            "Failed to locate internal IP address on the given network",
+            "Specified internal IP not found among instance interfaces",
+            "Cannot split into blocks",
+            "Backend not available",
+            "Error while waiting for instance to become running",
+            "Empty profile, requirements or instance_configuration",
+            "Unable to locate the internal ip-address for the given network",
+            "Private SSH key is encrypted, password required",
+            "Cannot parse private key, key type is not supported",
+        ) or v.startswith("Error to parse profile, requirements or instance_configuration:"):
+            return cls.ERROR
+        if v in (
+            "All offers failed",
+            "No offers found",
+            "There were no offers found",
+            "Retry duration expired",
+            "The retry's duration expired",
+        ):
+            return cls.NO_OFFERS
+        if v == "Master instance failed to start":
+            return cls.MASTER_FAILED
+        if v == "Instance job finished":
+            return cls.JOB_FINISHED
+        if v == "Termination deadline":
+            return cls.TERMINATION_TIMEOUT
+        if v == "Instance has not become running in time":
+            return cls.STARTING_TIMEOUT
+        if v == "Fleet has too many instances":
+            return cls.MAX_INSTANCES_LIMIT
+        if v == "Low account balance":
+            return cls.NO_BALANCE
+        logger.warning("Unexpected instance termination reason string: %r", v)
+        return cls.ERROR
 
 
 class Instance(CoreModel):

src/dstack/_internal/server/background/tasks/process_fleets.py

@@ -8,7 +8,7 @@
 from sqlalchemy.orm import joinedload, load_only, selectinload
 
 from dstack._internal.core.models.fleets import FleetSpec, FleetStatus
-from dstack._internal.core.models.instances import InstanceStatus
+from dstack._internal.core.models.instances import InstanceStatus, InstanceTerminationReason
 from dstack._internal.server.db import get_db, get_session_ctx
 from dstack._internal.server.models import (
     FleetModel,
@@ -213,7 +213,8 @@ def _maintain_fleet_nodes_in_min_max_range(
                 break
             if instance.status in [InstanceStatus.IDLE]:
                 instance.status = InstanceStatus.TERMINATING
-                instance.termination_reason = "Fleet has too many instances"
+                instance.termination_reason = InstanceTerminationReason.MAX_INSTANCES_LIMIT
+                instance.termination_reason_message = "Fleet has too many instances"
                 nodes_redundant -= 1
                 logger.info(
                     "Terminating instance %s: %s",

src/dstack/_internal/server/background/tasks/process_instances.py

@@ -275,7 +275,7 @@ def _check_and_mark_terminating_if_idle_duration_expired(instance: InstanceModel
     delta = datetime.timedelta(seconds=idle_seconds)
     if idle_duration > delta:
         instance.status = InstanceStatus.TERMINATING
-        instance.termination_reason = InstanceTerminationReason.IDLE_TIMEOUT.value
+        instance.termination_reason = InstanceTerminationReason.IDLE_TIMEOUT
         logger.info(
             "Instance %s idle duration expired: idle time %ss. Terminating",
             instance.name,
@@ -311,7 +311,7 @@ async def _add_remote(instance: InstanceModel) -> None:
     retry_duration_deadline = instance.created_at + timedelta(seconds=PROVISIONING_TIMEOUT_SECONDS)
     if retry_duration_deadline < get_current_datetime():
         instance.status = InstanceStatus.TERMINATED
-        instance.termination_reason = InstanceTerminationReason.PROOVISIONING_TIMEOUT.value
+        instance.termination_reason = InstanceTerminationReason.PROVISIONING_TIMEOUT
         logger.warning(
             "Failed to start instance %s in %d seconds. Terminating...",
             instance.name,
@@ -334,7 +334,7 @@ async def _add_remote(instance: InstanceModel) -> None:
                 ssh_proxy_pkeys = None
         except (ValueError, PasswordRequiredException):
             instance.status = InstanceStatus.TERMINATED
-            instance.termination_reason = InstanceTerminationReason.ERROR.value
+            instance.termination_reason = InstanceTerminationReason.ERROR
             instance.termination_reason_message = "Unsupported private SSH key type"
             logger.warning(
                 "Failed to add instance %s: unsupported private SSH key type",
@@ -393,7 +393,7 @@ async def _add_remote(instance: InstanceModel) -> None:
         )
     if instance_network is not None and internal_ip is None:
         instance.status = InstanceStatus.TERMINATED
-        instance.termination_reason = InstanceTerminationReason.ERROR.value
+        instance.termination_reason = InstanceTerminationReason.ERROR
         instance.termination_reason_message = (
             "Failed to locate internal IP address on the given network"
         )
@@ -409,7 +409,7 @@ async def _add_remote(instance: InstanceModel) -> None:
     if internal_ip is not None:
         if not is_ip_among_addresses(ip_address=internal_ip, addresses=host_network_addresses):
             instance.status = InstanceStatus.TERMINATED
-            instance.termination_reason = InstanceTerminationReason.ERROR.value
+            instance.termination_reason = InstanceTerminationReason.ERROR
             instance.termination_reason_message = (
                 "Specified internal IP not found among instance interfaces"
             )
@@ -432,7 +432,7 @@ async def _add_remote(instance: InstanceModel) -> None:
         instance.total_blocks = blocks
     else:
         instance.status = InstanceStatus.TERMINATED
-        instance.termination_reason = InstanceTerminationReason.ERROR.value
+        instance.termination_reason = InstanceTerminationReason.ERROR
         instance.termination_reason_message = "Cannot split into blocks"
         logger.warning(
             "Failed to add instance %s: cannot split into blocks",
@@ -552,7 +552,7 @@ async def _create_instance(session: AsyncSession, instance: InstanceModel) -> No
         requirements = get_instance_requirements(instance)
     except ValidationError as e:
         instance.status = InstanceStatus.TERMINATED
-        instance.termination_reason = InstanceTerminationReason.ERROR.value
+        instance.termination_reason = InstanceTerminationReason.ERROR
         instance.termination_reason_message = (
             f"Error to parse profile, requirements or instance_configuration: {e}"
         )
@@ -679,17 +679,19 @@ async def _create_instance(session: AsyncSession, instance: InstanceModel) -> No
             )
         return
 
-    _mark_terminated(instance, InstanceTerminationReason.NO_OFFERS.value)
+    _mark_terminated(instance, InstanceTerminationReason.NO_OFFERS)
     if instance.fleet and is_fleet_master_instance(instance) and is_cloud_cluster(instance.fleet):
         # Do not attempt to deploy other instances, as they won't determine the correct cluster
         # backend, region, and placement group without a successfully deployed master instance
         for sibling_instance in instance.fleet.instances:
             if sibling_instance.id == instance.id:
                 continue
-            _mark_terminated(sibling_instance, InstanceTerminationReason.MASTER_FAILED.value)
+            _mark_terminated(sibling_instance, InstanceTerminationReason.MASTER_FAILED)
 
 
-def _mark_terminated(instance: InstanceModel, termination_reason: str) -> None:
+def _mark_terminated(
+    instance: InstanceModel, termination_reason: InstanceTerminationReason
+) -> None:
     instance.status = InstanceStatus.TERMINATED
     instance.termination_reason = termination_reason
     logger.info(
@@ -711,7 +713,7 @@ async def _check_instance(session: AsyncSession, instance: InstanceModel) -> Non
     ):
         # A busy instance could have no active jobs due to this bug: https://github.com/dstackai/dstack/issues/2068
         instance.status = InstanceStatus.TERMINATING
-        instance.termination_reason = InstanceTerminationReason.JOB_FINISHED.value
+        instance.termination_reason = InstanceTerminationReason.JOB_FINISHED
         logger.info(
             "Detected busy instance %s with finished job. Marked as TERMINATING",
             instance.name,
@@ -840,7 +842,7 @@ async def _check_instance(session: AsyncSession, instance: InstanceModel) -> Non
         deadline = instance.termination_deadline
         if get_current_datetime() > deadline:
             instance.status = InstanceStatus.TERMINATING
-            instance.termination_reason = InstanceTerminationReason.TERMINATION_TIMEOUT.value
+            instance.termination_reason = InstanceTerminationReason.TERMINATION_TIMEOUT
             logger.warning(
                 "Instance %s shim waiting timeout. Marked as TERMINATING",
                 instance.name,
@@ -869,7 +871,7 @@ async def _wait_for_instance_provisioning_data(
             "Instance %s failed because instance has not become running in time", instance.name
         )
         instance.status = InstanceStatus.TERMINATING
-        instance.termination_reason = InstanceTerminationReason.STARTING_TIMEOUT.value
+        instance.termination_reason = InstanceTerminationReason.STARTING_TIMEOUT
         return
 
     backend = await backends_services.get_project_backend_by_type(
@@ -882,7 +884,7 @@ async def _wait_for_instance_provisioning_data(
             instance.name,
         )
         instance.status = InstanceStatus.TERMINATING
-        instance.termination_reason = InstanceTerminationReason.ERROR.value
+        instance.termination_reason = InstanceTerminationReason.ERROR
         instance.termination_reason_message = "Backend not available"
         return
     try:
@@ -900,7 +902,7 @@ async def _wait_for_instance_provisioning_data(
             repr(e),
         )
         instance.status = InstanceStatus.TERMINATING
-        instance.termination_reason = InstanceTerminationReason.ERROR.value
+        instance.termination_reason = InstanceTerminationReason.ERROR
         instance.termination_reason_message = "Error while waiting for instance to become running"
     except Exception:
         logger.exception(

src/dstack/_internal/server/background/tasks/process_running_jobs.py

@@ -393,7 +393,7 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
                     if (
                         job_model.instance is not None
                         and job_model.instance.termination_reason
-                        == InstanceTerminationReason.NO_BALANCE.value
+                        == InstanceTerminationReason.NO_BALANCE
                     ):
                         # if instance was terminated due to no balance, set job termination reason accodingly
                         job_model.termination_reason = JobTerminationReason.NO_BALANCE

src/dstack/_internal/server/models.py

@@ -1,7 +1,7 @@
 import enum
 import uuid
 from datetime import datetime, timezone
-from typing import Callable, List, Optional, Union
+from typing import Callable, Generic, List, Optional, TypeVar, Union
 
 from sqlalchemy import (
     BigInteger,
@@ -30,7 +30,7 @@
 from dstack._internal.core.models.fleets import FleetStatus
 from dstack._internal.core.models.gateways import GatewayStatus
 from dstack._internal.core.models.health import HealthStatus
-from dstack._internal.core.models.instances import InstanceStatus
+from dstack._internal.core.models.instances import InstanceStatus, InstanceTerminationReason
 from dstack._internal.core.models.profiles import (
     DEFAULT_FLEET_TERMINATION_IDLE_TIME,
     TerminationPolicy,
@@ -141,26 +141,45 @@ def process_result_value(self, value: Optional[str], dialect) -> Optional[Decryp
             return DecryptedString(plaintext=None, decrypted=False, exc=e)
 
 
-class EnumAsString(TypeDecorator):
+E = TypeVar("E", bound=enum.Enum)
+
+
+class EnumAsString(TypeDecorator, Generic[E]):
     """
     A custom type decorator that stores enums as strings in the DB.
     """
 
     impl = String
     cache_ok = True
 
-    def __init__(self, enum_class: type[enum.Enum], *args, **kwargs):
+    def __init__(
+        self,
+        enum_class: type[E],
+        *args,
+        fallback_deserializer: Optional[Callable[[str], E]] = None,
+        **kwargs,
+    ):
+        """
+        Args:
+            enum_class: The enum class to be stored.
+            fallback_deserializer: An optional function used when the string
+                from the DB does not match any enum member name. If not
+                provided, an exception will be raised in such cases.
+        """
         self.enum_class = enum_class
+        self.fallback_deserializer = fallback_deserializer
         super().__init__(*args, **kwargs)
 
-    def process_bind_param(self, value: Optional[enum.Enum], dialect) -> Optional[str]:
+    def process_bind_param(self, value: Optional[E], dialect) -> Optional[str]:
         if value is None:
             return None
         return value.name
 
-    def process_result_value(self, value: Optional[str], dialect) -> Optional[enum.Enum]:
+    def process_result_value(self, value: Optional[str], dialect) -> Optional[E]:
         if value is None:
             return None
+        if value not in self.enum_class.__members__ and self.fallback_deserializer is not None:
+            return self.fallback_deserializer(value)
         return self.enum_class[value]
 
 
@@ -641,8 +660,16 @@ class InstanceModel(BaseModel):
 
     # instance termination handling
     termination_deadline: Mapped[Optional[datetime]] = mapped_column(NaiveDateTime)
-    # TODO: Migrate to EnumAsString(InstanceTerminationReason, 100) after enough releases to ensure backward compatibility
-    termination_reason: Mapped[Optional[str]] = mapped_column(String(4000))
+    # dstack versions prior to 0.20.1 represented instance termination reasons as raw strings.
+    # Such strings may still be stored in the database, so we are using a wide column (4000 chars)
+    # and a fallback deserializer to convert them to relevant enum members.
+    termination_reason: Mapped[Optional[InstanceTerminationReason]] = mapped_column(
+        EnumAsString(
+            InstanceTerminationReason,
+            4000,
+            fallback_deserializer=InstanceTerminationReason.from_legacy_str,
+        )
+    )
     termination_reason_message: Mapped[Optional[str]] = mapped_column(String(4000))
     # Deprecated since 0.19.22, not used
     health_status: Mapped[Optional[str]] = mapped_column(String(4000), deferred=True)