Fix utilization_policy in profiles (#2385)

Fixes: #2383

Author: un-def

src/dstack/_internal/cli/services/configurators/run.py

@@ -95,6 +95,7 @@ def apply_configuration(
                 reservation=profile.reservation,
                 spot_policy=profile.spot_policy,
                 retry_policy=profile.retry_policy,
+                utilization_policy=profile.utilization_policy,
                 max_duration=profile.max_duration,
                 stop_duration=profile.stop_duration,
                 max_price=profile.max_price,

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

@@ -23,6 +23,7 @@
     ProfileRetryPolicy,
     RetryEvent,
     SpotPolicy,
+    UtilizationPolicy,
 )
 from dstack._internal.core.models.repos import AnyRunRepoData
 from dstack._internal.core.models.resources import Memory, ResourcesSpec
@@ -192,6 +193,7 @@ class JobSpec(CoreModel):
     single_branch: Optional[bool] = None
     max_duration: Optional[int]
     stop_duration: Optional[int] = None
+    utilization_policy: Optional[UtilizationPolicy] = None
     registry_auth: Optional[RegistryAuth]
     requirements: Requirements
     retry: Optional[Retry]

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

@@ -348,7 +348,7 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
             job_model.termination_reason = JobTerminationReason.GATEWAY_ERROR
 
     if job_model.status == JobStatus.RUNNING:
-        await _check_gpu_utilization(session, run_model, job_model)
+        await _check_gpu_utilization(session, job_model, job)
 
     job_model.last_processed_at = common_utils.get_current_datetime()
     await session.commit()
@@ -679,10 +679,8 @@ def _terminate_if_inactivity_duration_exceeded(
         )
 
 
-async def _check_gpu_utilization(
-    session: AsyncSession, run_model: RunModel, job_model: JobModel
-) -> None:
-    policy = RunSpec.__response__.parse_raw(run_model.run_spec).configuration.utilization_policy
+async def _check_gpu_utilization(session: AsyncSession, job_model: JobModel, job: Job) -> None:
+    policy = job.job_spec.utilization_policy
     if policy is None:
         return
     after = common_utils.get_current_datetime() - timedelta(seconds=policy.time_window)

src/dstack/_internal/server/services/jobs/configurators/base.py

@@ -13,7 +13,11 @@
     PythonVersion,
     RunConfigurationType,
 )
-from dstack._internal.core.models.profiles import DEFAULT_STOP_DURATION, SpotPolicy
+from dstack._internal.core.models.profiles import (
+    DEFAULT_STOP_DURATION,
+    SpotPolicy,
+    UtilizationPolicy,
+)
 from dstack._internal.core.models.runs import (
     AppSpec,
     JobSpec,
@@ -113,6 +117,7 @@ async def _get_job_spec(
             single_branch=self._single_branch(),
             max_duration=self._max_duration(),
             stop_duration=self._stop_duration(),
+            utilization_policy=self._utilization_policy(),
             registry_auth=self._registry_auth(),
             requirements=self._requirements(),
             retry=self._retry(),
@@ -201,6 +206,9 @@ def _stop_duration(self) -> Optional[int]:
         # pydantic validator ensures this is int
         return self.run_spec.merged_profile.stop_duration
 
+    def _utilization_policy(self) -> Optional[UtilizationPolicy]:
+        return self.run_spec.merged_profile.utilization_policy
+
     def _registry_auth(self) -> Optional[RegistryAuth]:
         return self.run_spec.configuration.registry_auth
 

src/dstack/api/_public/runs.py

@@ -25,6 +25,7 @@
     ProfileRetryPolicy,
     SpotPolicy,
     TerminationPolicy,
+    UtilizationPolicy,
 )
 from dstack._internal.core.models.repos.base import Repo
 from dstack._internal.core.models.resources import ResourcesSpec
@@ -485,6 +486,7 @@ def get_plan(
         resources: Optional[ResourcesSpec] = None,
         spot_policy: Optional[SpotPolicy] = None,
         retry_policy: Optional[ProfileRetryPolicy] = None,
+        utilization_policy: Optional[UtilizationPolicy] = None,
         max_duration: Optional[Union[int, str]] = None,
         max_price: Optional[float] = None,
         working_dir: Optional[str] = None,
@@ -535,6 +537,7 @@ def get_plan(
             spot_policy=spot_policy,
             retry=None,
             retry_policy=retry_policy,
+            utilization_policy=utilization_policy,
             max_duration=max_duration,
             stop_duration=stop_duration,
             max_price=max_price,

src/tests/_internal/server/routers/test_runs.py

@@ -198,6 +198,7 @@ def get_dev_env_run_plan_dict(
                     "single_branch": False,
                     "max_duration": None,
                     "stop_duration": 300,
+                    "utilization_policy": None,
                     "registry_auth": None,
                     "requirements": {
                         "resources": {
@@ -357,6 +358,7 @@ def get_dev_env_run_dict(
                     "single_branch": False,
                     "max_duration": None,
                     "stop_duration": 300,
+                    "utilization_policy": None,
                     "registry_auth": None,
                     "requirements": {
                         "resources": {