dstack 0.0.9__py3-none-any.whl → 0.20.7__py3-none-any.whl

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

@@ -0,0 +1,317 @@
+import json
+from typing import List, Optional, Tuple
+
+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.server.models import JobModel, RunModel
+from dstack._internal.server.services import events
+from dstack._internal.server.services.jobs import (
+    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,
+    logger,
+)
+from dstack._internal.server.services.secrets import get_project_secrets_mapping
+
+
+async def retry_run_replica_jobs(
+    session: AsyncSession, run_model: RunModel, latest_jobs: List[JobModel], *, only_failed: bool
+):
+    # FIXME: Handle getting image configuration errors or skip it.
+    secrets = await get_project_secrets_mapping(
+        session=session,
+        project=run_model.project,
+    )
+
+    # 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)
+    replica_group_name = job_spec.replica_group
+
+    new_jobs = await get_jobs_from_run_spec(
+        run_spec=run_spec,
+        secrets=secrets,
+        replica_num=latest_jobs[0].replica_num,
+        replica_group_name=replica_group_name,
+    )
+    assert len(new_jobs) == len(latest_jobs), (
+        "Changing the number of jobs within a replica is not yet supported"
+    )
+    for job_model, new_job in zip(latest_jobs, new_jobs):
+        if not (job_model.status.is_finished() or job_model.status == JobStatus.TERMINATING):
+            if only_failed:
+                # No need to resubmit, skip
+                continue
+            # The job is not finished, but we have to retry all jobs. Terminate it
+            job_model.termination_reason = JobTerminationReason.TERMINATED_BY_SERVER
+            job_model.termination_reason_message = "Replica is to be retried"
+            switch_job_status(session, job_model, JobStatus.TERMINATING)
+
+        new_job_model = create_job_model_for_new_submission(
+            run_model=run_model,
+            job=new_job,
+            status=JobStatus.SUBMITTED,
+        )
+        # dirty hack to avoid passing all job submissions
+        new_job_model.submission_num = job_model.submission_num + 1
+        session.add(new_job_model)
+        events.emit(
+            session,
+            f"Job created when re-running replica. Status: {new_job_model.status.upper()}",
+            actor=events.SystemActor(),
+            targets=[events.Target.from_model(new_job_model)],
+        )
+
+
+def is_replica_registered(jobs: list[JobModel]) -> bool:
+    # Only job_num=0 is supposed to receive service requests
+    return jobs[0].registered
+
+
+async def scale_run_replicas(session: AsyncSession, run_model: RunModel, replicas_diff: int):
+    if replicas_diff == 0:
+        return
+
+    logger.info(
+        "%s: scaling %s %s replica(s)",
+        fmt(run_model),
+        "UP" if replicas_diff > 0 else "DOWN",
+        abs(replicas_diff),
+    )
+
+    active_replicas, inactive_replicas = build_replica_lists(run_model)
+    run_spec = RunSpec.__response__.parse_raw(run_model.run_spec)
+
+    if replicas_diff < 0:
+        scale_down_replicas(session, active_replicas, abs(replicas_diff))
+    else:
+        await _scale_up_replicas(
+            session,
+            run_model,
+            inactive_replicas,
+            replicas_diff,
+            run_spec,
+            group_name=None,
+        )
+
+
+def build_replica_lists(
+    run_model: RunModel,
+    group_filter: Optional[str] = None,
+) -> Tuple[
+    List[Tuple[int, bool, int, List[JobModel]]], List[Tuple[int, bool, int, List[JobModel]]]
+]:
+    # lists of (importance, is_out_of_date, replica_num, jobs)
+    active_replicas: list[tuple[int, bool, int, list[JobModel]]] = []
+    inactive_replicas: list[tuple[int, bool, int, list[JobModel]]] = []
+
+    for replica_num, replica_jobs in group_jobs_by_replica_latest(run_model.jobs):
+        # Filter by group if specified
+        if group_filter is not None:
+            if not job_belongs_to_group(replica_jobs[0], group_filter):
+                continue
+
+        statuses = set(job.status for job in replica_jobs)
+        deployment_num = replica_jobs[0].deployment_num  # same for all jobs
+        is_out_of_date = deployment_num < run_model.deployment_num
+
+        if {JobStatus.TERMINATING, *JobStatus.finished_statuses()} & statuses:
+            # if there are any terminating or finished jobs, the replica is inactive
+            inactive_replicas.append((0, is_out_of_date, replica_num, replica_jobs))
+        elif JobStatus.SUBMITTED in statuses:
+            # if there are any submitted jobs, the replica is active and has the importance of 0
+            active_replicas.append((0, is_out_of_date, replica_num, replica_jobs))
+        elif {JobStatus.PROVISIONING, JobStatus.PULLING} & statuses:
+            # if there are any provisioning or pulling jobs, the replica is active and has the importance of 1
+            active_replicas.append((1, is_out_of_date, replica_num, replica_jobs))
+        elif not is_replica_registered(replica_jobs):
+            # all jobs are running, but not receiving traffic, the replica is active and has the importance of 2
+            active_replicas.append((2, is_out_of_date, replica_num, replica_jobs))
+        else:
+            # all jobs are running and ready, the replica is active and has the importance of 3
+            active_replicas.append((3, is_out_of_date, replica_num, replica_jobs))
+
+    # Sort by is_out_of_date (up-to-date first), importance (desc), and replica_num (asc)
+    active_replicas.sort(key=lambda r: (r[1], -r[0], r[2]))
+
+    return active_replicas, inactive_replicas
+
+
+def scale_down_replicas(
+    session: AsyncSession,
+    active_replicas: List[Tuple[int, bool, int, List[JobModel]]],
+    count: int,
+) -> None:
+    """Scale down by terminating the least important replicas"""
+    if count <= 0:
+        return
+
+    for _, _, _, replica_jobs in reversed(active_replicas[-count:]):
+        for job in replica_jobs:
+            if job.status.is_finished() or job.status == JobStatus.TERMINATING:
+                continue
+            job.termination_reason = JobTerminationReason.SCALED_DOWN
+            switch_job_status(session, job, JobStatus.TERMINATING, events.SystemActor())
+            # background task will process the job later
+
+
+async def _scale_up_replicas(
+    session: AsyncSession,
+    run_model: RunModel,
+    inactive_replicas: List[Tuple[int, bool, int, List[JobModel]]],
+    replicas_diff: int,
+    run_spec: RunSpec,
+    group_name: Optional[str] = None,
+) -> None:
+    """Scale up by retrying inactive replicas and creating new ones"""
+    if replicas_diff <= 0:
+        return
+
+    scheduled_replicas = 0
+
+    # Retry inactive replicas first
+    for _, _, _, replica_jobs in inactive_replicas:
+        if scheduled_replicas == replicas_diff:
+            break
+        await retry_run_replica_jobs(session, run_model, replica_jobs, only_failed=False)
+        scheduled_replicas += 1
+
+    # Create new replicas
+    if scheduled_replicas < replicas_diff:
+        secrets = await get_project_secrets_mapping(
+            session=session,
+            project=run_model.project,
+        )
+
+        max_replica_num = max((job.replica_num for job in run_model.jobs), default=-1)
+
+        new_replicas_needed = replicas_diff - scheduled_replicas
+        for i in range(new_replicas_needed):
+            new_replica_num = max_replica_num + 1 + i
+            jobs = await get_jobs_from_run_spec(
+                run_spec=run_spec,
+                secrets=secrets,
+                replica_num=new_replica_num,
+                replica_group_name=group_name,
+            )
+            for job in jobs:
+                job_model = create_job_model_for_new_submission(
+                    run_model=run_model,
+                    job=job,
+                    status=JobStatus.SUBMITTED,
+                )
+                session.add(job_model)
+                events.emit(
+                    session,
+                    f"Job created on new replica submission. Status: {job_model.status.upper()}",
+                    actor=events.SystemActor(),
+                    targets=[events.Target.from_model(job_model)],
+                )
+                # Append to run_model.jobs so that when processing later replica groups in the same
+                # transaction, run_model.jobs includes jobs from previously processed groups.
+                run_model.jobs.append(job_model)
+
+
+async def scale_run_replicas_per_group(
+    session: AsyncSession,
+    run_model: RunModel,
+    replicas: List[ReplicaGroup],
+) -> None:
+    """Scale each replica group independently"""
+    if not replicas:
+        return
+
+    desired_replica_counts = (
+        json.loads(run_model.desired_replica_counts) if run_model.desired_replica_counts else {}
+    )
+
+    for group in replicas:
+        assert group.name is not None, "Group name is always set"
+        group_desired = desired_replica_counts.get(group.name, group.count.min or 0)
+
+        # Build replica lists filtered by this group
+        active_replicas, inactive_replicas = build_replica_lists(
+            run_model=run_model, group_filter=group.name
+        )
+
+        # Count active replicas
+        active_group_count = len(active_replicas)
+        group_diff = group_desired - active_group_count
+
+        if group_diff != 0:
+            # Check if rolling deployment is in progress for THIS GROUP
+
+            group_has_out_of_date = has_out_of_date_replicas(run_model, group_filter=group.name)
+
+            # During rolling deployment, don't scale down old replicas
+            # Let rolling deployment handle stopping old replicas
+            if group_diff < 0 and group_has_out_of_date:
+                # Skip scaling down during rolling deployment
+                continue
+            await scale_run_replicas_for_group(
+                session=session,
+                run_model=run_model,
+                group=group,
+                replicas_diff=group_diff,
+                run_spec=RunSpec.__response__.parse_raw(run_model.run_spec),
+                active_replicas=active_replicas,
+                inactive_replicas=inactive_replicas,
+            )
+
+
+async def scale_run_replicas_for_group(
+    session: AsyncSession,
+    run_model: RunModel,
+    group: ReplicaGroup,
+    replicas_diff: int,
+    run_spec: RunSpec,
+    active_replicas: List[Tuple[int, bool, int, List[JobModel]]],
+    inactive_replicas: List[Tuple[int, bool, int, List[JobModel]]],
+) -> None:
+    """Scale a specific replica group up or down"""
+    if replicas_diff == 0:
+        return
+
+    logger.info(
+        "%s: scaling %s %s replica(s) for group '%s'",
+        fmt(run_model),
+        "UP" if replicas_diff > 0 else "DOWN",
+        abs(replicas_diff),
+        group.name,
+    )
+
+    if replicas_diff < 0:
+        scale_down_replicas(session, active_replicas, abs(replicas_diff))
+    else:
+        await _scale_up_replicas(
+            session=session,
+            run_model=run_model,
+            inactive_replicas=inactive_replicas,
+            replicas_diff=replicas_diff,
+            run_spec=run_spec,
+            group_name=group.name,
+        )
+
+
+def job_belongs_to_group(job: JobModel, group_name: str) -> bool:
+    job_spec = JobSpec.__response__.parse_raw(job.job_spec_data)
+    return job_spec.replica_group == group_name
+
+
+def has_out_of_date_replicas(run: RunModel, group_filter: Optional[str] = None) -> bool:
+    for job in run.jobs:
+        # Filter jobs by group if specified
+        if group_filter is not None:
+            if not job_belongs_to_group(job, group_filter):
+                continue
+        if job.deployment_num < run.deployment_num and not (
+            job.status.is_finished() or job.termination_reason == JobTerminationReason.SCALED_DOWN
+        ):
+            return True
+    return False

dstack/_internal/server/services/runs/spec.py

@@ -0,0 +1,191 @@
+from dstack._internal.core.errors import ServerClientError
+from dstack._internal.core.models.configurations import RUN_PRIORITY_DEFAULT, ServiceConfiguration
+from dstack._internal.core.models.repos.virtual import DEFAULT_VIRTUAL_REPO_ID, VirtualRunRepoData
+from dstack._internal.core.models.runs import LEGACY_REPO_DIR, AnyRunConfiguration, RunSpec
+from dstack._internal.core.models.volumes import InstanceMountPoint
+from dstack._internal.core.services import validate_dstack_resource_name
+from dstack._internal.core.services.diff import diff_models
+from dstack._internal.server import settings
+from dstack._internal.server.models import UserModel
+from dstack._internal.server.services.docker import is_valid_docker_volume_target
+from dstack._internal.server.services.resources import set_resources_defaults
+from dstack._internal.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+_UPDATABLE_SPEC_FIELDS = ["configuration_path", "configuration"]
+_TYPE_SPECIFIC_UPDATABLE_SPEC_FIELDS = {
+    "service": [
+        # rolling deployment
+        "repo_data",
+        "repo_code_hash",
+        "file_archives",
+        "working_dir",
+    ],
+}
+_CONF_UPDATABLE_FIELDS = ["priority"]
+_TYPE_SPECIFIC_CONF_UPDATABLE_FIELDS = {
+    "dev-environment": ["inactivity_duration"],
+    "service": [
+        # in-place
+        "replicas",
+        "scaling",
+        # rolling deployment
+        # NOTE: keep this list in sync with the "Rolling deployment" section in services.md
+        "port",
+        "probes",
+        "resources",
+        "volumes",
+        "docker",
+        "files",
+        "image",
+        "user",
+        "privileged",
+        "entrypoint",
+        "working_dir",
+        "python",
+        "nvcc",
+        "single_branch",
+        "env",
+        "shell",
+        "commands",
+    ],
+}
+
+
+def validate_run_spec_and_set_defaults(
+    user: UserModel, run_spec: RunSpec, legacy_repo_dir: bool = False
+):
+    # This function may set defaults for null run_spec values,
+    # although most defaults are resolved when building job_spec
+    # so that we can keep both the original user-supplied value (null in run_spec)
+    # and the default in job_spec.
+    # If a property is stored in job_spec - resolve the default there.
+    # Server defaults are preferable over client defaults so that
+    # the defaults depend on the server version, not the client version.
+    if run_spec.run_name is not None:
+        validate_dstack_resource_name(run_spec.run_name)
+    for mount_point in run_spec.configuration.volumes:
+        if not is_valid_docker_volume_target(mount_point.path):
+            raise ServerClientError(f"Invalid volume mount path: {mount_point.path}")
+    if run_spec.repo_id is None and run_spec.repo_data is not None:
+        raise ServerClientError("repo_data must not be set if repo_id is not set")
+    if run_spec.repo_id is not None and run_spec.repo_data is None:
+        raise ServerClientError("repo_id must not be set if repo_data is not set")
+    # Some run_spec parameters have to be set here and not in the model defaults since
+    # the client may not pass them or pass null, but they must be always present, e.g. for runner.
+    if run_spec.repo_id is None:
+        run_spec.repo_id = DEFAULT_VIRTUAL_REPO_ID
+    if run_spec.repo_data is None:
+        run_spec.repo_data = VirtualRunRepoData()
+    if (
+        run_spec.merged_profile.utilization_policy is not None
+        and run_spec.merged_profile.utilization_policy.time_window
+        > settings.SERVER_METRICS_RUNNING_TTL_SECONDS
+    ):
+        raise ServerClientError(
+            f"Maximum utilization_policy.time_window is {settings.SERVER_METRICS_RUNNING_TTL_SECONDS}s"
+        )
+    if isinstance(run_spec.configuration, ServiceConfiguration):
+        if run_spec.merged_profile.schedule and all(
+            group.count.min == 0 for group in run_spec.configuration.replica_groups
+        ):
+            raise ServerClientError(
+                "Scheduled services with autoscaling to zero are not supported"
+            )
+        if len(run_spec.configuration.probes) > settings.MAX_PROBES_PER_JOB:
+            raise ServerClientError(
+                f"Cannot configure more than {settings.MAX_PROBES_PER_JOB} probes"
+            )
+        if any(
+            p.timeout is not None and p.timeout > settings.MAX_PROBE_TIMEOUT
+            for p in run_spec.configuration.probes
+        ):
+            raise ServerClientError(
+                f"Probe timeout cannot be longer than {settings.MAX_PROBE_TIMEOUT}s"
+            )
+    if run_spec.configuration.priority is None:
+        run_spec.configuration.priority = RUN_PRIORITY_DEFAULT
+    set_resources_defaults(run_spec.configuration.resources)
+    if run_spec.ssh_key_pub is None:
+        if user.ssh_public_key:
+            run_spec.ssh_key_pub = user.ssh_public_key
+        else:
+            raise ServerClientError("ssh_key_pub must be set if the user has no ssh_public_key")
+    if run_spec.configuration.working_dir is None and legacy_repo_dir:
+        run_spec.configuration.working_dir = LEGACY_REPO_DIR
+
+
+def check_can_update_run_spec(current_run_spec: RunSpec, new_run_spec: RunSpec):
+    spec_diff = diff_models(current_run_spec, new_run_spec)
+    changed_spec_fields = list(spec_diff.keys())
+    updatable_spec_fields = _UPDATABLE_SPEC_FIELDS + _TYPE_SPECIFIC_UPDATABLE_SPEC_FIELDS.get(
+        new_run_spec.configuration.type, []
+    )
+    for key in changed_spec_fields:
+        if key not in updatable_spec_fields:
+            raise ServerClientError(
+                f"Failed to update fields {changed_spec_fields}."
+                f" Can only update {updatable_spec_fields}."
+            )
+    # We don't allow update if the order of archives has been changed, as even if the archives
+    # are the same (the same id => hash => content and the same container path), the order of
+    # unpacking matters when one path is a subpath of another.
+    ignore_files = current_run_spec.file_archives == new_run_spec.file_archives
+    _check_can_update_configuration(
+        current_run_spec.configuration, new_run_spec.configuration, ignore_files
+    )
+
+
+def can_update_run_spec(current_run_spec: RunSpec, new_run_spec: RunSpec) -> bool:
+    try:
+        check_can_update_run_spec(current_run_spec, new_run_spec)
+    except ServerClientError as e:
+        logger.debug("Run cannot be updated: %s", repr(e))
+        return False
+    return True
+
+
+def get_nodes_required_num(run_spec: RunSpec) -> int:
+    nodes_required_num = 1
+    if run_spec.configuration.type == "task":
+        nodes_required_num = run_spec.configuration.nodes
+    elif run_spec.configuration.type == "service":
+        nodes_required_num = sum(
+            group.count.min or 0 for group in run_spec.configuration.replica_groups
+        )
+    return nodes_required_num
+
+
+def check_run_spec_requires_instance_mounts(run_spec: RunSpec) -> bool:
+    return any(
+        isinstance(mp, InstanceMountPoint) and not mp.optional
+        for mp in run_spec.configuration.volumes
+    )
+
+
+def _check_can_update_configuration(
+    current: AnyRunConfiguration, new: AnyRunConfiguration, ignore_files: bool
+) -> None:
+    if current.type != new.type:
+        raise ServerClientError(
+            f"Configuration type changed from {current.type} to {new.type}, cannot update"
+        )
+    updatable_fields = _CONF_UPDATABLE_FIELDS + _TYPE_SPECIFIC_CONF_UPDATABLE_FIELDS.get(
+        new.type, []
+    )
+    if ignore_files:
+        # We ignore files diff if the file archives are the same. It allows the user to move
+        # local files/dirs as long as their name(*), content, and the container path stay the same.
+        # (*) We could also ignore local name changes if the names didn't change in the tarballs.
+        # Currently, the client preserves the original file/dir name it the tarball, but it could
+        # use some generic names like "file"/"directory" instead.
+        updatable_fields.append("files")
+    diff = diff_models(current, new)
+    changed_fields = list(diff.keys())
+    for key in changed_fields:
+        if key not in updatable_fields:
+            raise ServerClientError(
+                f"Failed to update fields {changed_fields}. Can only update {updatable_fields}"
+            )