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

dstack/_internal/server/background/tasks/process_metrics.py

@@ -0,0 +1,172 @@
+import asyncio
+import json
+from typing import Dict, List, Optional
+
+from sqlalchemy import Delete, delete, select
+from sqlalchemy.orm import joinedload
+
+from dstack._internal.core.consts import DSTACK_RUNNER_HTTP_PORT
+from dstack._internal.core.models.runs import JobStatus
+from dstack._internal.server import settings
+from dstack._internal.server.db import get_session_ctx
+from dstack._internal.server.models import InstanceModel, JobMetricsPoint, JobModel, ProjectModel
+from dstack._internal.server.schemas.runner import MetricsResponse
+from dstack._internal.server.services.instances import get_instance_ssh_private_keys
+from dstack._internal.server.services.jobs import get_job_provisioning_data, get_job_runtime_data
+from dstack._internal.server.services.runner import client
+from dstack._internal.server.services.runner.ssh import runner_ssh_tunnel
+from dstack._internal.server.utils import sentry_utils
+from dstack._internal.utils.common import batched, get_current_datetime, get_or_error, run_async
+from dstack._internal.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+MAX_JOBS_FETCHED = 100
+BATCH_SIZE = 10
+MIN_COLLECT_INTERVAL_SECONDS = 9
+
+
+@sentry_utils.instrument_background_task
+async def collect_metrics():
+    async with get_session_ctx() as session:
+        res = await session.execute(
+            select(JobModel)
+            .where(JobModel.status.in_([JobStatus.RUNNING]))
+            .options(
+                joinedload(JobModel.instance)
+                .joinedload(InstanceModel.project)
+                .load_only(ProjectModel.ssh_private_key)
+            )
+            .order_by(JobModel.last_processed_at.asc())
+            .limit(MAX_JOBS_FETCHED)
+        )
+        job_models = res.unique().scalars().all()
+
+    for batch in batched(job_models, BATCH_SIZE):
+        await _collect_jobs_metrics(batch)
+
+
+@sentry_utils.instrument_background_task
+async def delete_metrics():
+    now_timestamp_micro = int(get_current_datetime().timestamp() * 1_000_000)
+    running_timestamp_micro_cutoff = (
+        now_timestamp_micro - settings.SERVER_METRICS_RUNNING_TTL_SECONDS * 1_000_000
+    )
+    finished_timestamp_micro_cutoff = (
+        now_timestamp_micro - settings.SERVER_METRICS_FINISHED_TTL_SECONDS * 1_000_000
+    )
+    await asyncio.gather(
+        _execute_delete_statement(
+            delete(JobMetricsPoint).where(
+                JobMetricsPoint.job_id.in_(
+                    select(JobModel.id).where(JobModel.status.in_([JobStatus.RUNNING]))
+                ),
+                JobMetricsPoint.timestamp_micro < running_timestamp_micro_cutoff,
+            )
+        ),
+        _execute_delete_statement(
+            delete(JobMetricsPoint).where(
+                JobMetricsPoint.job_id.in_(
+                    select(JobModel.id).where(JobModel.status.in_(JobStatus.finished_statuses()))
+                ),
+                JobMetricsPoint.timestamp_micro < finished_timestamp_micro_cutoff,
+            )
+        ),
+    )
+
+
+async def _execute_delete_statement(stmt: Delete) -> None:
+    async with get_session_ctx() as session:
+        await session.execute(stmt)
+        await session.commit()
+
+
+async def _collect_jobs_metrics(job_models: List[JobModel]):
+    filtered_job_models = await _filter_recently_collected_jobs(job_models)
+    tasks = []
+    for job_model in filtered_job_models:
+        tasks.append(_collect_job_metrics(job_model))
+    points = await asyncio.gather(*tasks)
+    async with get_session_ctx() as session:
+        for point in points:
+            if point is not None:
+                session.add(point)
+        await session.commit()
+
+
+async def _filter_recently_collected_jobs(job_models: List[JobModel]) -> List[JobModel]:
+    # Skip metrics collection if another replica collected it recently.
+    # Two replicas can still collect metrics simultaneously – that's fine since
+    # we'll just store some extra metric points in the db.
+    async with get_session_ctx() as session:
+        res = await session.execute(
+            select(JobMetricsPoint).where(
+                JobMetricsPoint.job_id.in_([j.id for j in job_models]),
+                JobMetricsPoint.timestamp_micro > _get_recently_collected_metric_cutoff(),
+            )
+        )
+        recent_points = res.scalars().all()
+        recent_job_ids = [p.job_id for p in recent_points]
+    return [j for j in job_models if j.id not in recent_job_ids]
+
+
+def _get_recently_collected_metric_cutoff() -> int:
+    now = int(get_current_datetime().timestamp() * 1_000_000)
+    cutoff = now - (MIN_COLLECT_INTERVAL_SECONDS * 1_000_000)
+    return cutoff
+
+
+async def _collect_job_metrics(job_model: JobModel) -> Optional[JobMetricsPoint]:
+    ssh_private_keys = get_instance_ssh_private_keys(get_or_error(job_model.instance))
+    jpd = get_job_provisioning_data(job_model)
+    jrd = get_job_runtime_data(job_model)
+    if jpd is None:
+        return None
+    try:
+        res = await run_async(
+            _pull_runner_metrics,
+            ssh_private_keys,
+            jpd,
+            jrd,
+        )
+    except Exception:
+        logger.exception("Failed to collect job %s metrics", job_model.job_name)
+        return None
+
+    if isinstance(res, bool):
+        # The job may already be terminated when collecting metrics - that's ok.
+        logger.warning("Failed to connect to job %s to collect metrics", job_model.job_name)
+        return None
+
+    if res is None:
+        logger.debug(
+            (
+                "Failed to collect job %s metrics."
+                " Either runner version does not support metrics API"
+                " or metrics collector is not available."
+            ),
+            job_model.job_name,
+        )
+        return None
+
+    gpus_memory_usage_bytes = [g.gpu_memory_usage_bytes for g in res.gpus]
+    gpus_util_percent = [g.gpu_util_percent for g in res.gpus]
+
+    return JobMetricsPoint(
+        job_id=job_model.id,
+        timestamp_micro=res.timestamp_micro,
+        cpu_usage_micro=res.cpu_usage_micro,
+        memory_usage_bytes=res.memory_usage_bytes,
+        memory_working_set_bytes=res.memory_working_set_bytes,
+        gpus_memory_usage_bytes=json.dumps(gpus_memory_usage_bytes),
+        gpus_util_percent=json.dumps(gpus_util_percent),
+    )
+
+
+@runner_ssh_tunnel(ports=[DSTACK_RUNNER_HTTP_PORT], retries=1)
+def _pull_runner_metrics(
+    ports: Dict[int, int],
+) -> Optional[MetricsResponse]:
+    runner_client = client.RunnerClient(port=ports[DSTACK_RUNNER_HTTP_PORT])
+    return runner_client.get_metrics()

dstack/_internal/server/background/tasks/process_placement_groups.py

@@ -0,0 +1,104 @@
+from typing import List
+from uuid import UUID
+
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import joinedload
+
+from dstack._internal.core.backends.base.compute import ComputeWithPlacementGroupSupport
+from dstack._internal.core.errors import PlacementGroupInUseError
+from dstack._internal.server.db import get_db, get_session_ctx
+from dstack._internal.server.models import PlacementGroupModel, ProjectModel
+from dstack._internal.server.services import backends as backends_services
+from dstack._internal.server.services.locking import get_locker
+from dstack._internal.server.services.placement import placement_group_model_to_placement_group
+from dstack._internal.server.utils import sentry_utils
+from dstack._internal.utils.common import get_current_datetime, run_async
+from dstack._internal.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+@sentry_utils.instrument_background_task
+async def process_placement_groups():
+    lock, lockset = get_locker(get_db().dialect_name).get_lockset(
+        PlacementGroupModel.__tablename__
+    )
+    async with get_session_ctx() as session:
+        async with lock:
+            res = await session.execute(
+                select(PlacementGroupModel)
+                .where(
+                    PlacementGroupModel.fleet_deleted == True,
+                    PlacementGroupModel.deleted == False,
+                    PlacementGroupModel.id.not_in(lockset),
+                )
+                .order_by(PlacementGroupModel.id)  # take locks in order
+                .with_for_update(skip_locked=True, key_share=True)
+            )
+            placement_group_models = res.scalars().all()
+            if len(placement_group_models) == 0:
+                return
+            placement_group_models_ids = [pg.id for pg in placement_group_models]
+            lockset.update(placement_group_models_ids)
+        try:
+            await _delete_placement_groups(
+                session=session,
+                placement_group_models_ids=placement_group_models_ids,
+            )
+        finally:
+            lockset.difference_update(placement_group_models_ids)
+
+
+async def _delete_placement_groups(
+    session: AsyncSession,
+    placement_group_models_ids: List[UUID],
+):
+    res = await session.execute(
+        select(PlacementGroupModel)
+        .where(
+            PlacementGroupModel.id.in_(placement_group_models_ids),
+            PlacementGroupModel.deleted == False,
+        )
+        .options(joinedload(PlacementGroupModel.project).joinedload(ProjectModel.backends))
+        .execution_options(populate_existing=True)
+    )
+    placement_group_models = res.unique().scalars().all()
+    for pg in placement_group_models:
+        await _delete_placement_group(pg)
+    await session.commit()
+
+
+async def _delete_placement_group(placement_group_model: PlacementGroupModel):
+    logger.debug("Deleting placement group %s", placement_group_model.name)
+    placement_group = placement_group_model_to_placement_group(placement_group_model)
+    if placement_group.provisioning_data is None:
+        logger.error(
+            "Failed to delete placement group %s. provisioning_data is None.", placement_group.name
+        )
+        return
+    backend = await backends_services.get_project_backend_by_type(
+        project=placement_group_model.project,
+        backend_type=placement_group.provisioning_data.backend,
+    )
+    if backend is None:
+        logger.error(
+            "Failed to delete placement group %s. Backend not available.", placement_group.name
+        )
+        return
+    compute = backend.compute()
+    assert isinstance(compute, ComputeWithPlacementGroupSupport)
+    try:
+        await run_async(compute.delete_placement_group, placement_group)
+    except PlacementGroupInUseError:
+        logger.info(
+            "Placement group %s is still in use. Skipping deletion for now.", placement_group.name
+        )
+        return
+    except Exception:
+        logger.exception("Got exception when deleting placement group %s", placement_group.name)
+        return
+
+    placement_group_model.deleted = True
+    placement_group_model.deleted_at = get_current_datetime()
+    logger.info("Deleted placement group %s", placement_group_model.name)

dstack/_internal/server/background/tasks/process_probes.py

@@ -0,0 +1,164 @@
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from datetime import timedelta
+from functools import partial
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import httpx
+from apscheduler.schedulers.asyncio import AsyncIOScheduler
+from httpx import AsyncClient, AsyncHTTPTransport
+from sqlalchemy import select, update
+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.services.ssh.tunnel import (
+    SSH_DEFAULT_OPTIONS,
+    IPSocket,
+    SocketPair,
+    UnixSocket,
+)
+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.locking import get_locker
+from dstack._internal.server.services.logging import fmt
+from dstack._internal.server.services.ssh import container_ssh_tunnel
+from dstack._internal.utils.common import get_current_datetime, get_or_error
+from dstack._internal.utils.logging import get_logger
+
+logger = get_logger(__name__)
+BATCH_SIZE = 100
+SSH_CONNECT_TIMEOUT = timedelta(seconds=10)
+PROCESSING_OVERHEAD_TIMEOUT = timedelta(minutes=1)
+PROBES_SCHEDULER = AsyncIOScheduler()
+
+
+async def process_probes():
+    probe_lock, probe_lockset = get_locker(get_db().dialect_name).get_lockset(
+        ProbeModel.__tablename__
+    )
+    async with get_session_ctx() as session:
+        async with probe_lock:
+            res = await session.execute(
+                select(ProbeModel.id)
+                .where(ProbeModel.id.not_in(probe_lockset))
+                .where(ProbeModel.active == True)
+                .where(ProbeModel.due <= get_current_datetime())
+                .order_by(ProbeModel.due.asc())
+                .limit(BATCH_SIZE)
+                .with_for_update(skip_locked=True, key_share=True)
+            )
+            probe_ids = res.unique().scalars().all()
+            probe_lockset.update(probe_ids)
+
+        try:
+            # Refetch to load all attributes.
+            # joinedload produces LEFT OUTER JOIN that can't be used with FOR UPDATE.
+            res = await session.execute(
+                select(ProbeModel)
+                .where(ProbeModel.id.in_(probe_ids))
+                .options(
+                    joinedload(ProbeModel.job)
+                    .joinedload(JobModel.instance)
+                    .joinedload(InstanceModel.project)
+                )
+                .options(joinedload(ProbeModel.job))
+                .execution_options(populate_existing=True)
+            )
+            probes = res.unique().scalars().all()
+            for probe in probes:
+                if probe.job.status != JobStatus.RUNNING:
+                    probe.active = False
+                else:
+                    job_spec: JobSpec = JobSpec.__response__.parse_raw(probe.job.job_spec_data)
+                    probe_spec = job_spec.probes[probe.probe_num]
+                    # Schedule the next probe execution in case this execution is interrupted
+                    probe.due = get_current_datetime() + _get_probe_async_processing_timeout(
+                        probe_spec
+                    )
+                    # Execute the probe asynchronously outside of the DB session
+                    PROBES_SCHEDULER.add_job(partial(_process_probe_async, probe, probe_spec))
+            await session.commit()
+        finally:
+            probe_lockset.difference_update(probe_ids)
+
+
+async def _process_probe_async(probe: ProbeModel, probe_spec: ProbeSpec) -> None:
+    start = get_current_datetime()
+    logger.debug("%s: processing probe", fmt(probe))
+    success = await _execute_probe(probe, probe_spec)
+
+    async with get_session_ctx() as session:
+        async with get_locker(get_db().dialect_name).lock_ctx(
+            ProbeModel.__tablename__, [probe.id]
+        ):
+            await session.execute(
+                update(ProbeModel)
+                .where(ProbeModel.id == probe.id)
+                .values(
+                    success_streak=0 if not success else ProbeModel.success_streak + 1,
+                    due=get_current_datetime() + timedelta(seconds=probe_spec.interval),
+                )
+            )
+    logger.debug(
+        "%s: probe processing took %ss",
+        fmt(probe),
+        (get_current_datetime() - start).total_seconds(),
+    )
+
+
+async def _execute_probe(probe: ProbeModel, probe_spec: ProbeSpec) -> bool:
+    """
+    Returns:
+        Whether probe execution was successful.
+    """
+
+    try:
+        async with _get_service_replica_client(probe.job) as client:
+            resp = await client.request(
+                method=probe_spec.method,
+                url="http://dstack" + probe_spec.url,
+                headers=[(h.name, h.value) for h in probe_spec.headers],
+                content=probe_spec.body,
+                timeout=probe_spec.timeout,
+                follow_redirects=False,
+            )
+            logger.debug("%s: probe status code: %s", fmt(probe), resp.status_code)
+            return resp.is_success
+    except (SSHError, httpx.RequestError) as e:
+        logger.debug("%s: probe failed: %r", fmt(probe), e)
+        return False
+
+
+def _get_probe_async_processing_timeout(probe_spec: ProbeSpec) -> timedelta:
+    return (
+        timedelta(seconds=probe_spec.timeout)
+        + SSH_CONNECT_TIMEOUT
+        + PROCESSING_OVERHEAD_TIMEOUT  # slow db queries and other unforeseen conditions
+    )
+
+
+@asynccontextmanager
+async def _get_service_replica_client(job: JobModel) -> AsyncGenerator[AsyncClient, None]:
+    options = {
+        **SSH_DEFAULT_OPTIONS,
+        "ConnectTimeout": str(int(SSH_CONNECT_TIMEOUT.total_seconds())),
+    }
+    job_spec: JobSpec = JobSpec.__response__.parse_raw(job.job_spec_data)
+    with TemporaryDirectory() as temp_dir:
+        app_socket_path = (Path(temp_dir) / "replica.sock").absolute()
+        async with container_ssh_tunnel(
+            job=job,
+            forwarded_sockets=[
+                SocketPair(
+                    remote=IPSocket("localhost", get_or_error(job_spec.service_port)),
+                    local=UnixSocket(app_socket_path),
+                ),
+            ],
+            options=options,
+        ):
+            async with AsyncClient(
+                transport=AsyncHTTPTransport(uds=str(app_socket_path))
+            ) as client:
+                yield client

dstack/_internal/server/background/tasks/process_prometheus_metrics.py

@@ -0,0 +1,150 @@
+import uuid
+from datetime import datetime, timedelta
+from typing import Optional
+
+import sqlalchemy.exc
+from sqlalchemy import delete, or_, select, update
+from sqlalchemy.orm import joinedload
+
+from dstack._internal.core.consts import DSTACK_SHIM_HTTP_PORT
+from dstack._internal.core.models.runs import JobStatus
+from dstack._internal.server.db import get_session_ctx
+from dstack._internal.server.models import (
+    InstanceModel,
+    JobModel,
+    JobPrometheusMetrics,
+    ProjectModel,
+)
+from dstack._internal.server.services.instances import get_instance_ssh_private_keys
+from dstack._internal.server.services.jobs import get_job_provisioning_data, get_job_runtime_data
+from dstack._internal.server.services.runner import client
+from dstack._internal.server.services.runner.ssh import runner_ssh_tunnel
+from dstack._internal.server.utils import sentry_utils
+from dstack._internal.server.utils.common import gather_map_async
+from dstack._internal.utils.common import batched, get_current_datetime, get_or_error, run_async
+from dstack._internal.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+MAX_JOBS_FETCHED = 100
+BATCH_SIZE = 10
+MIN_COLLECT_INTERVAL_SECONDS = 9
+# 10 minutes should be more than enough to scrape metrics, and, in any case,
+# 10 minutes old metrics has little to no value
+METRICS_TTL_SECONDS = 600
+
+
+@sentry_utils.instrument_background_task
+async def collect_prometheus_metrics():
+    now = get_current_datetime()
+    cutoff = now - timedelta(seconds=MIN_COLLECT_INTERVAL_SECONDS)
+    async with get_session_ctx() as session:
+        res = await session.execute(
+            select(JobModel)
+            .join(JobPrometheusMetrics, isouter=True)
+            .where(
+                JobModel.status.in_([JobStatus.RUNNING]),
+                or_(
+                    JobPrometheusMetrics.job_id.is_(None),
+                    JobPrometheusMetrics.collected_at < cutoff,
+                ),
+            )
+            .options(
+                joinedload(JobModel.instance)
+                .joinedload(InstanceModel.project)
+                .load_only(ProjectModel.ssh_private_key)
+            )
+            .order_by(JobModel.last_processed_at.asc())
+            .limit(MAX_JOBS_FETCHED)
+        )
+        job_models = res.unique().scalars().all()
+    for batch in batched(job_models, BATCH_SIZE):
+        await _collect_jobs_metrics(batch, now)
+
+
+@sentry_utils.instrument_background_task
+async def delete_prometheus_metrics():
+    now = get_current_datetime()
+    cutoff = now - timedelta(seconds=METRICS_TTL_SECONDS)
+    async with get_session_ctx() as session:
+        await session.execute(
+            delete(JobPrometheusMetrics).where(JobPrometheusMetrics.collected_at < cutoff)
+        )
+        await session.commit()
+
+
+async def _collect_jobs_metrics(job_models: list[JobModel], collected_at: datetime):
+    results = await gather_map_async(job_models, _collect_job_metrics, return_exceptions=True)
+    async with get_session_ctx() as session:
+        for job_model, result in results:
+            if result is None:
+                continue
+            if isinstance(result, BaseException):
+                logger.error(
+                    "Failed to collect job %s Prometheus metrics: %r", job_model.job_name, result
+                )
+                continue
+            res = await session.execute(
+                update(JobPrometheusMetrics)
+                .where(JobPrometheusMetrics.job_id == job_model.id)
+                .values(
+                    collected_at=collected_at,
+                    text=result,
+                )
+                .returning(JobPrometheusMetrics)
+            )
+            metrics = res.scalar()
+            if metrics is None:
+                metrics = JobPrometheusMetrics(
+                    job_id=job_model.id,
+                    collected_at=collected_at,
+                    text=result,
+                )
+                try:
+                    async with session.begin_nested():
+                        session.add(metrics)
+                except sqlalchemy.exc.IntegrityError:
+                    # Concurrent server replica already committed, ignoring
+                    pass
+        await session.commit()
+
+
+async def _collect_job_metrics(job_model: JobModel) -> Optional[str]:
+    jpd = get_job_provisioning_data(job_model)
+    if jpd is None:
+        return None
+    if not jpd.dockerized:
+        # Container-based backend, no shim
+        return None
+    ssh_private_keys = get_instance_ssh_private_keys(get_or_error(job_model.instance))
+    jrd = get_job_runtime_data(job_model)
+    try:
+        res = await run_async(
+            _pull_job_metrics,
+            ssh_private_keys,
+            jpd,
+            jrd,
+            job_model.id,
+        )
+    except Exception:
+        logger.exception("Failed to collect job %s Prometheus metrics", job_model.job_name)
+        return None
+
+    if isinstance(res, bool):
+        logger.warning(
+            "Failed to connect to job %s to collect Prometheus metrics", job_model.job_name
+        )
+        return None
+
+    if res is None:
+        # Either not supported by shim or exporter is not available
+        return None
+
+    return res
+
+
+@runner_ssh_tunnel(ports=[DSTACK_SHIM_HTTP_PORT], retries=1)
+def _pull_job_metrics(ports: dict[int, int], task_id: uuid.UUID) -> Optional[str]:
+    shim_client = client.ShimClient(port=ports[DSTACK_SHIM_HTTP_PORT])
+    return shim_client.get_task_metrics(task_id)