diff --git a/pathwaysutils/experimental/shared_pathways_service/yamls/pw-proxy.yaml b/pathwaysutils/experimental/shared_pathways_service/yamls/pw-proxy.yaml index e0ee244..1cbfe4a 100644 --- a/pathwaysutils/experimental/shared_pathways_service/yamls/pw-proxy.yaml +++ b/pathwaysutils/experimental/shared_pathways_service/yamls/pw-proxy.yaml @@ -21,7 +21,10 @@ spec: - --resource_manager_address=${PATHWAYS_HEAD_HOSTNAME}:${PATHWAYS_HEAD_PORT} - --gcs_scratch_location=${GCS_SCRATCH_LOCATION} - --virtual_slices=${EXPECTED_INSTANCES}${PROXY_ARGS} + - --pathways_pipe_unreachable_timeout=60s env: + - name: TPU_SKIP_MDS_QUERY + value: "true" ${PROXY_ENV} ports: - containerPort: ${PROXY_SERVER_PORT} diff --git a/pathwaysutils/profiling.py b/pathwaysutils/profiling.py index e31b1ca..86b00e2 100644 --- a/pathwaysutils/profiling.py +++ b/pathwaysutils/profiling.py @@ -356,13 +356,20 @@ def stop_trace() -> None: try: with _profile_state.lock: if _profile_state.executable is None: - raise RuntimeError("stop_trace called before a trace is being taken!") + _logger.warning( + "stop_trace called before a trace was started; ignoring." + ) + return try: _profile_state.call_profile_executable() finally: _profile_state.reset() finally: - _original_stop_trace() + try: + _original_stop_trace() + except RuntimeError as e: + if "No profile started" not in str(e): + raise _profiler_thread: threading.Thread | None = None diff --git a/pathwaysutils/test/profiling_test.py b/pathwaysutils/test/profiling_test.py index c146234..73c980e 100644 --- a/pathwaysutils/test/profiling_test.py +++ b/pathwaysutils/test/profiling_test.py @@ -403,11 +403,12 @@ def test_stop_trace_with_xprof_options_passes_out_avals(self): self.assertEqual(out_aval.shape, (1,)) self.assertEqual(out_aval.dtype, jnp.object_) - def test_stop_trace_before_start_error(self): - with self.assertRaisesRegex( - RuntimeError, "stop_trace called before a trace is being taken!" - ): + def test_stop_trace_before_start_warns_and_returns(self): + with mock.patch.object(profiling._logger, "warning") as mock_warn: profiling.stop_trace() + mock_warn.assert_called_once_with( + "stop_trace called before a trace was started; ignoring." + ) def test_start_server_starts_thread(self): mock_thread = self.enter_context(