diff --git a/e2e/auditor/test_audit_job.py b/e2e/auditor/test_audit_job.py index b375842c71..35f7fad537 100644 --- a/e2e/auditor/test_audit_job.py +++ b/e2e/auditor/test_audit_job.py @@ -21,6 +21,7 @@ import pytest from nemo_platform import NeMoPlatform from nmp.testing import add_mock_provider, short_unique_name +from nmp.testing.e2e import cleanup_platform_job from e2e.auditor.utils import minimal_audit_config, unique_name @@ -56,10 +57,13 @@ def _wait_for_audit_job(sdk: NeMoPlatform, job_name: str, workspace: str) -> str def _cleanup_audit_job(sdk: NeMoPlatform, job_name: str, workspace: str) -> None: - with suppress(Exception): - sdk.jobs.cancel(name=job_name, workspace=workspace) - with suppress(Exception): - sdk.jobs.delete(name=job_name, workspace=workspace) + cleanup_platform_job( + sdk, + job_name, + workspace, + timeout=AUDIT_JOB_TIMEOUT_SECONDS, + poll_interval=AUDIT_JOB_POLL_INTERVAL_SECONDS, + ) def _add_mock_provider_or_skip(sdk: NeMoPlatform, workspace: str, name: str) -> str: diff --git a/e2e/test_anonymizer_plugin.py b/e2e/test_anonymizer_plugin.py index 2a58030710..52ded6c479 100644 --- a/e2e/test_anonymizer_plugin.py +++ b/e2e/test_anonymizer_plugin.py @@ -39,6 +39,7 @@ from nemo_platform_plugin.files.client import FilesClient from nemo_platform_plugin.files.types import CreateFilesetRequest from nmp.testing import MockProviderResponse, add_mock_provider, short_unique_name +from nmp.testing.e2e import cleanup_platform_job pytestmark = [ pytest.mark.container_only, @@ -323,10 +324,13 @@ def _wait_for_anonymizer_job(job: AnonymizerJobResource, *, timeout_seconds: flo def _cleanup_anonymizer_job(sdk: NeMoPlatform, job_name: str) -> None: - with suppress(Exception): - sdk.jobs.cancel(name=job_name, workspace=sdk.workspace) - with suppress(Exception): - sdk.jobs.delete(name=job_name, workspace=sdk.workspace) + cleanup_platform_job( + sdk, + job_name, + str(sdk.workspace), + timeout=ANONYMIZER_JOB_TIMEOUT_SECONDS, + poll_interval=ANONYMIZER_POLL_INTERVAL_SECONDS, + ) @pytest.fixture(scope="module") diff --git a/e2e/test_evaluator_plugin.py b/e2e/test_evaluator_plugin.py index 3b2a9f2330..b214b21d86 100644 --- a/e2e/test_evaluator_plugin.py +++ b/e2e/test_evaluator_plugin.py @@ -49,7 +49,7 @@ from nemo_platform import APIConnectionError, APIStatusError, NeMoPlatform from nemo_platform.types.inference import ModelProvider from nmp.testing import add_mock_provider, short_unique_name, wait_for_model_entity -from nmp.testing.e2e import wait_for_platform_job +from nmp.testing.e2e import cleanup_platform_job, wait_for_platform_job from nmp.testing.utils import ensure_passthrough_virtual_model pytestmark = [ @@ -286,10 +286,13 @@ def _create_ready_mock_model( def _cleanup_evaluator_job(sdk: NeMoPlatform, job_name: str) -> None: - with suppress(Exception): - sdk.jobs.cancel(name=job_name, workspace=sdk.workspace) - with suppress(Exception): - sdk.jobs.delete(name=job_name, workspace=sdk.workspace) + cleanup_platform_job( + sdk, + job_name, + str(sdk.workspace), + timeout=EVALUATOR_JOB_TIMEOUT_SECONDS, + poll_interval=EVALUATOR_POLL_INTERVAL_SECONDS, + ) def _wait_for_evaluator_job(job: EvaluatorJobResource) -> None: diff --git a/e2e/test_safe_synthesizer.py b/e2e/test_safe_synthesizer.py index bf2049ebef..b731efb17c 100644 --- a/e2e/test_safe_synthesizer.py +++ b/e2e/test_safe_synthesizer.py @@ -345,14 +345,21 @@ def _wait_for_job_absent( def _delete_nss_job(sdk: NeMoPlatform, workspace: str, name: str, *, verify: bool = True) -> None: - response = sdk._client.delete( - _nss_url(sdk, workspace, f"jobs/{name}"), - headers=_string_headers(sdk), - timeout=60.0, - ) - if response.status_code not in {200, 202, 204, 404}: - response.raise_for_status() - if verify: + deadline = time.monotonic() + DELETE_VERIFY_TIMEOUT_SECONDS + while True: + response = sdk._client.delete( + _nss_url(sdk, workspace, f"jobs/{name}"), + headers=_string_headers(sdk), + timeout=60.0, + ) + if response.status_code in {200, 202, 204, 404}: + break + if response.status_code != 409: + response.raise_for_status() + if time.monotonic() >= deadline: + raise TimeoutError(f"Timed out deleting Safe Synthesizer job {name!r}; last response={response.text}") + time.sleep(2.0) + if verify and response.status_code != 404: _wait_for_job_absent(sdk, workspace, name) @@ -583,9 +590,42 @@ def create(prefix: str, data_source: str, config: dict[str, Any]) -> dict[str, A try: yield create finally: + cleanup_errors: list[Exception] = [] for job_name in reversed(job_names): - _cancel_nss_job(sdk, workspace, job_name) - _delete_nss_job(sdk, workspace, job_name) + try: + cancel_response = _cancel_nss_job(sdk, workspace, job_name) + if cancel_response is not None: + _wait_for_status( + sdk, + workspace, + job_name, + timeout_seconds=SMOKE_JOB_TIMEOUT_SECONDS, + poll_interval_seconds=2.0, + ) + except Exception as exc: + cleanup_errors.append(exc) + _capture_nss_debug_artifacts( + sdk, + workspace, + job_name, + "cleanup cancel/wait failed", + history=[], + error=exc, + ) + try: + _delete_nss_job(sdk, workspace, job_name) + except Exception as exc: + cleanup_errors.append(exc) + _capture_nss_debug_artifacts( + sdk, + workspace, + job_name, + "cleanup delete failed", + history=[], + error=exc, + ) + if cleanup_errors: + raise ExceptionGroup("Safe Synthesizer job cleanup failed", cleanup_errors) def test_safe_synthesizer_api_health(sdk: NeMoPlatform, workspace: str) -> None: diff --git a/openapi/ga/individual/platform.openapi.yaml b/openapi/ga/individual/platform.openapi.yaml index 7b6b825abd..42e64aa09c 100644 --- a/openapi/ga/individual/platform.openapi.yaml +++ b/openapi/ga/individual/platform.openapi.yaml @@ -5719,6 +5719,9 @@ paths: description: Successful Response '404': description: Job not Found + '409': + description: Job, attempt, or step is not terminal; paused or cancelling + jobs must reach a terminal state before deletion. '422': description: Validation Error content: diff --git a/openapi/ga/openapi.yaml b/openapi/ga/openapi.yaml index 7b6b825abd..42e64aa09c 100644 --- a/openapi/ga/openapi.yaml +++ b/openapi/ga/openapi.yaml @@ -5719,6 +5719,9 @@ paths: description: Successful Response '404': description: Job not Found + '409': + description: Job, attempt, or step is not terminal; paused or cancelling + jobs must reach a terminal state before deletion. '422': description: Validation Error content: diff --git a/openapi/openapi.yaml b/openapi/openapi.yaml index 7b6b825abd..42e64aa09c 100644 --- a/openapi/openapi.yaml +++ b/openapi/openapi.yaml @@ -5719,6 +5719,9 @@ paths: description: Successful Response '404': description: Job not Found + '409': + description: Job, attempt, or step is not terminal; paused or cancelling + jobs must reach a terminal state before deletion. '422': description: Validation Error content: diff --git a/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/api_factory.py b/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/api_factory.py index 65306b22af..af9eee50e9 100644 --- a/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/api_factory.py +++ b/packages/nemo_platform_plugin/src/nemo_platform_plugin/jobs/api_factory.py @@ -51,6 +51,8 @@ from nemo_platform_plugin.api.parsed_filter import ParsedFilter, make_filter_dep from nemo_platform_plugin.authz import AuthzScope, CallerKind, path_rule from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import NemoHTTPError +from nemo_platform_plugin.client.types import RetryPolicy from nemo_platform_plugin.dependencies import get_entity_client, get_sdk_client from nemo_platform_plugin.entities import EntityClient from nemo_platform_plugin.jobs.client import AsyncJobsClient @@ -1083,6 +1085,10 @@ async def get_job_status( @router.delete( "/jobs/{name}", status_code=status.HTTP_204_NO_CONTENT, + responses={ + status.HTTP_404_NOT_FOUND: {"description": "Job not Found"}, + status.HTTP_409_CONFLICT: {"description": "Job is not in a terminal state"}, + }, ) async def delete_job( workspace: str, @@ -1090,7 +1096,11 @@ async def delete_job( sdk: AsyncNeMoPlatform = Depends(get_sdk_client), ) -> None: f"""Delete a job by name for the {service_name} microservice.""" - await client_from_platform(sdk, AsyncJobsClient).delete_job(name=name, workspace=workspace) + try: + jobs_client = client_from_platform(sdk, AsyncJobsClient).with_retry(RetryPolicy(max_retries=0)) + await jobs_client.delete_job(name=name, workspace=workspace) + except NemoHTTPError as exc: + raise HTTPException(status_code=exc.status_code, detail=exc.detail) from exc return None @router.post( diff --git a/packages/nmp_common/src/nmp/common/docker/gpu_pool.py b/packages/nmp_common/src/nmp/common/docker/gpu_pool.py index d912de3fe8..7c98c44576 100644 --- a/packages/nmp_common/src/nmp/common/docker/gpu_pool.py +++ b/packages/nmp_common/src/nmp/common/docker/gpu_pool.py @@ -17,7 +17,30 @@ class GPUAllocationError(Exception): """Raised when GPU allocation fails due to insufficient resources.""" - pass + def __init__( + self, + message: str, + *, + requested: int | None = None, + available: int | None = None, + total: int | None = None, + ) -> None: + super().__init__(message) + self.message = message + self.requested = requested + self.available = available + self.total = total + + @property + def is_transient_capacity_exhaustion(self) -> bool: + """Whether a retry could succeed after another workload releases GPUs.""" + return ( + self.requested is not None + and self.available is not None + and self.total is not None + and self.requested > 0 + and self.available < self.requested <= self.total + ) @dataclass @@ -92,9 +115,13 @@ def allocate_gpu(self, workload_id: str, num_requested: int = 1) -> list[int]: available_gpus = {gpu for gpu, workload in self.gpu_to_workload_id.items() if workload is None} if len(available_gpus) < num_requested: + available_count = len(available_gpus) raise GPUAllocationError( f"Not enough GPUs available. Requested {num_requested}, " - f"available {len(available_gpus)} out of {self.num_reserved_gpus} total." + f"available {available_count} out of {self.num_reserved_gpus} total.", + requested=num_requested, + available=available_count, + total=self.num_reserved_gpus, ) gpu_ids = [] for _ in range(num_requested): diff --git a/packages/nmp_common/tests/api_factory/test_api_factory.py b/packages/nmp_common/tests/api_factory/test_api_factory.py index 2ca14a6f92..ad029eb406 100644 --- a/packages/nmp_common/tests/api_factory/test_api_factory.py +++ b/packages/nmp_common/tests/api_factory/test_api_factory.py @@ -355,6 +355,7 @@ def _job_routes_app(): mock_sdk = MagicMock() mock_jobs = MagicMock() + mock_jobs.with_retry.return_value = mock_jobs router = job_route_factory( service_name="test_service", job_type="TestJob", diff --git a/packages/nmp_common/tests/docker/test_gpu_pool.py b/packages/nmp_common/tests/docker/test_gpu_pool.py index 73b3e3f20a..53386af0d8 100644 --- a/packages/nmp_common/tests/docker/test_gpu_pool.py +++ b/packages/nmp_common/tests/docker/test_gpu_pool.py @@ -71,14 +71,17 @@ def test_allocate_sequential_workloads(self): assert set(gpu_ids_1).isdisjoint(set(gpu_ids_2)) @pytest.mark.parametrize( - "pool_size,pre_allocate,num_requested,expected_error_fragment", + "pool_size,pre_allocate,num_requested,expected_error_fragment,expected_transient", [ - pytest.param(2, 1, 2, "Requested 2", id="insufficient_remaining"), - pytest.param(1, 1, 1, "Requested 1", id="none_available"), - pytest.param(0, 0, 1, "Requested 1", id="empty_pool"), + pytest.param(2, 1, 2, "Requested 2", True, id="insufficient_remaining"), + pytest.param(1, 1, 1, "Requested 1", True, id="none_available"), + pytest.param(1, 0, 2, "Requested 2", False, id="request_exceeds_total"), + pytest.param(0, 0, 1, "Requested 1", False, id="empty_pool"), ], ) - def test_allocate_raises_when_insufficient(self, pool_size, pre_allocate, num_requested, expected_error_fragment): + def test_allocate_raises_when_insufficient( + self, pool_size, pre_allocate, num_requested, expected_error_fragment, expected_transient + ): """Test that allocation raises GPUAllocationError when not enough GPUs available.""" pool = DockerGPUPool(reserved_gpu_device_ids=list(range(pool_size))) if pre_allocate > 0: @@ -88,6 +91,10 @@ def test_allocate_raises_when_insufficient(self, pool_size, pre_allocate, num_re pool.allocate_gpu("workload-new", num_requested=num_requested) assert expected_error_fragment in str(exc_info.value) + assert exc_info.value.requested == num_requested + assert exc_info.value.available == pool_size - pre_allocate + assert exc_info.value.total == pool_size + assert exc_info.value.is_transient_capacity_exhaustion is expected_transient @pytest.mark.parametrize( "invalid_value", @@ -107,6 +114,7 @@ def test_allocate_raises_on_invalid_num_requested(self, invalid_value): assert "Invalid GPU request" in str(exc_info.value) assert "Must be a positive integer" in str(exc_info.value) + assert exc_info.value.is_transient_capacity_exhaustion is False class TestDockerGPUPoolRelease: diff --git a/packages/nmp_testing/src/nmp/testing/e2e/__init__.py b/packages/nmp_testing/src/nmp/testing/e2e/__init__.py index 3681335237..bd20393f55 100644 --- a/packages/nmp_testing/src/nmp/testing/e2e/__init__.py +++ b/packages/nmp_testing/src/nmp/testing/e2e/__init__.py @@ -43,7 +43,7 @@ load_config, ) from .docker import Docker -from .jobs import wait_for_job_completion, wait_for_job_logs, wait_for_platform_job +from .jobs import cleanup_platform_job, wait_for_job_completion, wait_for_job_logs, wait_for_platform_job from .kubernetes import Kubernetes __all__ = [ @@ -59,6 +59,7 @@ "discover_configs", "infer_backend", "load_config", + "cleanup_platform_job", "wait_for_job_completion", "wait_for_job_logs", "wait_for_platform_job", diff --git a/packages/nmp_testing/src/nmp/testing/e2e/jobs.py b/packages/nmp_testing/src/nmp/testing/e2e/jobs.py index 5dd12838e4..aa7e233eb7 100644 --- a/packages/nmp_testing/src/nmp/testing/e2e/jobs.py +++ b/packages/nmp_testing/src/nmp/testing/e2e/jobs.py @@ -13,6 +13,7 @@ from nemo_platform import NeMoPlatform from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import ConflictError, NotFoundError from nemo_platform_plugin.jobs.client import JobsClient from nemo_platform_plugin.jobs.schemas import PlatformJobLogPage from nemo_platform_plugin.jobs.types import PlatformJobResponse @@ -85,6 +86,11 @@ def poll_until_terminal( TERMINAL_STATUSES = {"completed", "error", "cancelled"} +def _job_status_value(status: object) -> str: + value = getattr(status, "value", status) + return str(value or "").lower() + + def wait_for_platform_job( sdk: NeMoPlatform, job_name: str, @@ -129,7 +135,7 @@ def wait_for_platform_job( def get_status() -> str: nonlocal last_job last_job = client_from_platform(sdk, JobsClient).get_job(name=job_name, workspace=workspace).data() - current = last_job.status.lower() if last_job.status else "" + current = _job_status_value(last_job.status) if not status_history or status_history[-1] != current: status_history.append(current) return current @@ -157,6 +163,84 @@ def get_status() -> str: return last_job +def wait_for_platform_job_terminal_or_absent( + sdk: NeMoPlatform, + job_name: str, + workspace: str, + timeout: float = 120.0, + poll_interval: float = 1.0, + terminal_statuses: set[str] | None = None, +) -> PlatformJobResponse | None: + """Wait until a platform job is terminal or already absent.""" + jobs = client_from_platform(sdk, JobsClient) + terminal = terminal_statuses or TERMINAL_STATUSES + deadline = time.monotonic() + timeout + last_status = "" + + while True: + try: + job = jobs.get_job(name=job_name, workspace=workspace).data() + except NotFoundError: + return None + + last_status = _job_status_value(job.status) + if last_status in terminal: + return job + + if time.monotonic() >= deadline: + raise TimeoutError(f"'{job_name}' did not reach a terminal status within {timeout}s. Status: {last_status}") + time.sleep(poll_interval) + + +def cleanup_platform_job( + sdk: NeMoPlatform, + job_name: str, + workspace: str, + timeout: float = 120.0, + poll_interval: float = 1.0, +) -> None: + """Cancel, wait for terminal state, and delete a platform job. + + A missing job is treated as already cleaned up. A delete conflict is not + swallowed; the helper keeps polling DELETE until it succeeds, the job is + absent, or the bounded timeout expires. + """ + jobs = client_from_platform(sdk, JobsClient) + try: + job = jobs.get_job(name=job_name, workspace=workspace).data() + except NotFoundError: + return + + if _job_status_value(job.status) not in TERMINAL_STATUSES: + try: + jobs.cancel_job(name=job_name, workspace=workspace).data() + except NotFoundError: + return + except ConflictError: + pass + wait_for_platform_job_terminal_or_absent( + sdk, + job_name, + workspace, + timeout=timeout, + poll_interval=poll_interval, + ) + + deadline = time.monotonic() + timeout + while True: + try: + jobs.delete_job(name=job_name, workspace=workspace).data() + return + except NotFoundError: + return + except ConflictError as exc: + if time.monotonic() >= deadline: + raise TimeoutError( + f"'{job_name}' could not be deleted within {timeout}s because it remained non-terminal" + ) from exc + time.sleep(poll_interval) + + def wait_for_job_completion( sdk: NeMoPlatform, service: str, diff --git a/packages/nmp_testing/tests/unit/test_jobs.py b/packages/nmp_testing/tests/unit/test_jobs.py index 8b43435cdc..f627e35b78 100644 --- a/packages/nmp_testing/tests/unit/test_jobs.py +++ b/packages/nmp_testing/tests/unit/test_jobs.py @@ -18,8 +18,10 @@ from contextlib import contextmanager from unittest.mock import MagicMock, patch +import httpx import pytest -from nmp.testing.e2e.jobs import TERMINAL_STATUSES, wait_for_platform_job +from nemo_platform_plugin.client.errors import ConflictError, NotFoundError +from nmp.testing.e2e.jobs import TERMINAL_STATUSES, cleanup_platform_job, wait_for_platform_job # --------------------------------------------------------------------------- # Helpers @@ -38,6 +40,12 @@ def _resp(data): return m +def _http_error(error_type, status_code: int, method: str = "GET"): + request = httpx.Request(method, "http://localhost/apis/jobs/v2/workspaces/ws/jobs/my-job") + response = httpx.Response(status_code, request=request, json={"detail": f"HTTP {status_code}"}) + return error_type(response) + + def _make_jobs_client(*statuses: str) -> MagicMock: """Return a typed jobs-client mock whose get_job() cycles through *statuses*. @@ -53,6 +61,8 @@ def _make_jobs_client(*statuses: str) -> MagicMock: responses.append(_resp(j)) jobs_client.get_job.side_effect = responses jobs_client.get_job_status.return_value = _resp(MagicMock(model_dump=MagicMock(return_value={}))) + jobs_client.cancel_job.return_value = _resp(None) + jobs_client.delete_job.return_value = _resp(None) return jobs_client @@ -251,3 +261,40 @@ def fake_poll(get_status, label, terminal, timeout, image_pull_timeout, poll_int wait_for_platform_job(_make_sdk(), "my-job", "ws", timeout=5.0) assert "Job status details:" in str(exc_info.value) + + +class TestCleanupPlatformJob: + def test_cancels_waits_and_deletes_non_terminal_job(self): + jobs_client = _make_jobs_client("active", "cancelled") + with _patch_client(jobs_client): + cleanup_platform_job(_make_sdk(), "my-job", "ws", timeout=5.0, poll_interval=0.0) + + jobs_client.cancel_job.assert_called_once_with(name="my-job", workspace="ws") + jobs_client.delete_job.assert_called_once_with(name="my-job", workspace="ws") + + def test_missing_job_is_already_cleaned_up(self): + jobs_client = _make_jobs_client() + jobs_client.get_job.side_effect = _http_error(NotFoundError, 404) + + with _patch_client(jobs_client): + cleanup_platform_job(_make_sdk(), "my-job", "ws", timeout=5.0, poll_interval=0.0) + + jobs_client.cancel_job.assert_not_called() + jobs_client.delete_job.assert_not_called() + + def test_retries_delete_conflict_until_success(self): + jobs_client = _make_jobs_client("completed") + jobs_client.delete_job.side_effect = [_http_error(ConflictError, 409, method="DELETE"), _resp(None)] + + with _patch_client(jobs_client): + cleanup_platform_job(_make_sdk(), "my-job", "ws", timeout=5.0, poll_interval=0.0) + + assert jobs_client.delete_job.call_count == 2 + + def test_delete_conflict_timeout_raises_timeout(self): + jobs_client = _make_jobs_client("completed") + jobs_client.delete_job.side_effect = _http_error(ConflictError, 409, method="DELETE") + + with _patch_client(jobs_client): + with pytest.raises(TimeoutError, match="could not be deleted"): + cleanup_platform_job(_make_sdk(), "my-job", "ws", timeout=0.0, poll_interval=0.0) diff --git a/plugins/nemo-agents/openapi/openapi.yaml b/plugins/nemo-agents/openapi/openapi.yaml index f6b8f63de6..b3d1d2ec9c 100644 --- a/plugins/nemo-agents/openapi/openapi.yaml +++ b/plugins/nemo-agents/openapi/openapi.yaml @@ -1043,6 +1043,10 @@ paths: responses: '204': description: Successful Response + '404': + description: Job not Found + '409': + description: Job is not in a terminal state '422': description: Validation Error content: @@ -1513,6 +1517,10 @@ paths: responses: '204': description: Successful Response + '404': + description: Job not Found + '409': + description: Job is not in a terminal state '422': description: Validation Error content: @@ -1791,6 +1799,10 @@ paths: responses: '204': description: Successful Response + '404': + description: Job not Found + '409': + description: Job is not in a terminal state '422': description: Validation Error content: @@ -2165,6 +2177,10 @@ paths: responses: '204': description: Successful Response + '404': + description: Job not Found + '409': + description: Job is not in a terminal state '422': description: Validation Error content: @@ -2635,6 +2651,10 @@ paths: responses: '204': description: Successful Response + '404': + description: Job not Found + '409': + description: Job is not in a terminal state '422': description: Validation Error content: @@ -2913,6 +2933,10 @@ paths: responses: '204': description: Successful Response + '404': + description: Job not Found + '409': + description: Job is not in a terminal state '422': description: Validation Error content: @@ -3287,6 +3311,10 @@ paths: responses: '204': description: Successful Response + '404': + description: Job not Found + '409': + description: Job is not in a terminal state '422': description: Validation Error content: diff --git a/plugins/nemo-anonymizer/openapi/openapi.yaml b/plugins/nemo-anonymizer/openapi/openapi.yaml index fdd8f47943..0d9ec1cae6 100644 --- a/plugins/nemo-anonymizer/openapi/openapi.yaml +++ b/plugins/nemo-anonymizer/openapi/openapi.yaml @@ -258,6 +258,10 @@ paths: responses: '204': description: Successful Response + '404': + description: Job not Found + '409': + description: Job is not in a terminal state '422': description: Validation Error content: diff --git a/plugins/nemo-auditor/openapi/openapi.yaml b/plugins/nemo-auditor/openapi/openapi.yaml index 50cddbcf89..cb9f271980 100644 --- a/plugins/nemo-auditor/openapi/openapi.yaml +++ b/plugins/nemo-auditor/openapi/openapi.yaml @@ -481,6 +481,10 @@ paths: responses: '204': description: Successful Response + '404': + description: Job not Found + '409': + description: Job is not in a terminal state '422': description: Validation Error content: diff --git a/plugins/nemo-customizer/openapi/openapi.yaml b/plugins/nemo-customizer/openapi/openapi.yaml index a2c0b06707..a701d9d8a8 100644 --- a/plugins/nemo-customizer/openapi/openapi.yaml +++ b/plugins/nemo-customizer/openapi/openapi.yaml @@ -246,6 +246,10 @@ paths: responses: '204': description: Successful Response + '404': + description: Job not Found + '409': + description: Job is not in a terminal state '422': description: Validation Error content: @@ -620,6 +624,10 @@ paths: responses: '204': description: Successful Response + '404': + description: Job not Found + '409': + description: Job is not in a terminal state '422': description: Validation Error content: @@ -994,6 +1002,10 @@ paths: responses: '204': description: Successful Response + '404': + description: Job not Found + '409': + description: Job is not in a terminal state '422': description: Validation Error content: diff --git a/plugins/nemo-data-designer/openapi/openapi.yaml b/plugins/nemo-data-designer/openapi/openapi.yaml index ddc30c5e16..c6bec26786 100644 --- a/plugins/nemo-data-designer/openapi/openapi.yaml +++ b/plugins/nemo-data-designer/openapi/openapi.yaml @@ -231,6 +231,10 @@ paths: responses: '204': description: Successful Response + '404': + description: Job not Found + '409': + description: Job is not in a terminal state '422': description: Validation Error content: diff --git a/plugins/nemo-evaluator/openapi/openapi.yaml b/plugins/nemo-evaluator/openapi/openapi.yaml index 85b63964fa..921967fafa 100644 --- a/plugins/nemo-evaluator/openapi/openapi.yaml +++ b/plugins/nemo-evaluator/openapi/openapi.yaml @@ -403,6 +403,10 @@ paths: responses: '204': description: Successful Response + '404': + description: Job not Found + '409': + description: Job is not in a terminal state '422': description: Validation Error content: @@ -907,6 +911,10 @@ paths: responses: '204': description: Successful Response + '404': + description: Job not Found + '409': + description: Job is not in a terminal state '422': description: Validation Error content: diff --git a/plugins/nemo-iron-swarm/openapi/openapi.yaml b/plugins/nemo-iron-swarm/openapi/openapi.yaml index ba772b06ca..13ed53b8ce 100644 --- a/plugins/nemo-iron-swarm/openapi/openapi.yaml +++ b/plugins/nemo-iron-swarm/openapi/openapi.yaml @@ -246,6 +246,10 @@ paths: responses: '204': description: Successful Response + '404': + description: Job not Found + '409': + description: Job is not in a terminal state '422': description: Validation Error content: @@ -1306,6 +1310,10 @@ paths: responses: '204': description: Successful Response + '404': + description: Job not Found + '409': + description: Job is not in a terminal state '422': description: Validation Error content: diff --git a/plugins/nemo-safe-synthesizer/openapi/openapi.yaml b/plugins/nemo-safe-synthesizer/openapi/openapi.yaml index 8b9792ba59..8580a5142e 100644 --- a/plugins/nemo-safe-synthesizer/openapi/openapi.yaml +++ b/plugins/nemo-safe-synthesizer/openapi/openapi.yaml @@ -370,6 +370,10 @@ paths: responses: '204': description: Successful Response + '404': + description: Job not Found + '409': + description: Job is not in a terminal state '422': description: Validation Error content: diff --git a/sdk/python/nemo-platform/.nmpcontext/openapi.yaml b/sdk/python/nemo-platform/.nmpcontext/openapi.yaml index ef4f22e2ae..f13403c86b 100644 --- a/sdk/python/nemo-platform/.nmpcontext/openapi.yaml +++ b/sdk/python/nemo-platform/.nmpcontext/openapi.yaml @@ -5722,6 +5722,9 @@ paths: description: Successful Response '404': description: Job not Found + '409': + description: Job, attempt, or step is not terminal; paused or cancelling + jobs must reach a terminal state before deletion. '422': description: Validation Error content: diff --git a/services/core/entities/src/nmp/core/entities/controllers/workspace_cleanup.py b/services/core/entities/src/nmp/core/entities/controllers/workspace_cleanup.py index fb7198ecdc..8d5b86c48e 100644 --- a/services/core/entities/src/nmp/core/entities/controllers/workspace_cleanup.py +++ b/services/core/entities/src/nmp/core/entities/controllers/workspace_cleanup.py @@ -7,6 +7,7 @@ from nemo_platform import AsyncNeMoPlatform from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import NotFoundError from nemo_platform_plugin.files.client import AsyncFilesClient from nemo_platform_plugin.jobs.client import AsyncJobsClient from nemo_platform_plugin.jobs.schemas import PlatformJobStatus @@ -26,6 +27,24 @@ _TERMINAL_JOB_STATUSES: frozenset[PlatformJobStatus] = frozenset( {PlatformJobStatus.COMPLETED, PlatformJobStatus.ERROR, PlatformJobStatus.CANCELLED} ) +_JOB_TERMINAL_WAIT_TIMEOUT_SECONDS = 300.0 +_JOB_TERMINAL_WAIT_POLL_SECONDS = 2.0 + + +class WorkspaceJobCleanupError(RuntimeError): + """Raised when workspace cleanup cannot delete every job.""" + + +def _job_status_value(status: object) -> PlatformJobStatus | str: + value = getattr(status, "value", status) + try: + return PlatformJobStatus(str(value)) + except ValueError: + return str(value or "") + + +def _job_status_is_terminal(status: object) -> bool: + return _job_status_value(status) in _TERMINAL_JOB_STATUSES class WorkspaceCleanup(HeartbeatMixin, Controller): @@ -132,32 +151,66 @@ async def _cleanup_jobs(self, workspace: Workspace) -> None: jobs_client = client_from_platform(self._nmp_sdk, AsyncJobsClient) jobs = [job async for job in (await jobs_client.list_jobs(workspace=workspace.name)).items()] + cleanup_errors: list[Exception] = [] for job in jobs: - if job.status not in _TERMINAL_JOB_STATUSES: - try: - logger.info(f"Cancelling job: {job.name}") - await jobs_client.cancel_job( - name=job.name, - workspace=workspace.name, - ) - except Exception as e: - logger.warning(f"Failed to cancel job {job.name}: {e}") - try: + if not _job_status_is_terminal(job.status): + cancel_succeeded = False + try: + logger.info(f"Cancelling job: {job.name}") + await jobs_client.cancel_job( + name=job.name, + workspace=workspace.name, + ) + cancel_succeeded = True + except Exception as e: + logger.warning(f"Failed to cancel job {job.name}: {e}") + if cancel_succeeded: + await self._wait_for_terminal_job(jobs_client, workspace.name, job.name) + logger.info(f"Deleting job: {job.name}") await jobs_client.delete_job( name=job.name, workspace=workspace.name, ) + except NotFoundError: + logger.info(f"Job already deleted: {job.name}") except Exception as e: logger.warning(f"Failed to delete job {job.name}: {e}") + cleanup_errors.append(e) finally: self.emit_heartbeat() + if cleanup_errors: + raise WorkspaceJobCleanupError( + f"Failed to delete {len(cleanup_errors)} job(s) while cleaning workspace {workspace.name}" + ) + except Exception as e: - logger.error(f"Failed to list jobs for workspace {workspace.name}: {e}") + logger.error(f"Failed to cleanup jobs for workspace {workspace.name}: {e}") raise + async def _wait_for_terminal_job(self, jobs_client: AsyncJobsClient, workspace: str, job_name: str) -> None: + """Wait until a cancelled job is terminal before hard deletion.""" + deadline = asyncio.get_running_loop().time() + _JOB_TERMINAL_WAIT_TIMEOUT_SECONDS + last_status: PlatformJobStatus | str = "" + while True: + try: + response = await jobs_client.get_job_status(name=job_name, workspace=workspace) + except NotFoundError: + return + status_info = response.data() + last_status = _job_status_value(status_info.status) + if last_status in _TERMINAL_JOB_STATUSES: + return + if asyncio.get_running_loop().time() >= deadline: + raise TimeoutError( + f"Timed out waiting for job {job_name} in workspace {workspace} to reach a terminal status; " + f"last status was {last_status}" + ) + self.emit_heartbeat() + await asyncio.sleep(_JOB_TERMINAL_WAIT_POLL_SECONDS) + @tracer.start_as_current_span("workspace_cleanup/cleanup_deployments") async def _cleanup_deployments(self, workspace: Workspace) -> None: logger.info(f"Cleaning up deployments for workspace: {workspace.name}") diff --git a/services/core/entities/tests/controllers/test_workspace_cleanup.py b/services/core/entities/tests/controllers/test_workspace_cleanup.py index 7944ccc223..e087a4360b 100644 --- a/services/core/entities/tests/controllers/test_workspace_cleanup.py +++ b/services/core/entities/tests/controllers/test_workspace_cleanup.py @@ -6,7 +6,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest -from nmp.core.entities.controllers.workspace_cleanup import WorkspaceCleanup +from nmp.core.entities.controllers.workspace_cleanup import WorkspaceCleanup, WorkspaceJobCleanupError from nmp.core.entities.entities import Workspace, WorkspaceDeletionStage @@ -40,6 +40,18 @@ def _make_mock_files_client(filesets: list | None = None) -> AsyncMock: return mock_files +def _make_response(data) -> MagicMock: + response = MagicMock() + response.data.return_value = data + return response + + +def _make_job_status(status: str) -> MagicMock: + status_response = MagicMock() + status_response.status = status + return status_response + + def _make_jobs_client(jobs: list | None = None) -> MagicMock: """Build a mock typed AsyncJobsClient. @@ -52,6 +64,7 @@ def _make_jobs_client(jobs: list | None = None) -> MagicMock: jobs_client.list_jobs = AsyncMock(return_value=_MockAsyncPaginatedResponse(jobs or [])) jobs_client.cancel_job = AsyncMock() jobs_client.delete_job = AsyncMock() + jobs_client.get_job_status = AsyncMock(return_value=_make_response(_make_job_status("cancelled"))) return jobs_client @@ -307,7 +320,7 @@ async def test_deletes_completed_jobs_without_cancelling(self): jobs_client.delete_job.assert_awaited_once() @pytest.mark.asyncio - async def test_continues_on_individual_job_failure(self): + async def test_attempts_remaining_jobs_then_raises_on_job_delete_failure(self): workspace = _make_workspace() job1 = MagicMock() job1.name = "fail-job" @@ -321,10 +334,55 @@ async def test_continues_on_individual_job_failure(self): controller = _make_controller() with _patch_jobs_client(jobs_client): - await controller._cleanup_jobs(workspace) + with pytest.raises(WorkspaceJobCleanupError): + await controller._cleanup_jobs(workspace) assert jobs_client.delete_job.await_count == 2 + @pytest.mark.asyncio + async def test_non_terminal_job_after_cancel_is_not_treated_as_deleted(self, monkeypatch: pytest.MonkeyPatch): + from nmp.core.entities.controllers import workspace_cleanup as workspace_cleanup_module + + monkeypatch.setattr(workspace_cleanup_module, "_JOB_TERMINAL_WAIT_TIMEOUT_SECONDS", 0.0) + workspace = _make_workspace() + jobs_client = _make_jobs_client([_make_job("still-running-job", status="active")]) + jobs_client.get_job_status = AsyncMock(return_value=_make_response(_make_job_status("active"))) + + controller = _make_controller() + with _patch_jobs_client(jobs_client): + with pytest.raises(WorkspaceJobCleanupError): + await controller._cleanup_jobs(workspace) + + jobs_client.cancel_job.assert_awaited_once_with(name="still-running-job", workspace="test-workspace") + jobs_client.delete_job.assert_not_awaited() + + @pytest.mark.asyncio + async def test_wait_for_terminal_job_emits_heartbeat_between_polls(self, monkeypatch: pytest.MonkeyPatch): + from nmp.core.entities.controllers import workspace_cleanup as workspace_cleanup_module + + monkeypatch.setattr(workspace_cleanup_module, "_JOB_TERMINAL_WAIT_TIMEOUT_SECONDS", 60.0) + monkeypatch.setattr(workspace_cleanup_module, "_JOB_TERMINAL_WAIT_POLL_SECONDS", 1.0) + jobs_client = _make_jobs_client([]) + jobs_client.get_job_status = AsyncMock( + side_effect=[ + _make_response(_make_job_status("active")), + _make_response(_make_job_status("cancelled")), + ] + ) + controller = _make_controller() + + async def _mark_sleep(_delay: float) -> None: + return None + + with ( + patch.object(controller, "emit_heartbeat") as emit_heartbeat, + patch("nmp.core.entities.controllers.workspace_cleanup.asyncio.sleep", side_effect=_mark_sleep), + ): + await controller._wait_for_terminal_job(jobs_client, "test-workspace", "active-job") + + emit_heartbeat.assert_called_once() + assert jobs_client.get_job_status.await_count == 2 + @pytest.mark.asyncio async def test_raises_on_list_failure(self): workspace = _make_workspace() diff --git a/services/core/jobs/src/nmp/core/jobs/api/v2/jobs/endpoints.py b/services/core/jobs/src/nmp/core/jobs/api/v2/jobs/endpoints.py index 2fad1befba..8477456343 100644 --- a/services/core/jobs/src/nmp/core/jobs/api/v2/jobs/endpoints.py +++ b/services/core/jobs/src/nmp/core/jobs/api/v2/jobs/endpoints.py @@ -16,7 +16,7 @@ from nmp.common.api.utils import generate_openapi_extra_params, parse_deep_object from nmp.common.auth import AuthClient, AuthContext, get_auth_client from nmp.common.config import get_platform_config -from nmp.common.entities.client import EntityConflictError, EntityValidationError +from nmp.common.entities.client import EntityConflictError, EntityNotFoundError, EntityValidationError from nmp.common.jobs.docker import validate_gpu_available_for_docker from nmp.common.jobs.exceptions import PlatformJobCompilationError from nmp.common.jobs.log_client import JobLogsClient, dep_job_logs_client @@ -51,9 +51,11 @@ from nmp.core.jobs.app.ctx import JobContext from nmp.core.jobs.app.dispatcher import ( JobAlreadyExistsError, + JobDeletionConflictError, JobDispatcher, JobOutputLocationError, JobSecretValidationError, + JobStatusUpdateSkippedError, StateTransitionConflictError, ) from nmp.core.jobs.app.profiles import ExecutionProfileT @@ -452,6 +454,9 @@ async def resume_job( responses={ status.HTTP_204_NO_CONTENT: {"description": "Successful Response"}, status.HTTP_404_NOT_FOUND: {"description": "Job not Found"}, + status.HTTP_409_CONFLICT: { + "description": "Job, attempt, or step is not terminal; paused or cancelling jobs must reach a terminal state before deletion." + }, }, ) async def delete_job( @@ -461,7 +466,18 @@ async def delete_job( ) -> None: """Delete a platform job.""" with scoped_app_ctx(JobContext(id=name)): - deleted = await dispatcher.delete_job(name, workspace) + try: + deleted = await dispatcher.delete_job(name, workspace) + except JobDeletionConflictError as exc: + logger.info( + "Cannot delete job '%s' in workspace '%s'", + sanitize_for_log(name), + sanitize_for_log(workspace), + ) + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=str(exc), + ) from exc if not deleted: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, @@ -546,9 +562,9 @@ async def page_job_logs( ) try: - filters = { + filters: dict[str, str] = { "job": name, - "job_attempt": attempt_id if attempt_id is not None else job.attempt_id, + "job_attempt": str(attempt_id) if attempt_id is not None else job.attempt_id, } if step_id: filters["job_step"] = step_id @@ -594,13 +610,20 @@ async def create_job_result( status_code=status.HTTP_404_NOT_FOUND, detail=f"Job '{job}' not found in workspace '{workspace}'.", ) - result = await dispatcher.create_result( - job_id=job_entity.id, - result_name=name, - workspace=workspace, - artifact_url=request.artifact_url, - artifact_storage_type=request.artifact_storage_type, - ) + try: + result = await dispatcher.create_result( + job_id=job_entity.id, + job_name=job, + result_name=name, + workspace=workspace, + artifact_url=request.artifact_url, + artifact_storage_type=request.artifact_storage_type, + ) + except EntityNotFoundError as exc: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Job '{job}' not found in workspace '{workspace}'.", + ) from exc return result.to_response() @@ -830,6 +853,15 @@ async def update_job_step_status( status_code=status.HTTP_409_CONFLICT, detail="Conflict updating job step: it was modified by another request. Refresh the step and retry.", ) from exc + except JobStatusUpdateSkippedError: + logger.info( + "Skipping job step status update because the job was deleted", + extra={ + "job": sanitize_for_log(job), + "step": sanitize_for_log(name), + "workspace": sanitize_for_log(workspace), + }, + ) return step_entity @@ -893,13 +925,19 @@ async def update_job_step_task( detail=f"Step '{step}' for job '{job}' not found in workspace '{workspace}'.", ) - return await dispatcher.create_or_update_task( - job, - name, - workspace, - update, - step_entity, - ) + try: + return await dispatcher.create_or_update_task( + job, + name, + workspace, + update, + step_entity, + ) + except EntityNotFoundError as exc: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Step '{step}' for job '{job}' not found in workspace '{workspace}'.", + ) from exc @router.get( diff --git a/services/core/jobs/src/nmp/core/jobs/app/dispatcher.py b/services/core/jobs/src/nmp/core/jobs/app/dispatcher.py index 71e44c013f..bd96de0026 100644 --- a/services/core/jobs/src/nmp/core/jobs/app/dispatcher.py +++ b/services/core/jobs/src/nmp/core/jobs/app/dispatcher.py @@ -4,7 +4,8 @@ import asyncio import json import logging -from typing import Any, Dict, List, Optional, Tuple +import weakref +from typing import Any, Dict, List, Optional, Tuple, TypeVar from nemo_platform import AsyncNeMoPlatform from nemo_platform_plugin.client.adapter import client_from_platform @@ -62,6 +63,14 @@ class JobAlreadyExistsError(ValueError): """Exception raised when creating a job whose name is already in use.""" +class JobDeletionConflictError(ValueError): + """Exception raised when deleting a job that is not safe to hard-delete.""" + + +class JobStatusUpdateSkippedError(Exception): + """Raised when a status update loses a race with job deletion.""" + + class JobSecretValidationError(ValueError): """Exception raised when a job's secret references cannot be validated.""" @@ -77,6 +86,26 @@ class JobOutputLocationError(ValueError): description="Total number of job dispatcher operations performed", ) +_DELETE_PAGE_SIZE = 1000 +_JOB_MUTATION_LOCKS: weakref.WeakValueDictionary[tuple[str, str], asyncio.Lock] = weakref.WeakValueDictionary() +EntityT = TypeVar("EntityT") + + +def _get_job_mutation_lock(job_name: str, workspace: str) -> asyncio.Lock: + """Return the process-local lock for a job's mutable state. + + This serializes delete, rerun, status updates, and same-process associated + record writes inside one API worker. It does not coordinate across API + replicas; multi-replica deployments need a store-level lease or conditional + delete to close that gap durably. + """ + key = (workspace, job_name) + lock = _JOB_MUTATION_LOCKS.get(key) + if lock is None: + lock = asyncio.Lock() + _JOB_MUTATION_LOCKS[key] = lock + return lock + def create_platform_job_response(job: PlatformJob, attempt: PlatformJobAttempt) -> PlatformJobResponse: """Helper to create PlatformJobResponse from job and attempt entities.""" @@ -105,6 +134,12 @@ def create_platform_job_response(job: PlatformJob, attempt: PlatformJobAttempt) ) +def _format_status_for_message(status: PlatformJobStatus | str) -> str: + if isinstance(status, PlatformJobStatus): + return status.value + return str(status) + + # Status lives on PlatformJobAttempt, not PlatformJob, so it cannot be # resolved by the PlatformJob entity-store query. Instead, the full filter tree # is evaluated in-memory (op.apply(InMemoryFilterRepository(virtual_job))) against @@ -259,6 +294,20 @@ async def create_job( sdk: Optional[AsyncNeMoPlatform] = None, ) -> PlatformJobResponse: """Create a new job and its first step.""" + if job_req.name is None: + return await self._create_job(job_req, workspace, auth_context=auth_context, sdk=sdk) + + async with _get_job_mutation_lock(job_req.name, workspace): + return await self._create_job(job_req, workspace, auth_context=auth_context, sdk=sdk) + + async def _create_job( + self, + job_req: CreatePlatformJobRequest, + workspace: str, + auth_context: Optional[AuthContext] = None, + sdk: Optional[AsyncNeMoPlatform] = None, + ) -> PlatformJobResponse: + """Create a new job after the caller has acquired any needed name lock.""" job_name = job_req.name if job_name is not None: # Check if a job with the same name already exists @@ -359,7 +408,7 @@ async def get_job(self, job_name: str, workspace: str) -> PlatformJobResponse | except EntityNotFoundError: return None try: - attempt = await self.store.get_by_id(PlatformJobAttempt, job_entity.current_attempt_id) # type: ignore + attempt = await self.store.get_by_id(PlatformJobAttempt, job_entity.current_attempt_id) except EntityNotFoundError: return None return create_platform_job_response(job_entity, attempt) @@ -432,6 +481,11 @@ async def delete_job(self, job_name: str, workspace: str) -> bool: Returns: True if job was deleted, False if job was not found. """ + async with _get_job_mutation_lock(job_name, workspace): + return await self._delete_job_locked(job_name, workspace) + + async def _delete_job_locked(self, job_name: str, workspace: str) -> bool: + """Delete a terminal job while holding the per-job mutation lock.""" try: job_entity = await self.store.get(PlatformJob, job_name, workspace=workspace) except EntityNotFoundError: @@ -441,42 +495,55 @@ async def delete_job(self, job_name: str, workspace: str) -> bool: try: # Get all attempts - attempts_response = await self.store.list( + attempts = await self._list_all_entities( PlatformJobAttempt, filter_obj={"job": job_entity.id}, - page_size=1000, workspace=workspace, ) + steps_by_attempt: dict[str, list[PlatformJobStep]] = {} + for attempt in attempts: + if not attempt.status.is_terminal(): + raise JobDeletionConflictError( + f"Cannot delete job '{job_entity.name}' while it is " + f"'{_format_status_for_message(attempt.status)}'. Cancel the job and wait for it to reach a " + "terminal state before deleting." + ) + + steps = await self._list_all_entities( + PlatformJobStep, + filter_obj={"attempt_id": attempt.id}, + workspace=workspace, + ) + steps_by_attempt[attempt.id] = steps + + for step in steps: + if not step.status.is_terminal(): + raise JobDeletionConflictError( + f"Cannot delete job '{job_entity.name}' while step '{step.name}' is " + f"'{_format_status_for_message(step.status)}'. Cancel the job and wait for it to reach a " + "terminal state before deleting." + ) # Get all results - results_response = await self.store.list( + results = await self._list_all_entities( PlatformJobResult, filter_obj={"job": job_entity.id}, - page_size=1000, workspace=workspace, ) - for result in results_response.data: + for result in results: await self.store.delete_by_id(PlatformJobResult, result.id) # Delete steps and tasks for each attempt - for attempt in attempts_response.data: - steps_response = await self.store.list( - PlatformJobStep, - filter_obj={"attempt_id": attempt.id}, - page_size=1000, - workspace=workspace, - ) - - for step in steps_response.data: + for attempt in attempts: + for step in steps_by_attempt[attempt.id]: # Delete tasks - tasks_response = await self.store.list( + tasks = await self._list_all_entities( PlatformJobTask, filter_obj={"step_id": step.id}, - page_size=1000, workspace=workspace, ) - for task in tasks_response.data: + for task in tasks: await self.store.delete_by_id(PlatformJobTask, task.id) # Delete step @@ -508,10 +575,36 @@ async def delete_job(self, job_name: str, workspace: str) -> bool: operations_counter.add(1, attributes={"operation": "delete_job"}) return True + except JobDeletionConflictError: + logger.info("Refusing to delete non-terminal job", extra=extras) + raise except Exception as e: logger.exception("Error deleting job", extra=extras) raise e + async def _list_all_entities( + self, + entity_type: type[EntityT], + *, + filter_obj: dict[str, Any], + workspace: str, + ) -> list[EntityT]: + """Return every page for a filtered entity query.""" + page = 1 + entities: list[EntityT] = [] + while True: + response = await self.store.list( + entity_type, + filter_obj=filter_obj, + page=page, + page_size=_DELETE_PAGE_SIZE, + workspace=workspace, + ) + entities.extend(response.data) + if response.pagination.page >= response.pagination.total_pages: + return entities + page = response.pagination.page + 1 + # ========================================================================= # Attempt Operations # ========================================================================= @@ -849,6 +942,25 @@ async def update_job_status_from_step( On EntityConflictError (e.g. reconciler updated the step concurrently), refetches the step and retries once if the requested transition is still valid. """ + attempt = await self.get_attempt(step.attempt_id) + if attempt is None: + raise JobStatusUpdateSkippedError(f"Attempt does not exist: {step.attempt_id}") + try: + job = await self.store.get_by_id(PlatformJob, attempt.job) + except EntityNotFoundError as exc: + raise JobStatusUpdateSkippedError(f"Job does not exist: {attempt.job}") from exc + + async with _get_job_mutation_lock(job.name, job.workspace): + return await self._update_job_status_from_step_locked(step, status, status_details, error_details) + + async def _update_job_status_from_step_locked( + self, + step: PlatformJobStep, + status: PlatformJobStatus, + status_details: Optional[Dict[str, Any]] = None, + error_details: Optional[Dict[str, Any]] = None, + ) -> tuple[PlatformJobStep, PlatformJobAttempt]: + """Update a job from a step while holding the per-job mutation lock.""" step_to_save = step saved_step: PlatformJobStep | None = None for attempt in range(2): @@ -891,7 +1003,7 @@ async def update_job_status_from_step( # Update job / attempt status from step attempt = await self.get_attempt(saved_step.attempt_id) if attempt is None: - raise Exception(f"Attempt does not exist: {saved_step.attempt_id}") + raise JobStatusUpdateSkippedError(f"Attempt does not exist: {saved_step.attempt_id}") # Determine new attempt status new_attempt_status = attempt.status @@ -1074,7 +1186,11 @@ async def cancel_job(self, job_name: str, workspace: str) -> PlatformJobResponse async def rerun_job(self, job_name: str, workspace: str) -> PlatformJobResponse | None: """Re-run a job.""" + async with _get_job_mutation_lock(job_name, workspace): + return await self._rerun_job_locked(job_name, workspace) + async def _rerun_job_locked(self, job_name: str, workspace: str) -> PlatformJobResponse | None: + """Re-run a terminal job while holding the per-job mutation lock.""" try: job_entity = await self.store.get(PlatformJob, job_name, workspace=workspace) except EntityNotFoundError: @@ -1191,6 +1307,16 @@ async def create_or_update_task( self, job_name: str, task_name: str, workspace: str, task_update: PlatformJobTaskUpdate, step: PlatformJobStep ) -> PlatformJobTask: """Create or update a task for a job step.""" + async with _get_job_mutation_lock(job_name, workspace): + return await self._create_or_update_task_locked(job_name, task_name, workspace, task_update, step) + + async def _create_or_update_task_locked( + self, job_name: str, task_name: str, workspace: str, task_update: PlatformJobTaskUpdate, step: PlatformJobStep + ) -> PlatformJobTask: + """Create or update a task while holding the per-job mutation lock.""" + await self.store.get(PlatformJob, job_name, workspace=workspace) + await self.store.get_by_id(PlatformJobStep, step.id) + task = await self.get_task(step.id, task_name, workspace=workspace) if not task: @@ -1279,6 +1405,7 @@ def is_step_cancelled(self, step: PlatformJobStep) -> bool: async def create_result( self, job_id: str, + job_name: str, result_name: str, artifact_url: str, artifact_storage_type: Any, @@ -1293,14 +1420,18 @@ async def create_result( artifact_storage_type: Type of artifact storage workspace: Workspace for the result """ - result = PlatformJobResult( - name=result_name, - workspace=workspace, - job=job_id, # Use entity ID for parent FK - artifact_url=artifact_url, - artifact_storage_type=artifact_storage_type, - ) - return await self.store.create(result) + async with _get_job_mutation_lock(job_name, workspace): + job = await self.store.get(PlatformJob, job_name, workspace=workspace) + if job.id != job_id: + raise EntityNotFoundError(f"Job ID mismatch for {workspace}/{job_name}") + result = PlatformJobResult( + name=result_name, + workspace=workspace, + job=job_id, # Use entity ID for parent FK + artifact_url=artifact_url, + artifact_storage_type=artifact_storage_type, + ) + return await self.store.create(result) async def get_result(self, job_name: str, result_name: str, workspace: str) -> Optional[PlatformJobResult]: """Get a platform job result.""" diff --git a/services/core/jobs/src/nmp/core/jobs/controllers/backends/docker.py b/services/core/jobs/src/nmp/core/jobs/controllers/backends/docker.py index 10906df55d..a8f8a6744c 100644 --- a/services/core/jobs/src/nmp/core/jobs/controllers/backends/docker.py +++ b/services/core/jobs/src/nmp/core/jobs/controllers/backends/docker.py @@ -2325,6 +2325,8 @@ def configure_container(self, container_args: dict, executor_config: GPUExecutio try: gpu_ids = self.gpu_pool.allocate_gpu(container_args["labels"][JOB_STEP_ID_LABEL], num_requested=num_gpus) except GPUAllocationError as e: + if e.is_transient_capacity_exhaustion: + raise SchedulingDeferred(str(e)) from e raise ResourceAllocationError(str(e)) from e container_args["device_requests"] = [ diff --git a/services/core/jobs/src/nmp/core/jobs/controllers/scheduler.py b/services/core/jobs/src/nmp/core/jobs/controllers/scheduler.py index f78ed8819e..a60c300231 100644 --- a/services/core/jobs/src/nmp/core/jobs/controllers/scheduler.py +++ b/services/core/jobs/src/nmp/core/jobs/controllers/scheduler.py @@ -162,6 +162,22 @@ def step(self): "reason": e.message, }, ) + try: + self._update_step_status_with_timing( + step=step, + phase="scheduling_deferred", + status=PlatformJobStatus(step.status), + status_details={"message": e.message}, + ) + except Exception: + logger.exception( + "Could not persist scheduling deferral details", + extra={ + "job": step.job, + "step": step.name, + "workspace": step.workspace, + }, + ) except Exception as e: logger.exception("Could not schedule job step", exc_info=True) log_job_diagnostics_if_debug( diff --git a/services/core/jobs/tests/controllers/test_docker_backend.py b/services/core/jobs/tests/controllers/test_docker_backend.py index f99d79e3a5..16935cc549 100644 --- a/services/core/jobs/tests/controllers/test_docker_backend.py +++ b/services/core/jobs/tests/controllers/test_docker_backend.py @@ -1195,7 +1195,7 @@ def test_docker_workload_identity_token_timing_constraints(monkeypatch): def test_schedule_docker_gpu(mock_nmp_client, docker_client_mock): - """Test successful job scheduling.""" + """Test GPU job scheduling defers when the pool is temporarily full.""" gpus = 2 gpu_executor_config = GPUExecutionProvider.model_validate( @@ -1285,8 +1285,8 @@ def test_schedule_docker_gpu(mock_nmp_client, docker_client_mock): ) executor._client = docker_client_mock - # whichever one is the third will fail - with pytest.raises(ResourceAllocationError): + # Whichever one is the third will defer until another step releases GPUs. + with pytest.raises(SchedulingDeferred): executor.schedule(executor_config=gpu_executor_config, step=step) executor.schedule(executor_config=gpu_executor_config, step=step_two) executor.schedule(executor_config=gpu_executor_config, step=step_three) @@ -1333,6 +1333,69 @@ def test_schedule_docker_gpu(mock_nmp_client, docker_client_mock): assert len([v for v in executor.gpu_pool.gpu_to_workload_id.values() if v is None]) == 1 +def test_gpu_configure_container_defers_when_pool_is_temporarily_full(mock_nmp_client, docker_client_mock): + """Full but sufficient GPU pools should defer scheduling instead of erroring.""" + gpu_executor_config = GPUExecutionProvider.model_validate( + { + "provider": "gpu", + "profile": "default", + "container": {"image": "hello-world:latest"}, + "resources": {"num_gpus": 1}, + "config": {}, + } + ) + + with patch("nmp.core.jobs.controllers.backends.docker.SharedResourceManager") as mock_srm: + mock_pool = DockerGPUPool(reserved_gpu_device_ids=[0]) + mock_pool.allocate_gpu("already-running", num_requested=1) + mock_srm.get_instance.return_value.get_gpu_pool.return_value = mock_pool + + executor = GPUDockerJobBackend( + nmp_sdk=mock_nmp_client, + execution_profile_config=DockerJobExecutionProfileConfig( + storage=DockerJobStorageConfig(volume_name="test_jobs_storage"), + ), + profile_name="default", + ) + executor._client = docker_client_mock + + with pytest.raises(SchedulingDeferred, match="Not enough GPUs available"): + executor.configure_container({"labels": {JOB_STEP_ID_LABEL: "deferred-step"}}, gpu_executor_config) + + assert executor.gpu_pool.gpu_to_workload_id == {0: "already-running"} + + +def test_gpu_configure_container_errors_when_request_exceeds_pool(mock_nmp_client, docker_client_mock): + """GPU requests larger than the pool are permanent resource allocation errors.""" + gpu_executor_config = GPUExecutionProvider.model_validate( + { + "provider": "gpu", + "profile": "default", + "container": {"image": "hello-world:latest"}, + "resources": {"num_gpus": 2}, + "config": {}, + } + ) + + with patch("nmp.core.jobs.controllers.backends.docker.SharedResourceManager") as mock_srm: + mock_pool = DockerGPUPool(reserved_gpu_device_ids=[0]) + mock_srm.get_instance.return_value.get_gpu_pool.return_value = mock_pool + + executor = GPUDockerJobBackend( + nmp_sdk=mock_nmp_client, + execution_profile_config=DockerJobExecutionProfileConfig( + storage=DockerJobStorageConfig(volume_name="test_jobs_storage"), + ), + profile_name="default", + ) + executor._client = docker_client_mock + + with pytest.raises(ResourceAllocationError, match="Requested 2"): + executor.configure_container({"labels": {JOB_STEP_ID_LABEL: "oversized-step"}}, gpu_executor_config) + + assert executor.gpu_pool.gpu_to_workload_id == {0: None} + + def test_gpu_cleanup_on_job_completion(mock_nmp_client, docker_client_mock): """Test that GPU resources are released when a job completes successfully.""" diff --git a/services/core/jobs/tests/controllers/test_scheduler.py b/services/core/jobs/tests/controllers/test_scheduler.py index 932da1b03b..c34e3749ae 100644 --- a/services/core/jobs/tests/controllers/test_scheduler.py +++ b/services/core/jobs/tests/controllers/test_scheduler.py @@ -30,15 +30,20 @@ def job_scheduler(backend_registry: BackendRegistry, mock_nmp_client) -> JobSche return JobScheduler(backend_registry, mock_nmp_client) +@fixture +def test_step_created(test_step_pending: PlatformJobStepWithContext) -> PlatformJobStepWithContext: + return test_step_pending.model_copy(update={"status": PlatformJobStatus.CREATED}) + + def test_does_schedule_job( job_scheduler: JobScheduler, backend_registry: BackendRegistry, mock_nmp_client, mock_jobs_client, - test_step_pending: PlatformJobStepWithContext, + test_step_created: PlatformJobStepWithContext, ): # Mock the jobs list response - mock_jobs_client.list_steps.return_value = paginated_response([test_step_pending]) + mock_jobs_client.list_steps.return_value = paginated_response([test_step_created]) # Get the test backend from the registry backend = backend_registry.get_backend(provider="cpu", profile="default") @@ -61,32 +66,82 @@ def test_does_schedule_job( # Test backend should have received one schedule call for our test job assert len(test_backend.mock.schedule_calls) == 1 - assert test_backend.mock.schedule_calls[0]["step"].id == test_step_pending.id + assert test_backend.mock.schedule_calls[0]["step"].id == test_step_created.id assert test_backend.mock.sync_calls == [] -def test_scheduling_deferred_leaves_step_created( +def test_scheduling_deferred_keeps_step_created_with_visible_status_details( job_scheduler: JobScheduler, mock_nmp_client, mock_jobs_client, - test_step_pending: PlatformJobStepWithContext, + test_step_created: PlatformJobStepWithContext, +): + mock_jobs_client.list_steps.return_value = paginated_response([test_step_created]) + + with patch.object(job_scheduler, "schedule_step", side_effect=SchedulingDeferred("capacity full")): + job_scheduler.step() + + mock_jobs_client.update_job_step_status.assert_called_once() + call = mock_jobs_client.update_job_step_status.call_args + assert call.kwargs["name"] == test_step_created.name + assert call.kwargs["workspace"] == test_step_created.workspace + assert call.kwargs["job"] == test_step_created.job + body = call.kwargs["body"] + assert body.status == PlatformJobStatus.CREATED + assert body.status_details == {"message": "capacity full"} + + +def test_scheduling_deferred_preserves_step_resuming_with_visible_status_details( + job_scheduler: JobScheduler, + mock_nmp_client, + mock_jobs_client, + test_step_resuming: PlatformJobStepWithContext, ): - mock_jobs_client.list_steps.return_value = paginated_response([test_step_pending]) + mock_jobs_client.list_steps.return_value = paginated_response([test_step_resuming]) with patch.object(job_scheduler, "schedule_step", side_effect=SchedulingDeferred("capacity full")): job_scheduler.step() - mock_jobs_client.update_job_step_status.assert_not_called() + mock_jobs_client.update_job_step_status.assert_called_once() + call = mock_jobs_client.update_job_step_status.call_args + assert call.kwargs["name"] == test_step_resuming.name + assert call.kwargs["workspace"] == test_step_resuming.workspace + assert call.kwargs["job"] == test_step_resuming.job + body = call.kwargs["body"] + assert body.status == PlatformJobStatus.RESUMING + assert body.status_details == {"message": "capacity full"} + + +def test_scheduling_deferred_status_write_failure_does_not_skip_later_steps( + job_scheduler: JobScheduler, + mock_nmp_client, + mock_jobs_client, + test_step_created: PlatformJobStepWithContext, +): + next_step = test_step_created.model_copy(update={"id": "next-step-id", "name": "next-step", "job": "next-job"}) + mock_jobs_client.list_steps.return_value = paginated_response([test_step_created, next_step]) + mock_jobs_client.update_job_step_status.side_effect = [RuntimeError("persist failed"), data_response(None)] + + with patch.object(job_scheduler, "schedule_step", side_effect=SchedulingDeferred("capacity full")) as schedule_step: + job_scheduler.step() + + assert schedule_step.call_count == 2 + assert mock_jobs_client.update_job_step_status.call_count == 2 + first_body = mock_jobs_client.update_job_step_status.call_args_list[0].kwargs["body"] + second_body = mock_jobs_client.update_job_step_status.call_args_list[1].kwargs["body"] + assert first_body.status == PlatformJobStatus.CREATED + assert second_body.status == PlatformJobStatus.CREATED + assert second_body.status_details == {"message": "capacity full"} def test_resource_allocation_error_marks_step_as_error( job_scheduler: JobScheduler, mock_nmp_client, mock_jobs_client, - test_step_pending: PlatformJobStepWithContext, + test_step_created: PlatformJobStepWithContext, ): """When ResourceAllocationError is raised (e.g. no GPUs), scheduler marks step as error with error_details.""" - mock_jobs_client.list_steps.return_value = paginated_response([test_step_pending]) + mock_jobs_client.list_steps.return_value = paginated_response([test_step_created]) error_message = "No GPUs available on this system. GPU jobs require a system with NVIDIA GPUs." with patch.object(job_scheduler, "schedule_step", side_effect=ResourceAllocationError(error_message)): @@ -94,9 +149,9 @@ def test_resource_allocation_error_marks_step_as_error( mock_jobs_client.update_job_step_status.assert_called_once() call = mock_jobs_client.update_job_step_status.call_args - assert call.kwargs["name"] == test_step_pending.name - assert call.kwargs["workspace"] == test_step_pending.workspace - assert call.kwargs["job"] == test_step_pending.job + assert call.kwargs["name"] == test_step_created.name + assert call.kwargs["workspace"] == test_step_created.workspace + assert call.kwargs["job"] == test_step_created.job body = call.kwargs["body"] assert body.status == PlatformJobStatus.ERROR assert body.status_details == {"message": error_message} @@ -107,9 +162,9 @@ def test_scheduler_logs_diagnostics_for_unexpected_schedule_error_in_debug_mode( job_scheduler: JobScheduler, mock_nmp_client, mock_jobs_client, - test_step_pending: PlatformJobStepWithContext, + test_step_created: PlatformJobStepWithContext, ): - mock_jobs_client.list_steps.return_value = paginated_response([test_step_pending]) + mock_jobs_client.list_steps.return_value = paginated_response([test_step_created]) with ( patch.object(job_scheduler, "schedule_step", side_effect=RuntimeError("boom")), @@ -120,7 +175,7 @@ def test_scheduler_logs_diagnostics_for_unexpected_schedule_error_in_debug_mode( log_diagnostics.assert_called_once_with( mock_nmp_client, - test_step_pending, + test_step_created, logger=job_scheduler._logger, context="unexpected scheduling error", ) @@ -130,14 +185,14 @@ def test_scheduler_does_not_mark_step_error_when_pending_update_conflicts_with_c job_scheduler: JobScheduler, mock_nmp_client, mock_jobs_client, - test_step_pending: PlatformJobStepWithContext, + test_step_created: PlatformJobStepWithContext, ): - mock_jobs_client.list_steps.return_value = paginated_response([test_step_pending]) + mock_jobs_client.list_steps.return_value = paginated_response([test_step_created]) conflict = _conflict_error( "Invalid status transition from PlatformJobStatus.ACTIVE to PlatformJobStatus.PENDING for step test-step-id" ) - active_step = test_step_pending.model_copy(update={"status": PlatformJobStatus.ACTIVE}) + active_step = test_step_created.model_copy(update={"status": PlatformJobStatus.ACTIVE}) mock_jobs_client.update_job_step_status.side_effect = [conflict] get_step_resp = active_step mock_jobs_client.get_job_step.return_value.data.return_value = get_step_resp @@ -146,15 +201,15 @@ def test_scheduler_does_not_mark_step_error_when_pending_update_conflicts_with_c mock_jobs_client.update_job_step_status.assert_called_once() update_call = mock_jobs_client.update_job_step_status.call_args - assert update_call.kwargs["name"] == test_step_pending.name - assert update_call.kwargs["workspace"] == test_step_pending.workspace - assert update_call.kwargs["job"] == test_step_pending.job + assert update_call.kwargs["name"] == test_step_created.name + assert update_call.kwargs["workspace"] == test_step_created.workspace + assert update_call.kwargs["job"] == test_step_created.job assert update_call.kwargs["body"].status == PlatformJobStatus.PENDING mock_jobs_client.get_job_step.assert_called_once_with( - name=test_step_pending.name, - workspace=test_step_pending.workspace, - job=test_step_pending.job, + name=test_step_created.name, + workspace=test_step_created.workspace, + job=test_step_created.job, ) diff --git a/services/core/jobs/tests/test_dispatcher.py b/services/core/jobs/tests/test_dispatcher.py index 750223b18d..18cc2e6303 100644 --- a/services/core/jobs/tests/test_dispatcher.py +++ b/services/core/jobs/tests/test_dispatcher.py @@ -7,6 +7,8 @@ (create, cancel, pause, resume, rerun, delete) using the EntityStore pattern. """ +import asyncio +import gc import json from unittest.mock import AsyncMock, MagicMock, patch @@ -20,13 +22,13 @@ EntityConflictError, EntityNotFoundError, ) -from nmp.common.jobs.schemas import PlatformJobStatus +from nmp.common.jobs.schemas import FileStorageType, PlatformJobStatus from nmp.core.jobs.api.v2.jobs.schemas import ( CreatePlatformJobRequest, PlatformJobResponse, PlatformJobTaskUpdate, ) -from nmp.core.jobs.app.dispatcher import JobDispatcher +from nmp.core.jobs.app.dispatcher import JobDeletionConflictError, JobDispatcher, JobStatusUpdateSkippedError from nmp.core.jobs.app.schemas import ( PlatformJobStepSpec, ) @@ -113,7 +115,7 @@ async def create_test_job_data( workspace=DEFAULT_WORKSPACE, job=saved_job.id, artifact_url="default/test-fileset#artifact", - artifact_storage_type="fileset", + artifact_storage_type=FileStorageType.FILESET, ) saved_result = await store.add(result) @@ -188,6 +190,318 @@ async def test_delete_job_nonexistent_job(mock_dispatcher: JobDispatcher): assert deleted is False +@pytest.mark.asyncio +async def test_delete_job_non_terminal_job_raises_without_deleting_data( + mock_dispatcher: JobDispatcher, + mock_store: EntityClient, + _mock_files_client, +): + """Deleting an active job is refused before metadata or fileset cleanup.""" + job_id, job_name, attempt_id, step_id, _, _ = await create_test_job_data(mock_store, "active-delete-test-job") + + attempt = await mock_store.get_by_id(PlatformJobAttempt, attempt_id) + attempt.status = PlatformJobStatus.ACTIVE + await mock_store.update(attempt) + + step = await mock_store.get_by_id(PlatformJobStep, step_id) + step.status = PlatformJobStatus.ACTIVE + await mock_store.update(step) + + with pytest.raises(JobDeletionConflictError, match="Cancel the job and wait"): + await mock_dispatcher.delete_job(job_name, DEFAULT_WORKSPACE) + + await verify_job_data_exists(mock_store, job_id, should_exist=True) + assert await count_entities(mock_store, PlatformJobAttempt, {"job": job_id}) == 1 + assert await count_entities(mock_store, PlatformJobStep, {"attempt_id": attempt_id}) == 1 + assert await count_entities(mock_store, PlatformJobTask, {"step_id": step_id}) == 1 + assert await count_entities(mock_store, PlatformJobResult, {"job": job_id}) == 1 + _mock_files_client.delete_fileset.assert_not_called() + + +@pytest.mark.asyncio +async def test_delete_job_deletes_related_entities_across_all_pages( + mock_dispatcher: JobDispatcher, + mock_store: EntityClient, + monkeypatch: pytest.MonkeyPatch, +): + """Deleting a terminal job visits every attempt, step, task, and result page.""" + from nmp.core.jobs.app import dispatcher as dispatcher_module + + monkeypatch.setattr(dispatcher_module, "_DELETE_PAGE_SIZE", 1) + + job_id, job_name, attempt_id, step_id, task_id, result_id = await create_test_job_data( + mock_store, "paginated-delete-test-job" + ) + job = await mock_store.get_by_id(PlatformJob, job_id) + + second_attempt = await mock_store.add( + PlatformJobAttempt( + name="attempt-2", + workspace=DEFAULT_WORKSPACE, + job=job.id, + seq=2, + status=PlatformJobStatus.COMPLETED, + spec={}, + platform_spec=job.platform_spec, + ) + ) + second_step = await mock_store.add( + PlatformJobStep( + name="step-2", + workspace=DEFAULT_WORKSPACE, + attempt_id=second_attempt.id, + status=PlatformJobStatus.COMPLETED, + ) + ) + second_task = await mock_store.add( + PlatformJobTask( + name="task-2", + workspace=DEFAULT_WORKSPACE, + step_id=second_step.id, + status=PlatformJobStatus.COMPLETED, + ) + ) + second_result = await mock_store.add( + PlatformJobResult( + name="result-2", + workspace=DEFAULT_WORKSPACE, + job=job.id, + artifact_url="default/test-fileset#artifact-2", + artifact_storage_type=FileStorageType.FILESET, + ) + ) + + deleted = await mock_dispatcher.delete_job(job_name, DEFAULT_WORKSPACE) + + assert deleted is True + for entity_type, entity_id in ( + (PlatformJob, job_id), + (PlatformJobAttempt, attempt_id), + (PlatformJobAttempt, second_attempt.id), + (PlatformJobStep, step_id), + (PlatformJobStep, second_step.id), + (PlatformJobTask, task_id), + (PlatformJobTask, second_task.id), + (PlatformJobResult, result_id), + (PlatformJobResult, second_result.id), + ): + with pytest.raises(EntityNotFoundError): + await mock_store.get_by_id(entity_type, entity_id) + + +@pytest.mark.asyncio +async def test_delete_job_serializes_with_rerun_job( + mock_dispatcher: JobDispatcher, + mock_store: EntityClient, + mock_nmp_client, +): + """A rerun request cannot create a new attempt while deletion is cleaning up.""" + _, job_name, _, _, _, _ = await create_test_job_data(mock_store, "delete-rerun-lock-test-job") + other_dispatcher = JobDispatcher(store=mock_store, sdk=mock_nmp_client) + + delete_started = asyncio.Event() + allow_delete = asyncio.Event() + original_delete_by_id = mock_store.delete_by_id + + async def blocking_delete_by_id(entity_type, entity_id): + if entity_type is PlatformJobResult and not delete_started.is_set(): + delete_started.set() + await allow_delete.wait() + return await original_delete_by_id(entity_type, entity_id) + + with patch.object(mock_store, "delete_by_id", side_effect=blocking_delete_by_id): + delete_task = asyncio.create_task(mock_dispatcher.delete_job(job_name, DEFAULT_WORKSPACE)) + rerun_task = None + try: + await asyncio.wait_for(delete_started.wait(), timeout=1.0) + + rerun_task = asyncio.create_task(other_dispatcher.rerun_job(job_name, DEFAULT_WORKSPACE)) + await asyncio.sleep(0.05) + assert not rerun_task.done() + + allow_delete.set() + assert await delete_task is True + assert await rerun_task is None + finally: + allow_delete.set() + if rerun_task is not None: + await asyncio.gather(delete_task, rerun_task, return_exceptions=True) + else: + await asyncio.gather(delete_task, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_delete_job_serializes_with_same_name_create( + mock_dispatcher: JobDispatcher, + mock_store: EntityClient, + mock_nmp_client, + sample_platform_job_request: CreatePlatformJobRequest, +): + """A same-name create waits until delete finishes all cleanup for the old job.""" + job_id, job_name, _, _, _, _ = await create_test_job_data(mock_store, "delete-create-lock-test-job") + other_dispatcher = JobDispatcher(store=mock_store, sdk=mock_nmp_client) + create_request = sample_platform_job_request.model_copy(update={"name": job_name}) + + delete_started = asyncio.Event() + allow_delete = asyncio.Event() + original_delete_by_id = mock_store.delete_by_id + + async def blocking_delete_by_id(entity_type, entity_id): + if entity_type is PlatformJobResult and not delete_started.is_set(): + delete_started.set() + await allow_delete.wait() + return await original_delete_by_id(entity_type, entity_id) + + with patch.object(mock_store, "delete_by_id", side_effect=blocking_delete_by_id): + delete_task = asyncio.create_task(mock_dispatcher.delete_job(job_name, DEFAULT_WORKSPACE)) + create_task = None + try: + await asyncio.wait_for(delete_started.wait(), timeout=1.0) + + create_task = asyncio.create_task(other_dispatcher.create_job(create_request, DEFAULT_WORKSPACE)) + await asyncio.sleep(0.05) + assert not create_task.done() + + allow_delete.set() + assert await delete_task is True + recreated = await create_task + finally: + allow_delete.set() + if create_task is not None: + await asyncio.gather(delete_task, create_task, return_exceptions=True) + else: + await asyncio.gather(delete_task, return_exceptions=True) + + await verify_job_data_exists(mock_store, job_id, should_exist=False) + assert recreated.name == job_name + assert recreated.id != job_id + + +@pytest.mark.asyncio +async def test_delete_job_serializes_with_task_creation( + mock_dispatcher: JobDispatcher, + mock_store: EntityClient, + mock_nmp_client, +): + """A task update cannot create a late child row after delete cleanup starts.""" + job_id, job_name, _, step_id, _, _ = await create_test_job_data(mock_store, "delete-task-lock-test-job") + step = await mock_store.get_by_id(PlatformJobStep, step_id) + other_dispatcher = JobDispatcher(store=mock_store, sdk=mock_nmp_client) + + delete_started = asyncio.Event() + allow_delete = asyncio.Event() + original_delete_by_id = mock_store.delete_by_id + + async def blocking_delete_by_id(entity_type, entity_id): + if entity_type is PlatformJobResult and not delete_started.is_set(): + delete_started.set() + await allow_delete.wait() + return await original_delete_by_id(entity_type, entity_id) + + with patch.object(mock_store, "delete_by_id", side_effect=blocking_delete_by_id): + delete_task = asyncio.create_task(mock_dispatcher.delete_job(job_name, DEFAULT_WORKSPACE)) + task_create_task = None + try: + await asyncio.wait_for(delete_started.wait(), timeout=1.0) + + task_create_task = asyncio.create_task( + other_dispatcher.create_or_update_task( + job_name, + "late-task", + DEFAULT_WORKSPACE, + PlatformJobTaskUpdate(status=PlatformJobStatus.ACTIVE), + step, + ) + ) + await asyncio.sleep(0.05) + assert not task_create_task.done() + + allow_delete.set() + assert await delete_task is True + with pytest.raises(EntityNotFoundError): + await task_create_task + finally: + allow_delete.set() + if task_create_task is not None: + await asyncio.gather(delete_task, task_create_task, return_exceptions=True) + else: + await asyncio.gather(delete_task, return_exceptions=True) + + await verify_job_data_exists(mock_store, job_id, should_exist=False) + assert await count_entities(mock_store, PlatformJobTask, {"step_id": step_id}) == 0 + + +@pytest.mark.asyncio +async def test_delete_job_serializes_with_result_creation( + mock_dispatcher: JobDispatcher, + mock_store: EntityClient, + mock_nmp_client, +): + """A result create cannot recreate associated data after delete cleanup starts.""" + job_id, job_name, _, _, _, _ = await create_test_job_data(mock_store, "delete-result-lock-test-job") + other_dispatcher = JobDispatcher(store=mock_store, sdk=mock_nmp_client) + + delete_started = asyncio.Event() + allow_delete = asyncio.Event() + original_delete_by_id = mock_store.delete_by_id + + async def blocking_delete_by_id(entity_type, entity_id): + if entity_type is PlatformJobResult and not delete_started.is_set(): + delete_started.set() + await allow_delete.wait() + return await original_delete_by_id(entity_type, entity_id) + + with patch.object(mock_store, "delete_by_id", side_effect=blocking_delete_by_id): + delete_task = asyncio.create_task(mock_dispatcher.delete_job(job_name, DEFAULT_WORKSPACE)) + result_create_task = None + try: + await asyncio.wait_for(delete_started.wait(), timeout=1.0) + + result_create_task = asyncio.create_task( + other_dispatcher.create_result( + job_id=job_id, + job_name=job_name, + result_name="late-result", + artifact_url="default/test-fileset#late", + artifact_storage_type=FileStorageType.FILESET, + workspace=DEFAULT_WORKSPACE, + ) + ) + await asyncio.sleep(0.05) + assert not result_create_task.done() + + allow_delete.set() + assert await delete_task is True + with pytest.raises(EntityNotFoundError): + await result_create_task + finally: + allow_delete.set() + if result_create_task is not None: + await asyncio.gather(delete_task, result_create_task, return_exceptions=True) + else: + await asyncio.gather(delete_task, return_exceptions=True) + + await verify_job_data_exists(mock_store, job_id, should_exist=False) + assert await count_entities(mock_store, PlatformJobResult, {"job": job_id}) == 0 + + +@pytest.mark.asyncio +async def test_delete_job_does_not_leave_idle_mutation_lock( + mock_dispatcher: JobDispatcher, + mock_store: EntityClient, +): + """Idle job mutation locks are weakly held and do not grow forever.""" + from nmp.core.jobs.app import dispatcher as dispatcher_module + + _, job_name, _, _, _, _ = await create_test_job_data(mock_store, "delete-lock-prune-test-job") + + deleted = await mock_dispatcher.delete_job(job_name, DEFAULT_WORKSPACE) + + assert deleted is True + gc.collect() + assert (DEFAULT_WORKSPACE, job_name) not in dispatcher_module._JOB_MUTATION_LOCKS + + @pytest.mark.asyncio async def test_create_or_update_task_ignores_invalid_terminal_transition( mock_dispatcher: JobDispatcher, mock_store: EntityClient @@ -393,6 +707,10 @@ def _echo_create(*, body, workspace): job = await mock_dispatcher.create_job(sample_platform_job_request, DEFAULT_WORKSPACE) assert job.fileset == f"job-fileset-{job.name}" + cancelled = await mock_dispatcher.cancel_job(job.name, DEFAULT_WORKSPACE) + assert cancelled is not None + assert cancelled.status == PlatformJobStatus.CANCELLED + deleted = await mock_dispatcher.delete_job(job.name, DEFAULT_WORKSPACE) assert deleted is True @@ -614,6 +932,20 @@ async def counting_update(entity, *args, **kwargs): assert step_update_count == 0 +@pytest.mark.asyncio +async def test_update_job_status_from_step_missing_job_raises_typed_skip( + mock_dispatcher: JobDispatcher, + mock_store: EntityClient, +): + """A step status update racing a job delete does not raise a generic Exception.""" + job_id, _, _, step_id, _, _ = await create_test_job_data(mock_store, "missing-job-status-update") + step = await mock_store.get_by_id(PlatformJobStep, step_id) + await mock_store.delete_by_id(PlatformJob, job_id) + + with pytest.raises(JobStatusUpdateSkippedError): + await mock_dispatcher.update_job_status_from_step(step, PlatformJobStatus.COMPLETED) + + @pytest.mark.asyncio async def test_cancel_job_multiple_steps_only_cancels_active( mock_dispatcher: JobDispatcher, diff --git a/services/core/jobs/tests/test_jobs_api.py b/services/core/jobs/tests/test_jobs_api.py index 4d93f8d2d9..68942b5440 100644 --- a/services/core/jobs/tests/test_jobs_api.py +++ b/services/core/jobs/tests/test_jobs_api.py @@ -1368,6 +1368,50 @@ async def test_cancel_job_conflict_sanitizes_log_fields( assert "\n" not in message +@pytest.mark.asyncio +async def test_delete_non_terminal_job_returns_409_and_keeps_job( + test_client: AsyncClient, + sample_platform_job_request: CreatePlatformJobRequest, + caplog: pytest.LogCaptureFixture, +): + request = sample_platform_job_request.model_copy(update={"name": "non-terminal-delete"}) + create_response = await test_client.post("/apis/jobs/v2/workspaces/default/jobs", json=request.model_dump()) + assert create_response.status_code == 201, create_response.text + + with caplog.at_level(logging.INFO, logger="nmp.core.jobs.api.v2.jobs.endpoints"): + delete_response = await test_client.delete("/apis/jobs/v2/workspaces/default/jobs/non-terminal-delete") + + assert delete_response.status_code == 409 + assert "Cancel the job and wait for it to reach a terminal state" in delete_response.json()["detail"] + log_record = next(record for record in caplog.records if "Cannot delete job" in record.getMessage()) + assert log_record.exc_info is None + + get_response = await test_client.get("/apis/jobs/v2/workspaces/default/jobs/non-terminal-delete") + assert get_response.status_code == 200 + assert get_response.json()["status"] == PlatformJobStatus.CREATED.value + + +@pytest.mark.asyncio +async def test_factory_delete_non_terminal_job_propagates_409(test_client: AsyncClient): + create_response = await test_client.post( + "/apis/jobs/v2/workspaces/default/hello-world/jobs", + json={ + "name": "factory-non-terminal-delete", + "description": "factory delete conflict", + "spec": {"config": {"key": "value"}, "target": "str"}, + "ownership": {"user": "u", "service": "s"}, + }, + ) + assert create_response.status_code == 201, create_response.text + + delete_response = await test_client.delete( + "/apis/jobs/v2/workspaces/default/hello-world/jobs/factory-non-terminal-delete" + ) + + assert delete_response.status_code == 409 + assert "Cancel the job and wait for it to reach a terminal state" in delete_response.json()["detail"] + + @pytest.mark.asyncio async def test_update_job_step_conflict_sanitizes_log_fields( test_client: AsyncClient, diff --git a/services/core/jobs/tests/test_jobs_client.py b/services/core/jobs/tests/test_jobs_client.py index 00661dfa64..3c366ce57b 100644 --- a/services/core/jobs/tests/test_jobs_client.py +++ b/services/core/jobs/tests/test_jobs_client.py @@ -155,6 +155,8 @@ async def test_job_lifecycle_methods_round_trip( assert cancelled.status == PlatformJobStatus.CANCELLED deleted_job = await _create_job(jobs_client, sample_platform_job_request, "typed-delete") + delete_ready = (await jobs_client.cancel_job(workspace="default", name=deleted_job.name)).data() + assert delete_ready.status == PlatformJobStatus.CANCELLED deleted = await jobs_client.delete_job(workspace="default", name=deleted_job.name) assert deleted.http_response.status_code == 204