diff --git a/src/dstack/api/_public/runs.py b/src/dstack/api/_public/runs.py index 3fe274c98..18448aab7 100644 --- a/src/dstack/api/_public/runs.py +++ b/src/dstack/api/_public/runs.py @@ -208,6 +208,9 @@ def logs( Args: start_time: Minimal log timestamp. diagnose: Return runner logs if `True`. + replica_num: The replica number or `None` to use any running replica, + falling back to the lowest-numbered replica if no replica is running. + job_num: The job number inside the replica. Yields: Log messages. @@ -217,7 +220,7 @@ def logs( else: job = self._find_job(replica_num=replica_num, job_num=job_num) if job is None: - return [] + return next_token = None while True: resp = self._api_client.logs.poll( @@ -270,7 +273,8 @@ def attach( Args: ssh_identity_file: SSH keypair to access instances. - replica_num: replica_num or None to attach to any running replica. + replica_num: replica_num or None to attach to any running replica, falling back to + the lowest-numbered replica if no replica is running. Raises: dstack.api.PortUsedError: If ports are in use or the run is attached by another process. @@ -418,15 +422,13 @@ def detach(self): self._ssh_attach = None def _find_job(self, replica_num: Optional[int], job_num: int) -> Optional[Job]: - for j in self._run.jobs: - if ( - replica_num is not None - and j.job_spec.replica_num == replica_num - or replica_num is None - and j.job_submissions[-1].status == JobStatus.RUNNING - ) and j.job_spec.job_num == job_num: - return j - return None + jobs = [j for j in self._run.jobs if j.job_spec.job_num == job_num] + if replica_num is not None: + return next((j for j in jobs if j.job_spec.replica_num == replica_num), None) + running = [j for j in jobs if j.job_submissions[-1].status == JobStatus.RUNNING] + # Prefer a running replica, as attaching requires one. Fall back to the lowest-numbered + # replica so that logs remain readable once the run is finished. + return min(running or jobs, key=lambda j: j.job_spec.replica_num, default=None) def __str__(self) -> str: return f"" diff --git a/src/tests/api/test_runs.py b/src/tests/api/test_runs.py index bc52a7587..337750051 100644 --- a/src/tests/api/test_runs.py +++ b/src/tests/api/test_runs.py @@ -1,4 +1,22 @@ -from dstack.api._public.runs import RunCollection +import base64 +import uuid +from datetime import datetime, timezone + +from dstack._internal.core.models.configurations import TaskConfiguration +from dstack._internal.core.models.logs import JobSubmissionLogs, LogEvent, LogEventSource +from dstack._internal.core.models.resources import ResourcesSpec +from dstack._internal.core.models.runs import ( + Job, + JobSpec, + JobStatus, + JobSubmission, + Requirements, + RunSpec, + RunStatus, +) +from dstack._internal.core.models.runs import Run as RunModel +from dstack._internal.server.schemas.logs import PollLogsRequest +from dstack.api._public.runs import Run, RunCollection class _RunsAPI: @@ -27,3 +45,124 @@ def test_default_list_fallback_limits_job_submissions(self): assert api_client.runs.calls[0]["job_submissions_limit"] == 1 assert api_client.runs.calls[1]["job_submissions_limit"] == 1 + + +class _LogsAPI: + def __init__(self, logs_by_job_submission_id: dict[uuid.UUID, list[bytes]]): + self._logs_by_job_submission_id = logs_by_job_submission_id + self.requests: list[PollLogsRequest] = [] + + def poll(self, project_name: str, body: PollLogsRequest) -> JobSubmissionLogs: + self.requests.append(body) + messages = self._logs_by_job_submission_id.get(body.job_submission_id, []) + return JobSubmissionLogs( + logs=[ + LogEvent( + timestamp=datetime(2023, 1, 2, 3, 4, 5, tzinfo=timezone.utc), + log_source=LogEventSource.STDOUT, + message=base64.b64encode(message).decode(), + ) + for message in messages + ] + ) + + +def _get_job(replica_num: int, status: JobStatus, job_num: int = 0) -> Job: + return Job( + job_spec=JobSpec( + replica_num=replica_num, + job_num=job_num, + job_name=f"test-run-{replica_num}-{job_num}", + commands=["echo hello"], + env={}, + image_name="ubuntu:latest", + requirements=Requirements(resources=ResourcesSpec()), + ), + job_submissions=[ + JobSubmission( + id=uuid.uuid4(), + submission_num=0, + submitted_at=datetime(2023, 1, 2, 3, 4, 5, tzinfo=timezone.utc), + last_processed_at=datetime(2023, 1, 2, 3, 4, 5, tzinfo=timezone.utc), + status=status, + ) + ], + ) + + +def _get_run_model(status: RunStatus, jobs: list[Job]) -> RunModel: + return RunModel( + id=uuid.uuid4(), + project_name="main", + user="test", + submitted_at=datetime(2023, 1, 2, 3, 4, 5, tzinfo=timezone.utc), + last_processed_at=datetime(2023, 1, 2, 3, 4, 5, tzinfo=timezone.utc), + status=status, + run_spec=RunSpec( + run_name="test-run", + configuration=TaskConfiguration(commands=["echo hello"], image="ubuntu:latest"), + ), + jobs=jobs, + ) + + +def _get_run(run_model: RunModel, logs_api: _LogsAPI) -> Run: + api_client = _APIClient() + api_client.logs = logs_api + return Run(api_client=api_client, project="main", run=run_model) + + +class TestRunLogs: + def test_returns_logs_of_finished_run(self): + job = _get_job(replica_num=0, status=JobStatus.DONE) + run_model = _get_run_model(status=RunStatus.DONE, jobs=[job]) + job_submission_id = job.job_submissions[-1].id + logs_api = _LogsAPI({job_submission_id: [b"hello\n"]}) + run = _get_run(run_model, logs_api) + + assert b"".join(run.logs()) == b"hello\n" + assert [r.job_submission_id for r in logs_api.requests] == [job_submission_id] + + def test_prefers_running_replica(self): + terminated_job = _get_job(replica_num=0, status=JobStatus.TERMINATED) + running_job = _get_job(replica_num=1, status=JobStatus.RUNNING) + run_model = _get_run_model(status=RunStatus.RUNNING, jobs=[terminated_job, running_job]) + logs_api = _LogsAPI( + { + terminated_job.job_submissions[-1].id: [b"old replica\n"], + running_job.job_submissions[-1].id: [b"new replica\n"], + } + ) + run = _get_run(run_model, logs_api) + + assert b"".join(run.logs()) == b"new replica\n" + + def test_returns_logs_of_lowest_numbered_replica_if_no_replica_is_running(self): + replica_1_job = _get_job(replica_num=1, status=JobStatus.TERMINATED) + replica_0_job = _get_job(replica_num=0, status=JobStatus.TERMINATED) + run_model = _get_run_model( + status=RunStatus.TERMINATED, jobs=[replica_1_job, replica_0_job] + ) + logs_api = _LogsAPI( + { + replica_1_job.job_submissions[-1].id: [b"replica 1\n"], + replica_0_job.job_submissions[-1].id: [b"replica 0\n"], + } + ) + run = _get_run(run_model, logs_api) + + assert b"".join(run.logs()) == b"replica 0\n" + + def test_returns_logs_of_requested_replica(self): + running_job = _get_job(replica_num=0, status=JobStatus.RUNNING) + terminated_job = _get_job(replica_num=1, status=JobStatus.TERMINATED) + run_model = _get_run_model(status=RunStatus.RUNNING, jobs=[running_job, terminated_job]) + logs_api = _LogsAPI( + { + running_job.job_submissions[-1].id: [b"replica 0\n"], + terminated_job.job_submissions[-1].id: [b"replica 1\n"], + } + ) + run = _get_run(run_model, logs_api) + + assert b"".join(run.logs(replica_num=1)) == b"replica 1\n"