diff --git a/pytest.ini b/pytest.ini index ca71dc5cc..899f67a61 100644 --- a/pytest.ini +++ b/pytest.ini @@ -7,3 +7,4 @@ addopts = --allow-unix-socket markers = shim_version + dockerized diff --git a/src/dstack/_internal/server/background/tasks/process_prometheus_metrics.py b/src/dstack/_internal/server/background/tasks/process_prometheus_metrics.py index 78a953940..b9e15ca0b 100644 --- a/src/dstack/_internal/server/background/tasks/process_prometheus_metrics.py +++ b/src/dstack/_internal/server/background/tasks/process_prometheus_metrics.py @@ -99,11 +99,14 @@ async def _collect_jobs_metrics(job_models: list[JobModel], collected_at: dateti async def _collect_job_metrics(job_model: JobModel) -> Optional[str]: - 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 + 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, diff --git a/src/tests/_internal/server/background/tasks/test_process_prometheus_metrics.py b/src/tests/_internal/server/background/tasks/test_process_prometheus_metrics.py index f0228fe7c..232227d88 100644 --- a/src/tests/_internal/server/background/tasks/test_process_prometheus_metrics.py +++ b/src/tests/_internal/server/background/tasks/test_process_prometheus_metrics.py @@ -34,7 +34,13 @@ @pytest.mark.usefixtures("test_db", "image_config_mock") class TestCollectPrometheusMetrics: @pytest_asyncio.fixture - async def job(self, session: AsyncSession) -> JobModel: + async def job(self, request: pytest.FixtureRequest, session: AsyncSession) -> JobModel: + dockerized: bool + marker = request.node.get_closest_marker("dockerized") + if marker is None: + dockerized = True + else: + dockerized = marker.args[0] user = await create_user(session=session, global_role=GlobalRole.USER) project = await create_project(session=session, owner=user) await add_project_member( @@ -59,7 +65,7 @@ async def job(self, session: AsyncSession) -> JobModel: session=session, run=run, status=JobStatus.RUNNING, - job_provisioning_data=get_job_provisioning_data(), + job_provisioning_data=get_job_provisioning_data(dockerized=dockerized), instance_assigned=True, instance=instance, ) @@ -142,6 +148,21 @@ async def test_skips_recently_updated( assert metrics.text == "# prom old response" assert metrics.collected_at == datetime(2023, 1, 2, 3, 5, 15) + @freeze_time(datetime(2023, 1, 2, 3, 5, 20, tzinfo=timezone.utc)) + @pytest.mark.dockerized(False) + async def test_skips_non_dockerized_jobs( + self, session: AsyncSession, job: JobModel, ssh_tunnel_mock: Mock, shim_client_mock: Mock + ): + await collect_prometheus_metrics() + + ssh_tunnel_mock.assert_not_called() + shim_client_mock.get_task_metrics.assert_not_called() + res = await session.execute( + select(JobPrometheusMetrics).where(JobPrometheusMetrics.job_id == job.id) + ) + metrics = res.scalar_one_or_none() + assert metrics is None + @pytest.mark.asyncio @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)