diff --git a/nemo_run/run/job.py b/nemo_run/run/job.py index 12243fbc..35a6cf4e 100644 --- a/nemo_run/run/job.py +++ b/nemo_run/run/job.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import logging import sys import traceback from dataclasses import dataclass, field @@ -28,6 +29,7 @@ from nemo_run.core.execution.slurm import SlurmExecutor from nemo_run.core.frontend.console.api import CONSOLE from nemo_run.core.serialization.zlib_json import ZlibJSONSerializer + from nemo_run.run.logs import get_logs from nemo_run.run.plugin import ExperimentPlugin from nemo_run.run.task import direct_run_fn @@ -36,6 +38,8 @@ from nemo_run.run.torchx_backend.runner import Runner from nemo_run.run.torchx_backend.schedulers.api import get_executor_str +logger = logging.getLogger(__name__) + @dataclass class Job(ConfigurableMixin): @@ -98,9 +102,9 @@ def status(self, runner: Runner) -> AppState: status = runner.status(self.handle) state = status.state if status else None except Exception: - ... - finally: - return state or self.state + logger.exception("Failed to get status for job %s", self.handle) + state = None + return state or self.state def logs(self, runner: Runner, regex: str | None = None): get_logs( @@ -316,11 +320,11 @@ def status(self, runner: Runner) -> AppState: status = runner.status(handle) state = status.state if status else None except Exception: - ... - finally: - if not state: - state = AppState.UNKNOWN - new_states.append(state) + logger.exception("Failed to get status for job handle %s", handle) + state = None + if not state: + state = AppState.UNKNOWN + new_states.append(state) self.states = new_states return self.state diff --git a/test/run/test_job.py b/test/run/test_job.py index 896e4841..e4684b0c 100644 --- a/test/run/test_job.py +++ b/test/run/test_job.py @@ -102,7 +102,7 @@ def test_job_status_launched(simple_task, docker_executor, mock_runner): mock_runner.status.assert_called_once_with("test-handle") -def test_job_status_exception(simple_task, docker_executor, mock_runner): +def test_job_status_exception(simple_task, docker_executor, mock_runner, caplog): job = Job( id="test-job", task=simple_task, @@ -113,7 +113,10 @@ def test_job_status_exception(simple_task, docker_executor, mock_runner): ) mock_runner.status.side_effect = Exception("Test exception") - assert job.status(mock_runner) == AppState.RUNNING + with caplog.at_level("ERROR", logger="nemo_run.run.job"): + assert job.status(mock_runner) == AppState.RUNNING + assert "Failed to get status for job test-handle" in caplog.text + assert "Test exception" in caplog.text def test_job_logs(simple_task, docker_executor, mock_runner): @@ -437,7 +440,7 @@ def test_job_group_status_launched(simple_task, docker_executor, mock_runner): mock_runner.status.assert_called_once_with("handle1") -def test_job_group_status_exception(simple_task, docker_executor, mock_runner): +def test_job_group_status_exception(simple_task, docker_executor, mock_runner, caplog): job_group = JobGroup( id="test-group", tasks=[simple_task, simple_task], @@ -448,9 +451,12 @@ def test_job_group_status_exception(simple_task, docker_executor, mock_runner): ) mock_runner.status.side_effect = Exception("Test exception") - status = job_group.status(mock_runner) + with caplog.at_level("ERROR", logger="nemo_run.run.job"): + status = job_group.status(mock_runner) assert status == AppState.UNKNOWN assert job_group.states == [AppState.UNKNOWN] + assert "Failed to get status for job handle handle1" in caplog.text + assert "Test exception" in caplog.text def test_job_group_logs(simple_task, docker_executor, mock_runner):