Use new helpers get_{run,job}_spec

Author: un-def

src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py

@@ -59,6 +59,7 @@
     emit_job_status_change_event,
     get_job_provisioning_data,
     get_job_runtime_data,
+    get_job_spec,
 )
 from dstack._internal.server.services.locking import get_locker
 from dstack._internal.server.services.logging import fmt
@@ -797,7 +798,7 @@ async def _detach_volumes_from_job_instance(
     jpd: JobProvisioningData,
     run_termination_reason: Optional[RunTerminationReason],
 ) -> _VolumeDetachResult:
-    job_spec = JobSpec.__response__.parse_raw(job_model.job_spec_data)
+    job_spec = get_job_spec(job_model)
     backend = await backends_services.get_project_backend_by_type(
         project=instance_model.project,
         backend_type=jpd.backend,

src/dstack/_internal/server/background/scheduled_tasks/probes.py

@@ -12,7 +12,7 @@
 from sqlalchemy.orm import joinedload
 
 from dstack._internal.core.errors import SSHError
-from dstack._internal.core.models.runs import JobSpec, JobStatus, ProbeSpec
+from dstack._internal.core.models.runs import JobStatus, ProbeSpec
 from dstack._internal.core.services.ssh.tunnel import (
     SSH_DEFAULT_OPTIONS,
     IPSocket,
@@ -21,6 +21,7 @@
 )
 from dstack._internal.server.db import get_db, get_session_ctx
 from dstack._internal.server.models import InstanceModel, JobModel, ProbeModel
+from dstack._internal.server.services.jobs import get_job_spec
 from dstack._internal.server.services.locking import get_locker
 from dstack._internal.server.services.logging import fmt
 from dstack._internal.server.services.ssh import container_ssh_tunnel
@@ -71,7 +72,7 @@ async def process_probes():
                 if probe.job.status != JobStatus.RUNNING:
                     probe.active = False
                 else:
-                    job_spec: JobSpec = JobSpec.__response__.parse_raw(probe.job.job_spec_data)
+                    job_spec = get_job_spec(probe.job)
                     probe_spec = job_spec.probes[probe.probe_num]
                     if probe_spec.until_ready and probe.success_streak >= probe_spec.ready_after:
                         probe.active = False
@@ -148,7 +149,7 @@ async def _get_service_replica_client(job: JobModel) -> AsyncGenerator[AsyncClie
         **SSH_DEFAULT_OPTIONS,
         "ConnectTimeout": str(int(SSH_CONNECT_TIMEOUT.total_seconds())),
     }
-    job_spec: JobSpec = JobSpec.__response__.parse_raw(job.job_spec_data)
+    job_spec = get_job_spec(job)
     with TemporaryDirectory() as temp_dir:
         app_socket_path = (Path(temp_dir) / "replica.sock").absolute()
         async with container_ssh_tunnel(

src/dstack/_internal/server/background/scheduled_tasks/running_jobs.py

@@ -34,7 +34,6 @@
     JobTerminationReason,
     ProbeSpec,
     Run,
-    RunSpec,
     RunStatus,
 )
 from dstack._internal.core.models.volumes import InstanceMountPoint, Volume, VolumeMountPoint
@@ -67,6 +66,7 @@
     find_job,
     get_job_attached_volumes,
     get_job_runtime_data,
+    get_job_spec,
     is_master_job,
     job_model_to_job_submission,
     switch_job_status,
@@ -82,6 +82,7 @@
 from dstack._internal.server.services.runner import client
 from dstack._internal.server.services.runner.ssh import runner_ssh_tunnel
 from dstack._internal.server.services.runs import (
+    get_run_spec,
     is_job_ready,
     run_model_to_run,
 )
@@ -732,7 +733,7 @@ def _process_provisioning_with_shim(
     Returns:
         is successful
     """
-    job_spec = JobSpec.__response__.parse_raw(job_model.job_spec_data)
+    job_spec = get_job_spec(job_model)
 
     shim_client = client.ShimClient(port=ports[DSTACK_SHIM_HTTP_PORT])
 
@@ -980,7 +981,7 @@ def _terminate_if_inactivity_duration_exceeded(
     job_model: JobModel,
     no_connections_secs: Optional[int],
 ) -> None:
-    conf = RunSpec.__response__.parse_raw(run_model.run_spec).configuration
+    conf = get_run_spec(run_model).configuration
     if not isinstance(conf, DevEnvironmentConfiguration) or not isinstance(
         conf.inactivity_duration, int
     ):

src/dstack/_internal/server/background/scheduled_tasks/runs.py

@@ -13,7 +13,6 @@
 from dstack._internal.core.models.profiles import RetryEvent, StopCriteria
 from dstack._internal.core.models.runs import (
     Job,
-    JobSpec,
     JobStatus,
     JobTerminationReason,
     Run,
@@ -33,6 +32,7 @@
 from dstack._internal.server.services import events
 from dstack._internal.server.services.jobs import (
     find_job,
+    get_job_spec,
     get_job_specs_from_run_spec,
     group_jobs_by_replica_latest,
     is_master_job,
@@ -527,7 +527,7 @@ async def _handle_run_replicas(
             if job.status.is_finished():
                 continue
             try:
-                job_spec = JobSpec.__response__.parse_raw(job.job_spec_data)
+                job_spec = get_job_spec(job)
                 existing_group_names.add(job_spec.replica_group)
             except Exception:
                 continue
@@ -643,7 +643,7 @@ async def _update_jobs_to_new_deployment_in_place(
         replica_group_name = None
 
         if replicas:
-            job_spec = JobSpec.__response__.parse_raw(job_models[0].job_spec_data)
+            job_spec = get_job_spec(job_models[0])
             replica_group_name = job_spec.replica_group
 
         # FIXME: Handle getting image configuration errors or skip it.
@@ -658,7 +658,7 @@ async def _update_jobs_to_new_deployment_in_place(
         )
         can_update_all_jobs = True
         for old_job_model, new_job_spec in zip(job_models, new_job_specs):
-            old_job_spec = JobSpec.__response__.parse_raw(old_job_model.job_spec_data)
+            old_job_spec = get_job_spec(old_job_model)
             if new_job_spec != old_job_spec:
                 can_update_all_jobs = False
                 break

src/dstack/_internal/server/background/scheduled_tasks/terminating_jobs.py

@@ -39,6 +39,7 @@
 from dstack._internal.server.services.jobs import (
     get_job_provisioning_data,
     get_job_runtime_data,
+    get_job_spec,
     switch_job_status,
 )
 from dstack._internal.server.services.locking import get_locker
@@ -356,7 +357,7 @@ async def _detach_volumes_from_job_instance(
     instance_model: InstanceModel,
     volume_models: list[VolumeModel],
 ) -> bool:
-    job_spec = JobSpec.__response__.parse_raw(job_model.job_spec_data)
+    job_spec = get_job_spec(job_model)
     backend = await backends_services.get_project_backend_by_type(
         project=project,
         backend_type=jpd.backend,

src/dstack/_internal/server/services/prometheus/custom_metrics.py

@@ -13,7 +13,7 @@
 from sqlalchemy.orm import aliased, joinedload
 
 from dstack._internal.core.models.instances import InstanceStatus
-from dstack._internal.core.models.runs import JobStatus, RunSpec, RunStatus
+from dstack._internal.core.models.runs import JobStatus, RunStatus
 from dstack._internal.server.models import (
     InstanceModel,
     JobMetricsPoint,
@@ -25,6 +25,7 @@
 )
 from dstack._internal.server.services.instances import get_instance_offer
 from dstack._internal.server.services.jobs import get_job_provisioning_data, get_job_runtime_data
+from dstack._internal.server.services.runs import get_run_spec
 from dstack._internal.utils.common import get_current_datetime
 
 
@@ -152,7 +153,7 @@ async def get_job_metrics(session: AsyncSession) -> Iterable[Metric]:
             price = jrd.offer.price
         gpus = resources.gpus
         cpus = resources.cpus
-        run_spec = RunSpec.__response__.parse_raw(job.run.run_spec)
+        run_spec = get_run_spec(job.run)
         labels = {
             "dstack_project_name": job.project.name,
             "dstack_user_name": job.run.user.name,

src/dstack/_internal/server/services/proxy/repo.py

@@ -12,9 +12,7 @@
 from dstack._internal.core.models.instances import SSHConnectionParams
 from dstack._internal.core.models.runs import (
     JobProvisioningData,
-    JobSpec,
     JobStatus,
-    RunSpec,
     RunStatus,
     ServiceSpec,
     get_service_port,
@@ -32,6 +30,8 @@
 from dstack._internal.proxy.lib.repo import BaseProxyRepo
 from dstack._internal.server.models import InstanceModel, JobModel, ProjectModel, RunModel
 from dstack._internal.server.services.instances import get_instance_remote_connection_info
+from dstack._internal.server.services.jobs import get_job_spec
+from dstack._internal.server.services.runs import get_run_spec
 from dstack._internal.server.settings import DEFAULT_SERVICE_CLIENT_MAX_BODY_SIZE
 from dstack._internal.utils.common import get_or_error
 
@@ -68,7 +68,7 @@ async def get_service(self, project_name: str, run_name: str) -> Optional[Servic
         if not len(jobs):
             return None
         run = jobs[0].run
-        run_spec = RunSpec.__response__.parse_raw(run.run_spec)
+        run_spec = get_run_spec(run)
         if not isinstance(run_spec.configuration, ServiceConfiguration):
             return None
         replicas = []
@@ -108,7 +108,7 @@ async def get_service(self, project_name: str, run_name: str) -> Optional[Servic
             if rci is not None and rci.ssh_proxy is not None:
                 ssh_head_proxy = rci.ssh_proxy
                 ssh_head_proxy_private_key = get_or_error(rci.ssh_proxy_keys)[0].private
-            job_spec: JobSpec = JobSpec.__response__.parse_raw(job.job_spec_data)
+            job_spec = get_job_spec(job)
             replica = Replica(
                 id=job.id.hex,
                 app_port=get_service_port(job_spec, run_spec.configuration),

src/dstack/_internal/server/services/runs/__init__.py

@@ -24,7 +24,6 @@
 from dstack._internal.core.models.runs import (
     ApplyRunPlanInput,
     Job,
-    JobSpec,
     JobStatus,
     JobSubmission,
     JobTerminationReason,
@@ -54,6 +53,7 @@
     check_can_attach_job_volumes,
     delay_job_instance_termination,
     get_job_configured_volumes,
+    get_job_spec,
     get_jobs_from_run_spec,
     job_model_to_job_submission,
     remove_job_spec_sensitive_info,
@@ -835,7 +835,7 @@ def _get_run_jobs_with_submissions(
                     submissions.append(job_submission)
             if job_model is not None:
                 # Use the spec from the latest submission. Submissions can have different specs
-                job_spec = JobSpec.__response__.parse_raw(job_model.job_spec_data)
+                job_spec = get_job_spec(job_model)
                 if not include_sensitive:
                     remove_job_spec_sensitive_info(job_spec)
                 jobs.append(Job(job_spec=job_spec, job_submissions=submissions))
@@ -861,7 +861,7 @@ def _get_run_status_message(run_model: RunModel) -> str:
     if run_model.status in [RunStatus.SUBMITTED, RunStatus.PENDING]:
         # Show `retrying` if any job caused the run to retry
         for job_models in job_models_grouped_by_job:
-            last_job_spec = JobSpec.__response__.parse_raw(job_models[-1].job_spec_data)
+            last_job_spec = get_job_spec(job_models[-1])
             retry_on_events = last_job_spec.retry.on_events if last_job_spec.retry else []
             last_job_termination_reason = _get_last_job_termination_reason(job_models)
             if (

src/dstack/_internal/server/services/runs/replicas.py

@@ -4,17 +4,19 @@
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from dstack._internal.core.models.configurations import ReplicaGroup
-from dstack._internal.core.models.runs import JobSpec, JobStatus, JobTerminationReason, RunSpec
+from dstack._internal.core.models.runs import JobStatus, JobTerminationReason, RunSpec
 from dstack._internal.server.models import JobModel, RunModel
 from dstack._internal.server.services import events
 from dstack._internal.server.services.jobs import (
+    get_job_spec,
     get_jobs_from_run_spec,
     group_jobs_by_replica_latest,
     switch_job_status,
 )
 from dstack._internal.server.services.logging import fmt
 from dstack._internal.server.services.runs import (
     create_job_model_for_new_submission,
+    get_run_spec,
     logger,
 )
 from dstack._internal.server.services.secrets import get_project_secrets_mapping
@@ -30,8 +32,8 @@ async def retry_run_replica_jobs(
     )
 
     # Determine replica group from existing job
-    run_spec = RunSpec.__response__.parse_raw(run_model.run_spec)
-    job_spec = JobSpec.__response__.parse_raw(latest_jobs[0].job_spec_data)
+    run_spec = get_run_spec(run_model)
+    job_spec = get_job_spec(latest_jobs[0])
     replica_group_name = job_spec.replica_group
 
     new_jobs = await get_jobs_from_run_spec(
@@ -86,7 +88,7 @@ async def scale_run_replicas(session: AsyncSession, run_model: RunModel, replica
     )
 
     active_replicas, inactive_replicas = build_replica_lists(run_model)
-    run_spec = RunSpec.__response__.parse_raw(run_model.run_spec)
+    run_spec = get_run_spec(run_model)
 
     if replicas_diff < 0:
         scale_down_replicas(session, active_replicas, abs(replicas_diff))
@@ -259,7 +261,7 @@ async def scale_run_replicas_per_group(
                 run_model=run_model,
                 group=group,
                 replicas_diff=group_diff,
-                run_spec=RunSpec.__response__.parse_raw(run_model.run_spec),
+                run_spec=get_run_spec(run_model),
                 active_replicas=active_replicas,
                 inactive_replicas=inactive_replicas,
             )
@@ -300,7 +302,7 @@ async def scale_run_replicas_for_group(
 
 
 def job_belongs_to_group(job: JobModel, group_name: str) -> bool:
-    job_spec = JobSpec.__response__.parse_raw(job.job_spec_data)
+    job_spec = get_job_spec(job)
     return job_spec.replica_group == group_name
 
 

src/dstack/_internal/server/services/services/__init__.py

@@ -32,7 +32,7 @@
     RouterType,
     SGLangServiceRouterConfig,
 )
-from dstack._internal.core.models.runs import JobSpec, Run, RunSpec, ServiceModelSpec, ServiceSpec
+from dstack._internal.core.models.runs import Run, RunSpec, ServiceModelSpec, ServiceSpec
 from dstack._internal.core.models.services import OpenAIChatModel
 from dstack._internal.proxy.gateway.const import SERVICE_ALREADY_REGISTERED_ERROR_TEMPLATE
 from dstack._internal.server import settings
@@ -322,7 +322,7 @@ async def register_replica(
             async with conn.client() as client:
                 await client.register_replica(
                     run=run,
-                    job_spec=JobSpec.__response__.parse_raw(job_model.job_spec_data),
+                    job_spec=jobs_services.get_job_spec(job_model),
                     job_submission=job_submission,
                     instance_project_ssh_private_key=instance_project_ssh_private_key,
                     ssh_head_proxy=ssh_head_proxy,