Fix Pathways OSS container regression and pretraining test issues.

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}

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

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(