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 7353358f08..cad3806984 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 @@ -7,11 +7,13 @@ import logging import os import tarfile +import threading import time import uuid from abc import abstractmethod from concurrent.futures import ThreadPoolExecutor -from typing import Generic, Literal, TypeVar +from dataclasses import dataclass +from typing import Any, Generic, Literal, TypeVar import docker.types from docker.errors import APIError, ImageNotFound, NotFound @@ -79,6 +81,7 @@ FailedToScheduleError, JobStorageError, ResourceAllocationError, + SchedulingDeferred, ) from opentelemetry import trace from pydantic import BaseModel, Field @@ -87,6 +90,7 @@ tracer = trace.get_tracer(__name__) logger = logging.getLogger(__name__) +DOCKER_CONTAINER_START_WORKERS = 10 def k8s_shm_quantity_to_docker(quantity: str) -> str: @@ -116,6 +120,13 @@ def k8s_shm_quantity_to_docker(quantity: str) -> str: ProviderT = TypeVar("ProviderT", bound=ExecutionProviderT) +@dataclass(frozen=True, slots=True) +class DockerTimestampParseResult: + parsed: datetime.datetime | None + parse_error: str | None + is_zero: bool + + class DockerVolumeMount(BaseModel): volume_name: str = Field(description="Name of the Docker volume to mount") mount_path: str = Field(description="Path inside the container where the volume will be mounted") @@ -181,7 +192,8 @@ class DockerJobBackend(JobBackend[ProviderT, DockerJobExecutionProfileConfig], G BACKEND_NAME: str = "docker" def init(self) -> None: - self._container_run_threadpool = ThreadPoolExecutor(max_workers=10) + self._container_start_admission = threading.BoundedSemaphore(DOCKER_CONTAINER_START_WORKERS) + self._container_run_threadpool = ThreadPoolExecutor(max_workers=DOCKER_CONTAINER_START_WORKERS) self._client = docker.from_env(timeout=180) if NEMO_JOBS_IMAGE_REGISTRY: logger.info( @@ -555,17 +567,6 @@ def schedule_single_container( executor_config: ProviderT, step: PlatformJobStepWithContext, ) -> JobUpdate: - ttl_seconds = self._execution_profile_config.ttl_seconds_before_active - if self.should_enforce_before_active_ttl(step) and self.check_step_ttl_before_active(step, ttl_seconds): - # If we are here it is because the scheduler kept retrying the job, - # which means there was a resource issue (GPU contention). Nothing really - # imperative to clean up, just mark as ERROR. - return JobUpdate( - status=PlatformJobStatus.ERROR.value, - status_details={"message": f"Job timed out after reaching max TTL of {ttl_seconds} seconds"}, - error_details={"message": f"Job timed out after reaching max TTL of {ttl_seconds} seconds"}, - ) - platform_config = get_platform_config() # Profile-level env vars first (e.g. HOME=/tmp); system, step, and shared env override these @@ -642,11 +643,62 @@ def schedule_single_container( }, ) + if not self._container_start_admission.acquire(blocking=False): + logger.debug( + "Docker start admission full, deferring scheduling", + extra={"job": step.job, "step": step.name}, + ) + raise SchedulingDeferred("Docker start worker capacity is full") + + # The admission slot is owned by this method until submit succeeds. + # After that, run_container releases it when the start worker exits. + try: + container_args = self._prepare_container_args_for_start( + executor_config=executor_config, + step=step, + task_id=task_id, + env=env, + log_config=log_config, + job_storage_mount=job_storage_mount, + task_storage_mount=task_storage_mount, + step_config_json=step_config_json, + ) + submitted_to_threadpool_at = time.monotonic() + self._container_run_threadpool.submit(self.run_container, step, container_args, submitted_to_threadpool_at) + except Exception: + self._container_start_admission.release() + raise + logger.debug( + "Docker run_container submitted", + extra={ + "job": step.job, + "step": step.name, + "task": task_id, + }, + ) + return JobUpdate( + status=PlatformJobStatus.PENDING, + status_details={"message": "Container schedule pending, checking for existing image and container"}, + ) + + def _prepare_container_args_for_start( + self, + *, + executor_config: ProviderT, + step: PlatformJobStepWithContext, + task_id: str, + env: dict, + log_config: LogConfig, + job_storage_mount: str, + task_storage_mount: str, + step_config_json: str, + ) -> dict: storage_config = self._execution_profile_config.storage job_volume_name = storage_config.volume_name if storage_config is not None else "" task_volume_name = self.task_storage_volume_name(workspace=step.workspace, job=step.job, task=task_id) config_volume_name = self.task_config_volume_name(workspace=step.workspace, job=step.job, task=task_id) additional_volume_mounts = storage_config.additional_volume_mounts if storage_config else None + ensure_storage_started_at = time.monotonic() self.ensure_job_storage( # if the job storage mount is not used, pass empty string to avoid creating unnecessary job storage volume job_storage_volume_name=job_volume_name if job_storage_mount != "" else "", @@ -659,6 +711,15 @@ def schedule_single_container( additional_volumes_mounts=additional_volume_mounts, step_config_json=step_config_json, ) + logger.debug( + "Docker job storage ensured", + extra={ + "job": step.job, + "step": step.name, + "task": task_id, + "duration_seconds": time.monotonic() - ensure_storage_started_at, + }, + ) labels = { JOB_WORKSPACE_ID_LABEL: step.workspace, @@ -708,12 +769,7 @@ def schedule_single_container( } container_args["network"] = self._execution_profile_config.networking.job_container_network - container_args = self.configure_container(container_args, executor_config) - self._container_run_threadpool.submit(self.run_container, step, container_args) - return JobUpdate( - status=PlatformJobStatus.PENDING, - status_details={"message": "Container schedule pending, checking for existing image and container"}, - ) + return self.configure_container(container_args, executor_config) def cancel_scheduling(self, step: PlatformJobStepWithContext) -> bool: """Check if the job step is cancelling or pausing, and update status accordingly.""" @@ -770,10 +826,19 @@ def get_jobs_launcher_binary(self) -> io.BytesIO | None: return jobs_launcher_stream return None - def run_container(self, step: PlatformJobStepWithContext, container_args: dict): + def run_container( + self, + step: PlatformJobStepWithContext, + container_args: dict, + submitted_to_threadpool_at: float | None = None, + ): with start_span_with_ctx( tracer, "jobs_controller/docker_backend/run_container", JobContext(id=step.job, step_name=step.name) ): + log_extra = {"job": step.job, "step": step.name} + if submitted_to_threadpool_at is not None: + log_extra["queue_delay_seconds"] = time.monotonic() - submitted_to_threadpool_at + logger.debug("Docker run_container worker started", extra=log_extra) try: self._run_container_in_thread(step, container_args) except FailedToScheduleError as e: @@ -788,6 +853,8 @@ def run_container(self, step: PlatformJobStepWithContext, container_args: dict): logger.exception("Failed to schedule container for job step") except Exception: logger.exception("Unexpected error while scheduling container for job step") + finally: + self._container_start_admission.release() def _run_container_in_thread(self, step: PlatformJobStepWithContext, container_args: dict): status_details = {} @@ -796,12 +863,39 @@ def _run_container_in_thread(self, step: PlatformJobStepWithContext, container_a # If a request to pause or cancel came in while we were waiting for scheduling loop, # cancel scheduling the container logger.debug("Checking for cancellation or pausing before creating container") + cancel_check_started_at = time.monotonic() if self.cancel_scheduling(step): + logger.debug( + "Docker pre-create cancellation check stopped scheduling", + extra={ + "job": step.job, + "step": step.name, + "duration_seconds": time.monotonic() - cancel_check_started_at, + }, + ) return + logger.debug( + "Docker pre-create cancellation check completed", + extra={ + "job": step.job, + "step": step.name, + "duration_seconds": time.monotonic() - cancel_check_started_at, + }, + ) # For resuming containers, check if it already exists logger.debug("Checking for existing container for job step") + get_container_started_at = time.monotonic() container = self.get_container(step) + logger.debug( + "Docker existing container lookup completed", + extra={ + "job": step.job, + "step": step.name, + "found": container is not None, + "duration_seconds": time.monotonic() - get_container_started_at, + }, + ) if container is not None: logger.info("Container already exists, not creating a new one", extra={"container_name": container.name}) else: @@ -809,7 +903,17 @@ def _run_container_in_thread(self, step: PlatformJobStepWithContext, container_a logger.debug("Creating container for job step") # Find the jobs launcher binary inside this running python container and also include it in the job container + launcher_lookup_started_at = time.monotonic() jobs_launcher_stream = self.get_jobs_launcher_binary() + logger.debug( + "Docker jobs launcher lookup completed", + extra={ + "job": step.job, + "step": step.name, + "found": jobs_launcher_stream is not None, + "duration_seconds": time.monotonic() - launcher_lookup_started_at, + }, + ) if jobs_launcher_stream is not None: # Modify the container entrypoint and command to use the jobs-launcher original_entrypoint = container_args.get("entrypoint", []) @@ -822,9 +926,24 @@ def _run_container_in_thread(self, step: PlatformJobStepWithContext, container_a ) try: + create_started_at = time.monotonic() container = self._client.containers.create(**container_args) + create_duration_seconds = time.monotonic() - create_started_at # Container will create successfully only if image is found locally logger.info("Image found locally", extra={"image": container_args["image"]}) + logger.debug( + "Docker container create succeeded", + extra={ + "job": step.job, + "step": step.name, + "container_name": container.name, + "image": container_args["image"], + "duration_seconds": create_duration_seconds, + "image_source": "local", + "requested_auto_remove": container_args.get("auto_remove"), + "host_config_auto_remove": container.attrs.get("HostConfig", {}).get("AutoRemove"), + }, + ) except (ImageNotFound, NotFound): # Image not found locally, pull it logger.info("Image not found locally, pulling from registry", extra={"image": container_args["image"]}) @@ -871,7 +990,21 @@ def _run_container_in_thread(self, step: PlatformJobStepWithContext, container_a # Now create it with the pulled container image logger.debug("Creating container for job step after pulling image") try: + create_started_at = time.monotonic() container = self._client.containers.create(**container_args) + logger.debug( + "Docker container create succeeded", + extra={ + "job": step.job, + "step": step.name, + "container_name": container.name, + "image": container_args["image"], + "duration_seconds": time.monotonic() - create_started_at, + "image_source": "pulled", + "requested_auto_remove": container_args.get("auto_remove"), + "host_config_auto_remove": container.attrs.get("HostConfig", {}).get("AutoRemove"), + }, + ) except APIError as e: raise FailedToScheduleError( "Failed to create container for job step", @@ -886,8 +1019,17 @@ def _run_container_in_thread(self, step: PlatformJobStepWithContext, container_a # Insert the jobs-launcher into the container if the launcher exists if jobs_launcher_stream is not None: try: + put_archive_started_at = time.monotonic() container.put_archive(path="/", data=jobs_launcher_stream) - logger.debug("Jobs launcher inserted into container successfully") + logger.debug( + "Jobs launcher inserted into container successfully", + extra={ + "job": step.job, + "step": step.name, + "container_name": container.name, + "duration_seconds": time.monotonic() - put_archive_started_at, + }, + ) except APIError as e: raise FailedToScheduleError( "Failed to insert jobs-launcher into container", @@ -900,12 +1042,32 @@ def _run_container_in_thread(self, step: PlatformJobStepWithContext, container_a # If a request to pause or cancel came in while we were waiting for scheduling loop, # cancel scheduling the container logger.debug("Checking for cancellation or pausing before starting container") + pre_start_cancel_check_started_at = time.monotonic() if self.cancel_scheduling(step): + logger.debug( + "Docker pre-start cancellation check stopped scheduling", + extra={ + "job": step.job, + "step": step.name, + "container_name": container.name, + "duration_seconds": time.monotonic() - pre_start_cancel_check_started_at, + }, + ) return + logger.debug( + "Docker pre-start cancellation check completed", + extra={ + "job": step.job, + "step": step.name, + "container_name": container.name, + "duration_seconds": time.monotonic() - pre_start_cancel_check_started_at, + }, + ) try: # If no errors to this point, start the container status_details["message"] = "Starting container" + pre_start_status_write_started_at = time.monotonic() self._nmp_sdk.jobs.steps.update_status( step.name, workspace=step.workspace, @@ -913,9 +1075,19 @@ def _run_container_in_thread(self, step: PlatformJobStepWithContext, container_a status=status, status_details=status_details, ) + logger.debug( + "Docker pre-start status update succeeded", + extra={ + "job": step.job, + "step": step.name, + "container_name": container.name, + "duration_seconds": time.monotonic() - pre_start_status_write_started_at, + }, + ) started = False max_attempts = 3 attempts = 0 + start_started_at = time.monotonic() while not started and attempts < max_attempts: attempts += 1 try: @@ -932,8 +1104,15 @@ def _run_container_in_thread(self, step: PlatformJobStepWithContext, container_a raise e time.sleep(5) # brief pause before retrying - logger.info( - "Started container for job step", extra={"container_name": container.name, "attempts": attempts} + logger.debug( + "Started container for job step", + extra={ + "job": step.job, + "step": step.name, + "container_name": container.name, + "attempts": attempts, + "duration_seconds": time.monotonic() - start_started_at, + }, ) except Exception as e: raise FailedToScheduleError( @@ -953,6 +1132,8 @@ def _sync(self, step: PlatformJobStepWithContext) -> JobUpdate: return self._kill_container_with_error(step, container, message) return self.sync_active(step, container) elif step.status == PlatformJobStatus.PENDING: + if container is not None and container.status in ("running", "exited", "dead"): + return self.sync_pending(step, container) if result := self.enforce_sync_ttl( step, self._execution_profile_config.ttl_seconds_before_active, @@ -1121,6 +1302,67 @@ def sync_stop_container(self, step: PlatformJobStepWithContext, container: Conta else: raise e + @staticmethod + def parse_docker_timestamp(timestamp: str | None) -> DockerTimestampParseResult: + """Parse a timestamp from Docker container state. + + Docker stores lifecycle timestamps as Go time.Time values and exposes them + through inspect as formatted strings. An unset Go time.Time formats as + 0001-01-01T00:00:00Z, so treat that value as absent while keeping a + structured flag for debugging. + """ + if not timestamp: + return DockerTimestampParseResult(parsed=None, parse_error=None, is_zero=False) + is_zero_time = timestamp.startswith("0001-01-01") + if is_zero_time: + return DockerTimestampParseResult(parsed=None, parse_error=None, is_zero=True) + try: + parsed = datetime.datetime.fromisoformat(timestamp) + except ValueError as exc: + return DockerTimestampParseResult(parsed=None, parse_error=str(exc), is_zero=False) + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=datetime.UTC) + return DockerTimestampParseResult(parsed=parsed, parse_error=None, is_zero=False) + + def docker_state_debug_fields(self, container: Container) -> dict[str, Any]: + attrs = container.attrs or {} + state = attrs.get("State", {}) + now = datetime.datetime.now(datetime.UTC) + finished_at_raw = state.get("FinishedAt") + finished_at_result = self.parse_docker_timestamp(finished_at_raw) + finished_at = finished_at_result.parsed + cleanup_after_finished_at = ( + finished_at + datetime.timedelta(seconds=self._execution_profile_config.ttl_seconds_after_finished) + if finished_at + else None + ) + + return { + "container_status": container.status, + "docker_state_status": state.get("Status"), + "docker_state_started_at": state.get("StartedAt"), + "docker_state_finished_at": finished_at_raw, + "docker_state_finished_at_parsed": finished_at.isoformat() if finished_at else None, + "docker_state_finished_at_parse_error": finished_at_result.parse_error, + "docker_state_finished_at_is_zero": finished_at_result.is_zero, + "docker_state_finished_age_seconds": (now - finished_at).total_seconds() if finished_at else None, + "docker_state_exit_code": state.get("ExitCode"), + "docker_state_error": state.get("Error"), + "docker_state_oom_killed": state.get("OOMKilled"), + "docker_state_dead": state.get("Dead"), + "docker_state_running": state.get("Running"), + "docker_state_paused": state.get("Paused"), + "host_config_auto_remove": attrs.get("HostConfig", {}).get("AutoRemove"), + "cleanup_completed_jobs_immediately": self._execution_profile_config.cleanup_completed_jobs_immediately, + "ttl_seconds_after_finished": self._execution_profile_config.ttl_seconds_after_finished, + "cleanup_after_finished_at": cleanup_after_finished_at.isoformat() if cleanup_after_finished_at else None, + "cleanup_ttl_remaining_seconds": (cleanup_after_finished_at - now).total_seconds() + if cleanup_after_finished_at + else None, + "cleanup_ttl_due": cleanup_after_finished_at <= now if cleanup_after_finished_at else None, + "now_utc": now.isoformat(), + } + def create_step_update(self, step: PlatformJobStepWithContext, container: Container) -> JobUpdate: status, status_details, error_stack = self.map_docker_container_status_to_platform_status(step, container) task_id = self.get_label_from_container(container, JOB_TASK_ID_LABEL) @@ -1128,6 +1370,19 @@ def create_step_update(self, step: PlatformJobStepWithContext, container: Contai if status == PlatformJobStatus.ERROR: error_details["message"] = status_details.get("message", "Job encountered an error") + logger.debug( + "Docker container status mapped to platform status", + extra={ + "workspace": step.workspace, + "job": step.job, + "step": step.name, + "task": task_id, + "container_name": container.name, + "platform_status": status.value, + **self.docker_state_debug_fields(container), + }, + ) + # Upsert the task against the Jobs API. self._nmp_sdk.jobs.tasks.create_or_update( task_id, @@ -1258,7 +1513,9 @@ def cleanup_steps(self): if container.labels.get(JOB_TYPE_LABEL) != JOB_TYPE_JOB: continue if container.status in ("exited", "dead"): - exit_code = container.attrs.get("State", {}).get("ExitCode", 0) + state_debug = self.docker_state_debug_fields(container) + exit_code = state_debug.get("docker_state_exit_code") or 0 + auto_remove = state_debug.get("host_config_auto_remove") job = self.get_label_from_container(container, JOB_ID_LABEL) step_name = self.get_label_from_container(container, JOB_STEP_NAME_LABEL) workspace = self.get_label_from_container(container, JOB_WORKSPACE_ID_LABEL) @@ -1266,8 +1523,29 @@ def cleanup_steps(self): # Verify the step is terminal before cleaning up. # This prevents cleaning up resources that we last marked in active state, # were prematurely cleaned up, and then sync active to error because the resource is gone. - if not self.check_step_is_terminal(job=job, step_name=step_name, workspace=workspace): - logger.debug("Skipping cleanup for for job container because step is not in terminal state") + step_is_terminal = self.check_step_is_terminal(job=job, step_name=step_name, workspace=workspace) + cleanup_log_extra = { + "workspace": workspace, + "job": job, + "step": step_name, + "container_name": container.name, + "container_id": container.id[:16], + "exit_code": exit_code, + "host_config_auto_remove": auto_remove, + } + logger.debug( + "Docker cleanup inspected exited job container", + extra={ + **cleanup_log_extra, + "step_is_terminal": step_is_terminal, + **state_debug, + }, + ) + if not step_is_terminal: + logger.debug( + "Skipping cleanup for job container because step is not in terminal state", + extra=cleanup_log_extra, + ) continue # Always disconnect the container from its network first if not already done. @@ -1279,16 +1557,45 @@ def cleanup_steps(self): self._execution_profile_config.cleanup_completed_jobs_immediately and exit_code in TERMINAL_EXIT_CODES ): + logger.debug( + "Docker cleanup removing terminal job container immediately", + extra=cleanup_log_extra, + ) self.cleanup_single_container(container) continue # Otherwise, check if the TTL has expired for errored jobs or completed jobs if not cleaned up immediately last_transition_time_str = container.attrs.get("State", {}).get("FinishedAt") - if last_transition_time_str and ( - datetime.datetime.fromisoformat(last_transition_time_str) + finished_at_result = self.parse_docker_timestamp(last_transition_time_str) + cleanup_after_finished_at = ( + finished_at_result.parsed + datetime.timedelta(seconds=self._execution_profile_config.ttl_seconds_after_finished) - ) < datetime.datetime.now(datetime.UTC): + if finished_at_result.parsed + else None + ) + if cleanup_after_finished_at and cleanup_after_finished_at < datetime.datetime.now(datetime.UTC): + logger.debug( + "Docker cleanup removing expired job container", + extra={ + **cleanup_log_extra, + "finished_at": last_transition_time_str, + "finished_at_parse_error": finished_at_result.parse_error, + "finished_at_is_zero": finished_at_result.is_zero, + **state_debug, + }, + ) self.cleanup_single_container(container) + else: + logger.debug( + "Docker cleanup retaining terminal job container until TTL", + extra={ + **cleanup_log_extra, + "finished_at": last_transition_time_str, + "finished_at_parse_error": finished_at_result.parse_error, + "finished_at_is_zero": finished_at_result.is_zero, + **state_debug, + }, + ) except NotFound: # Container may disappear between list and inspect/attribute access. # Ignore and continue cleanup for remaining containers. diff --git a/services/core/jobs/src/nmp/core/jobs/controllers/backends/exceptions.py b/services/core/jobs/src/nmp/core/jobs/controllers/backends/exceptions.py index 0f3eb2dcd6..2b833ec05f 100644 --- a/services/core/jobs/src/nmp/core/jobs/controllers/backends/exceptions.py +++ b/services/core/jobs/src/nmp/core/jobs/controllers/backends/exceptions.py @@ -18,5 +18,13 @@ def __init__(self, message: str = "Failed to allocate resource"): self.message = message +class SchedulingDeferred(Exception): + """Exception raised when a backend has no immediate capacity to accept a step.""" + + def __init__(self, message: str = "Scheduling deferred"): + super().__init__(message) + self.message = message + + class JobStorageError(Exception): """Exception raised when there's an issue with job storage.""" diff --git a/services/core/jobs/src/nmp/core/jobs/controllers/backends/subprocess.py b/services/core/jobs/src/nmp/core/jobs/controllers/backends/subprocess.py index 72439f7352..6944e0cf21 100644 --- a/services/core/jobs/src/nmp/core/jobs/controllers/backends/subprocess.py +++ b/services/core/jobs/src/nmp/core/jobs/controllers/backends/subprocess.py @@ -161,15 +161,6 @@ def shutdown(self) -> None: self._finish_logs(metadata) def schedule(self, executor_config: SubprocessExecutionProvider, step: PlatformJobStepWithContext) -> JobUpdate: - ttl_seconds = self._execution_profile_config.ttl_seconds_before_active - if self.should_enforce_before_active_ttl(step) and self.check_step_ttl_before_active(step, ttl_seconds): - message = f"Job timed out after reaching max TTL of {ttl_seconds} seconds" - return JobUpdate( - status=PlatformJobStatus.ERROR.value, - status_details={"message": message}, - error_details={"message": message}, - ) - if not executor_config.command: return JobUpdate( status=PlatformJobStatus.ERROR.value, diff --git a/services/core/jobs/src/nmp/core/jobs/controllers/reconciler.py b/services/core/jobs/src/nmp/core/jobs/controllers/reconciler.py index 8dab47f099..423370132e 100644 --- a/services/core/jobs/src/nmp/core/jobs/controllers/reconciler.py +++ b/services/core/jobs/src/nmp/core/jobs/controllers/reconciler.py @@ -3,6 +3,7 @@ import logging import threading +import time from nemo_platform import APIError, APIStatusError, NeMoPlatform from nemo_platform.types.jobs import PlatformJobStepWithContext @@ -54,6 +55,7 @@ def step(self): logger.debug("Stop signal received, skipping reconciliation step") return + fetch_started_at = time.monotonic() with tracer.start_as_current_span("jobs_reconciler/fetch_steps_for_reconciliation"): try: statuses: list[SDKPlatformJobStatus] = [ @@ -73,6 +75,10 @@ def step(self): logger.info(f"Got {len(steps_to_reconcile)} job steps to reconcile") else: logger.debug("No job steps to reconcile") + logger.debug( + "Reconciler fetched job steps", + extra={"count": len(steps_to_reconcile), "duration_seconds": time.monotonic() - fetch_started_at}, + ) for step in steps_to_reconcile: with start_span_with_ctx( tracer, @@ -86,7 +92,21 @@ def step(self): with scoped_app_ctx( JobBackendContext(provider=provider, profile=profile, name=str(backend)), ): + sync_started_at = time.monotonic() job_update = backend.sync(step) + logger.debug( + "Reconciler backend sync completed", + extra={ + "job": step.job, + "step": step.name, + "workspace": step.workspace, + "provider": provider, + "profile": profile, + "from_status": step.status, + "to_status": job_update.status, + "duration_seconds": time.monotonic() - sync_started_at, + }, + ) logger.info(f"Updating job step status from '{step.status}' to '{job_update.status}'") if ( job_update.status == PlatformJobStatus.ERROR.value @@ -98,13 +118,13 @@ def step(self): logger=self._logger, context="step transitioned to error during reconciliation", ) - self._nmp_sdk.jobs.steps.update_status( - step.name, - workspace=step.workspace, - job=step.job, + self._update_step_status_with_timing( + step=step, + provider=provider, + profile=profile, status=job_update.status, - status_details=job_update.status_details, # type: ignore - error_details=job_update.error_details, # type: ignore + status_details=job_update.status_details, + error_details=job_update.error_details, ) except APIStatusError as e: # In cases when attempting to update job step status results in a conflict (409), @@ -149,6 +169,56 @@ def step(self): except Exception: logger.exception("Could not complete cleanup steps for backend", exc_info=True) + def _update_step_status_with_timing( + self, + *, + step: PlatformJobStepWithContext, + provider: str, + profile: str, + status: str, + status_details: dict | None = None, + error_details: dict | None = None, + ): + started_at = time.monotonic() + try: + response = self._nmp_sdk.jobs.steps.update_status( + step.name, + workspace=step.workspace, + job=step.job, + status=status, + status_details=status_details, # type: ignore + error_details=error_details, # type: ignore + ) + except Exception: + logger.warning( + "Reconciler step status update failed", + extra={ + "job": step.job, + "step": step.name, + "workspace": step.workspace, + "provider": provider, + "profile": profile, + "from_status": step.status, + "to_status": status, + "duration_seconds": time.monotonic() - started_at, + }, + ) + raise + logger.debug( + "Reconciler step status update succeeded", + extra={ + "job": step.job, + "step": step.name, + "workspace": step.workspace, + "provider": provider, + "profile": profile, + "from_status": step.status, + "to_status": status, + "duration_seconds": time.monotonic() - started_at, + }, + ) + return response + def get_steps_for_reconciliation(self, statuses: list[SDKPlatformJobStatus]) -> list[PlatformJobStepWithContext]: """ Return the list of steps to reconcile. 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 4d9b5a8e98..e2f2f2a495 100644 --- a/services/core/jobs/src/nmp/core/jobs/controllers/scheduler.py +++ b/services/core/jobs/src/nmp/core/jobs/controllers/scheduler.py @@ -3,10 +3,13 @@ import logging import threading +import time import traceback +from typing import TypedDict, cast, get_args import nemo_platform from nemo_platform import APIStatusError, NeMoPlatform +from nemo_platform.types import PlatformJobStatus as SDKPlatformJobStatus from nemo_platform.types.jobs import PlatformJobStepWithContext from nemo_platform.types.jobs.platform_job_steps_list_filter_param import PlatformJobStepsListFilterParam from nmp.common.controller import Controller @@ -14,7 +17,7 @@ from nmp.common.observability import start_span_with_ctx from nmp.core.jobs.app.ctx import JobBackendContext, JobContext from nmp.core.jobs.controllers.backends import JobUpdate, extract_provider_profile -from nmp.core.jobs.controllers.backends.exceptions import ResourceAllocationError +from nmp.core.jobs.controllers.backends.exceptions import ResourceAllocationError, SchedulingDeferred from nmp.core.jobs.controllers.backends.registry import BackendRegistry from nmp.core.jobs.controllers.diagnostics import log_job_diagnostics_if_debug from opentelemetry import metrics, trace @@ -25,6 +28,18 @@ DEFAULT_PROFILE = "default" DEFAULT_PROVIDER = "cpu" +SDK_PLATFORM_JOB_STATUSES = frozenset(get_args(SDKPlatformJobStatus)) + + +class StepStatusDetailParams(TypedDict, total=False): + status_details: dict[str, object] + error_details: dict[str, object] + + +def as_sdk_platform_job_status(status: str) -> SDKPlatformJobStatus: + if status not in SDK_PLATFORM_JOB_STATUSES: + raise ValueError(f"Unsupported platform job status: {status}") + return cast(SDKPlatformJobStatus, status) class JobScheduler(Controller): @@ -60,6 +75,7 @@ def step(self): return steps = [] + fetch_started_at = time.monotonic() with tracer.start_as_current_span("jobs_scheduler/fetch_steps_for_scheduling"): try: steps = self.get_steps_for_scheduling() @@ -73,21 +89,34 @@ def step(self): logger.info(f"Got {len(steps)} job steps to schedule") else: logger.debug("No job steps to schedule") + logger.debug( + "Scheduler fetched job steps", + extra={"count": len(steps), "duration_seconds": time.monotonic() - fetch_started_at}, + ) for step in steps: with start_span_with_ctx( tracer, "jobs_scheduler/schedule_step", JobContext(id=step.job, step_name=step.name) ): try: + schedule_started_at = time.monotonic() update = self.schedule_step(step) - logger.info("Scheduled job step") + logger.debug( + "Scheduled job step", + extra={ + "job": step.job, + "step": step.name, + "workspace": step.workspace, + "duration_seconds": time.monotonic() - schedule_started_at, + "status": update.status, + }, + ) try: - self._nmp_sdk.jobs.steps.update_status( - step.name, - workspace=step.workspace, - job=step.job, + self._update_step_status_with_timing( + step=step, + phase="schedule", status=update.status, - status_details=update.status_details, # type: ignore - error_details=update.error_details, # type: ignore + status_details=update.status_details, + error_details=update.error_details, ) except APIStatusError as e: # Stopgap for a scheduler/reconciler race: by the time the scheduler persists @@ -119,14 +148,23 @@ def step(self): context="resource allocation error during scheduling", ) self._step_scheduling_errors.add(1, attributes={"error_type": "resource_allocation"}) - self._nmp_sdk.jobs.steps.update_status( - step.name, - workspace=step.workspace, - job=step.job, + self._update_step_status_with_timing( + step=step, + phase="resource_allocation_error", status=PlatformJobStatus.ERROR.value, status_details={"message": e.message}, error_details={"message": e.message}, ) + except SchedulingDeferred as e: + logger.debug( + "Scheduling deferred for job step", + extra={ + "job": step.job, + "step": step.name, + "workspace": step.workspace, + "reason": e.message, + }, + ) except Exception as e: logger.exception("Could not schedule job step", exc_info=True) log_job_diagnostics_if_debug( @@ -136,15 +174,65 @@ def step(self): context="unexpected scheduling error", ) self._step_scheduling_errors.add(1, attributes={"error_type": "unknown"}) - self._nmp_sdk.jobs.steps.update_status( - step.name, - workspace=step.workspace, - job=step.job, + self._update_step_status_with_timing( + step=step, + phase="unexpected_error", status=PlatformJobStatus.ERROR.value, status_details={"message": str(e)}, error_details={"message": str(e), "error": traceback.format_exc()}, ) + def _update_step_status_with_timing( + self, + *, + step: PlatformJobStepWithContext, + phase: str, + status: str, + status_details: dict[str, object] | None = None, + error_details: dict[str, object] | None = None, + ): + started_at = time.monotonic() + detail_params: StepStatusDetailParams = {} + if status_details is not None: + detail_params["status_details"] = status_details + if error_details is not None: + detail_params["error_details"] = error_details + try: + response = self._nmp_sdk.jobs.steps.update_status( + step.name, + workspace=step.workspace, + job=step.job, + status=as_sdk_platform_job_status(status), + **detail_params, + ) + except Exception: + logger.warning( + "Scheduler step status update failed", + extra={ + "job": step.job, + "step": step.name, + "workspace": step.workspace, + "phase": phase, + "from_status": step.status, + "to_status": status, + "duration_seconds": time.monotonic() - started_at, + }, + ) + raise + logger.debug( + "Scheduler step status update succeeded", + extra={ + "job": step.job, + "step": step.name, + "workspace": step.workspace, + "phase": phase, + "from_status": step.status, + "to_status": status, + "duration_seconds": time.monotonic() - started_at, + }, + ) + return response + def get_steps_for_scheduling(self) -> list[PlatformJobStepWithContext]: """ Return the oldest set of steps to schedule. We using the @@ -159,7 +247,7 @@ def get_steps_for_scheduling(self) -> list[PlatformJobStepWithContext]: name="-", # Use "-" to indicate all jobs workspace="-", # Cross-workspace query filter=filter_params, - sort="-created_at", + sort="created_at", ): steps.append(step) return steps diff --git a/services/core/jobs/tests/controllers/test_docker_backend.py b/services/core/jobs/tests/controllers/test_docker_backend.py index 0a9e1e4b08..2f916b82b0 100644 --- a/services/core/jobs/tests/controllers/test_docker_backend.py +++ b/services/core/jobs/tests/controllers/test_docker_backend.py @@ -49,13 +49,14 @@ ) from nmp.core.jobs.controllers.backends.docker import ( DEFAULT_VOLUME_PERMISSIONS_IMAGE, + DOCKER_CONTAINER_START_WORKERS, CPUDockerJobBackend, DockerJobExecutionProfileConfig, DockerJobStorageConfig, DockerVolumeMount, GPUDockerJobBackend, ) -from nmp.core.jobs.controllers.backends.exceptions import ResourceAllocationError +from nmp.core.jobs.controllers.backends.exceptions import ResourceAllocationError, SchedulingDeferred from pydantic import ValidationError @@ -1538,32 +1539,42 @@ def test_cleanup_steps_by_ttl(docker_job, docker_client_mock, test_job_step, cle mock_network.disconnect.assert_any_call(mock_container_old_error) -def test_cleanup_created_by_ttl(docker_job, docker_client_mock, test_job_step): - """Test that schedule_single_container transitions to an ERROR state when step's created_at exceeds TTL.""" - # Get the TTL configuration (default is 30 minutes) +def test_created_step_does_not_ttl_before_backend_acceptance(docker_job, docker_client_mock, test_job_step): + """CREATED age should not fail a step before the backend accepts it.""" ttl_seconds = docker_job._execution_profile_config.ttl_seconds_before_active - - # Create a step with an created_at timestamp that exceeds the TTL (35 minutes ago) old_timestamp = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(seconds=ttl_seconds + 300) test_job_step.created_at = old_timestamp test_job_step.updated_at = old_timestamp - test_job_step.status = PlatformJobStatus.CREATED - - # Get the executor config from the step executor_config = test_job_step.step_spec.executor + docker_job._container_run_threadpool = MagicMock() - # Call schedule_single_container which should detect the timeout - result = docker_job.schedule_single_container(executor_config, test_job_step) + try: + result = docker_job.schedule_single_container(executor_config, test_job_step) + finally: + if docker_job._container_run_threadpool.submit.called: + docker_job._container_start_admission.release() - # Verify that it returns an ERROR status with timeout message - assert result.status == PlatformJobStatus.ERROR.value - assert result.status_details == {"message": "Job timed out after reaching max TTL of 1800 seconds"} - assert result.error_details == {"message": "Job timed out after reaching max TTL of 1800 seconds"} + assert result.status == PlatformJobStatus.PENDING + docker_job._container_run_threadpool.submit.assert_called_once() + + +def test_docker_schedule_defers_when_start_admission_full(docker_job, docker_client_mock, test_job_step): + """A full Docker start gate leaves the step CREATED and avoids per-attempt Docker setup.""" + acquired = 0 + try: + for _ in range(DOCKER_CONTAINER_START_WORKERS): + assert docker_job._container_start_admission.acquire(blocking=False) + acquired += 1 - # Verify that no container was created - docker_client_mock.containers.create.assert_not_called() - docker_client_mock.containers.run.assert_not_called() + with pytest.raises(SchedulingDeferred, match="Docker start worker capacity is full"): + docker_job.schedule_single_container(test_job_step.step_spec.executor, test_job_step) + + docker_client_mock.containers.create.assert_not_called() + docker_client_mock.containers.run.assert_not_called() + finally: + for _ in range(acquired): + docker_job._container_start_admission.release() def test_resuming_step_skips_before_active_ttl_enforcement(docker_job, test_job_step): @@ -1586,8 +1597,8 @@ def test_before_active_ttl_uses_latest_of_created_and_updated(docker_job, test_j assert docker_job.check_step_ttl_before_active(test_job_step, ttl_seconds) is False -def test_cleanup_pending_by_ttl(docker_job, docker_client_mock, test_job_step): - """Test that sync of a PENDING step transitions to an ERROR state when step's created_at exceeds TTL.""" +def test_cleanup_pending_created_container_by_ttl(docker_job, docker_client_mock, test_job_step): + """A stale PENDING step with a Docker-created container transitions to ERROR.""" # Get the TTL configuration (default is 30 minutes) ttl_seconds = docker_job._execution_profile_config.ttl_seconds_before_active @@ -1602,7 +1613,7 @@ def test_cleanup_pending_by_ttl(docker_job, docker_client_mock, test_job_step): # Create a mock container that the sync method will find (with managed-by label so we kill it) container_mock = MagicMock() container_mock.id = "16-character-uid" - container_mock.status = "running" + container_mock.status = "created" task_id = uuid.uuid4().hex container_mock.labels = { JOB_ID_LABEL: test_job_step.job, @@ -1638,6 +1649,66 @@ def test_cleanup_pending_by_ttl(docker_job, docker_client_mock, test_job_step): ) +def test_pending_running_container_preempts_before_active_ttl(docker_job, docker_client_mock, test_job_step): + """If Docker is running, reconcile PENDING to ACTIVE instead of timing out first.""" + ttl_seconds = docker_job._execution_profile_config.ttl_seconds_before_active + old_timestamp = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(seconds=ttl_seconds + 300) + test_job_step.status = PlatformJobStatus.PENDING + test_job_step.created_at = old_timestamp + test_job_step.updated_at = old_timestamp + + task_id = uuid.uuid4().hex + container_mock = MagicMock() + container_mock.id = "16-character-uid" + container_mock.name = "job-test-job-id-test-step" + container_mock.status = "running" + container_mock.labels = { + JOB_ID_LABEL: test_job_step.job, + JOB_STEP_NAME_LABEL: test_job_step.name, + JOB_TASK_ID_LABEL: task_id, + JOB_MANAGED_BY_LABEL: JOB_MANAGED_BY_JOBS_CONTROLLER, + } + container_mock.attrs = {"State": {"Status": "running", "Running": True}, "HostConfig": {}} + docker_client_mock.containers.get.side_effect = None + docker_client_mock.containers.get.return_value = container_mock + + result = docker_job.sync(test_job_step) + + assert result.status == PlatformJobStatus.ACTIVE.value + assert result.status_details == {"message": "Job is running"} + container_mock.kill.assert_not_called() + + +def test_pending_exited_container_preempts_before_active_ttl(docker_job, docker_client_mock, test_job_step): + """If Docker already exited successfully, reconcile PENDING to COMPLETED instead of timing out first.""" + ttl_seconds = docker_job._execution_profile_config.ttl_seconds_before_active + old_timestamp = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(seconds=ttl_seconds + 300) + test_job_step.status = PlatformJobStatus.PENDING + test_job_step.created_at = old_timestamp + test_job_step.updated_at = old_timestamp + + task_id = uuid.uuid4().hex + container_mock = MagicMock() + container_mock.id = "16-character-uid" + container_mock.name = "job-test-job-id-test-step" + container_mock.status = "exited" + container_mock.labels = { + JOB_ID_LABEL: test_job_step.job, + JOB_STEP_NAME_LABEL: test_job_step.name, + JOB_TASK_ID_LABEL: task_id, + JOB_MANAGED_BY_LABEL: JOB_MANAGED_BY_JOBS_CONTROLLER, + } + container_mock.attrs = {"State": {"Status": "exited", "ExitCode": 0}, "HostConfig": {}} + docker_client_mock.containers.get.side_effect = None + docker_client_mock.containers.get.return_value = container_mock + + result = docker_job.sync(test_job_step) + + assert result.status == PlatformJobStatus.COMPLETED.value + assert result.status_details == {"message": "Job completed successfully with exit code 0"} + container_mock.kill.assert_not_called() + + def test_cleanup_active_by_ttl(docker_job, docker_client_mock, test_job_step): """Test that sync of an ACTIVE step transitions to an ERROR state when step's created_at exceeds TTL.""" # Get the TTL configuration for active jobs (default is 24 hours) @@ -1689,8 +1760,10 @@ def test_cleanup_active_by_ttl(docker_job, docker_client_mock, test_job_step): ) -def test_ttl_enforcement_container_already_stopped_on_kill(docker_job, docker_client_mock, test_job_step): - """Test TTL enforcement handles gracefully when container.kill() is called on already stopped container.""" +def test_ttl_enforcement_handles_409_when_kill_races_with_stopped_container( + docker_job, docker_client_mock, test_job_step +): + """Test TTL enforcement handles gracefully when container.kill() races with a stopped container.""" # Get the TTL configuration for pending/created jobs ttl_seconds = docker_job._execution_profile_config.ttl_seconds_before_active @@ -1706,7 +1779,7 @@ def test_ttl_enforcement_container_already_stopped_on_kill(docker_job, docker_cl container_mock = MagicMock() container_mock.id = "16-character-uid" container_mock.name = "job-test-job-id-test-step" - container_mock.status = "exited" # Container is already stopped + container_mock.status = "created" task_id = uuid.uuid4().hex container_mock.labels = { JOB_WORKSPACE_ID_LABEL: test_job_step.workspace, diff --git a/services/core/jobs/tests/controllers/test_kubernetes_backend.py b/services/core/jobs/tests/controllers/test_kubernetes_backend.py index c0c1f75391..9fa58f34f0 100644 --- a/services/core/jobs/tests/controllers/test_kubernetes_backend.py +++ b/services/core/jobs/tests/controllers/test_kubernetes_backend.py @@ -700,6 +700,21 @@ def test_schedule_job_success(kubernetes_job, cpu_execution_provider, test_step_ assert env_vars["NMP_CONFIG_WARNINGS_DISABLED"] == "1" +def test_created_step_does_not_ttl_before_backend_acceptance(kubernetes_job, cpu_execution_provider, test_step_pending): + """CREATED age should not fail a step before the Kubernetes backend accepts it.""" + kubernetes_job._batch_v1.create_namespaced_job.return_value = MagicMock() + ttl_seconds = kubernetes_job._execution_profile_config.ttl_seconds_before_active + old_timestamp = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(seconds=ttl_seconds + 300) + test_step_pending.created_at = old_timestamp + test_step_pending.updated_at = old_timestamp + test_step_pending.status = PlatformJobStatus.CREATED + + update = kubernetes_job.schedule(cpu_execution_provider, test_step_pending) + + assert update.status == PlatformJobStatus.PENDING + kubernetes_job._batch_v1.create_namespaced_job.assert_called_once() + + def test_kubernetes_job_profile_environment_applied( mock_nmp_client, kubernetes_client_mock, diff --git a/services/core/jobs/tests/controllers/test_scheduler.py b/services/core/jobs/tests/controllers/test_scheduler.py index 5da34188e0..889f1bb420 100644 --- a/services/core/jobs/tests/controllers/test_scheduler.py +++ b/services/core/jobs/tests/controllers/test_scheduler.py @@ -7,7 +7,7 @@ from nemo_platform import ConflictError from nmp.common.jobs.schemas import PlatformJobStatus from nmp.core.jobs.api.v2.jobs.schemas import PlatformJobStepWithContext -from nmp.core.jobs.controllers.backends.exceptions import ResourceAllocationError +from nmp.core.jobs.controllers.backends.exceptions import ResourceAllocationError, SchedulingDeferred from nmp.core.jobs.controllers.backends.registry import BackendRegistry from nmp.core.jobs.controllers.backends.test import MockDockerCPUJobBackend from nmp.core.jobs.controllers.scheduler import JobScheduler @@ -42,7 +42,7 @@ def test_does_schedule_job( workspace="-", name="-", filter={"status": ["created", "resuming"]}, - sort="-created_at", + sort="created_at", ) # Test backend should have received one schedule call for our test job @@ -51,6 +51,19 @@ def test_does_schedule_job( assert test_backend.mock.sync_calls == [] +def test_scheduling_deferred_leaves_step_created( + job_scheduler: JobScheduler, + mock_nmp_client, + test_step_pending: PlatformJobStepWithContext, +): + mock_nmp_client.jobs.steps.list.return_value = [test_step_pending] + + with patch.object(job_scheduler, "schedule_step", side_effect=SchedulingDeferred("capacity full")): + job_scheduler.step() + + mock_nmp_client.jobs.steps.update_status.assert_not_called() + + def test_resource_allocation_error_marks_step_as_error( job_scheduler: JobScheduler, mock_nmp_client, @@ -132,8 +145,6 @@ def test_scheduler_does_not_mark_step_error_when_pending_update_conflicts_with_c workspace=test_step_pending.workspace, job=test_step_pending.job, status=PlatformJobStatus.PENDING, - status_details=None, - error_details=None, ) mock_nmp_client.jobs.steps.retrieve.assert_called_once_with( test_step_pending.name, diff --git a/services/core/jobs/tests/controllers/test_subprocess_backend.py b/services/core/jobs/tests/controllers/test_subprocess_backend.py index 97e2265a21..db7f945b84 100644 --- a/services/core/jobs/tests/controllers/test_subprocess_backend.py +++ b/services/core/jobs/tests/controllers/test_subprocess_backend.py @@ -4,7 +4,7 @@ import os import sys import time -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from types import SimpleNamespace from unittest.mock import patch @@ -73,6 +73,26 @@ def test_schedule_starts_process_and_stages_environment( mock_nmp_client.jobs.tasks.create_or_update.assert_called() +def test_created_step_does_not_ttl_before_backend_acceptance( + mock_nmp_client, tmp_path, mock_platform_config, test_step_pending +): + backend = _subprocess_backend(mock_nmp_client, tmp_path, mock_platform_config) + step = _step_with_command(test_step_pending, ["/bin/sh", "-c", "true"]) + ttl_seconds = backend._execution_profile_config.ttl_seconds_before_active + old_timestamp = datetime.now(timezone.utc) - timedelta(seconds=ttl_seconds + 300) + step.created_at = old_timestamp + step.updated_at = old_timestamp + step.status = PlatformJobStatus.CREATED + + update = _schedule_without_otel_export(backend, step) + + assert update.status == PlatformJobStatus.PENDING + key = SubprocessProcessKey(step.workspace, step.job, str(step.attempt_id), step.name) + metadata = backend._process_registry.get(key) + assert metadata is not None + assert metadata.process.wait(timeout=5) == 0 + + def test_subprocess_persistent_storage_is_shared_across_job_attempt( mock_nmp_client, tmp_path, mock_platform_config, test_step_pending ): diff --git a/services/core/jobs/tests/controllers/test_volcano_backend.py b/services/core/jobs/tests/controllers/test_volcano_backend.py index fcaf352ec2..2d1b6af2c2 100644 --- a/services/core/jobs/tests/controllers/test_volcano_backend.py +++ b/services/core/jobs/tests/controllers/test_volcano_backend.py @@ -277,6 +277,25 @@ def test_schedule_job_success( assert env_vars["NMP_CONFIG_WARNINGS_DISABLED"] == "1" +def test_created_step_does_not_ttl_before_backend_acceptance( + volcano_job: VolcanoJobBackend, + distributed_gpu_execution_provider, + test_step_pending: PlatformJobStepWithContext, +): + """CREATED age should not fail a step before the Volcano backend accepts it.""" + volcano_job._custom_v1.create_namespaced_custom_object.return_value = MagicMock() # ty: ignore[invalid-assignment] + ttl_seconds = volcano_job._execution_profile_config.ttl_seconds_before_active + old_timestamp = datetime.datetime.now(datetime.timezone.utc) - datetime.timedelta(seconds=ttl_seconds + 300) + test_step_pending.created_at = old_timestamp + test_step_pending.updated_at = old_timestamp + test_step_pending.status = PlatformJobStatus.CREATED + + update = volcano_job.schedule(distributed_gpu_execution_provider, test_step_pending) + + assert update.status == PlatformJobStatus.PENDING + volcano_job._custom_v1.create_namespaced_custom_object.assert_called_once() # ty: ignore[possibly-unbound-attribute] + + def test_volcano_job_profile_environment_applied( kubernetes_client_mock, mock_nmp_client,