Update environment variables for JAX backend

PiperOrigin-RevId: 886935376

Author: guptaaka

pathwaysutils/experimental/shared_pathways_service/isc_pathways.py

@@ -163,6 +163,8 @@ def __init__(
     self._proxy_port = None
     self.proxy_server_image = proxy_server_image
     self.proxy_options = proxy_options or ProxyOptions()
+    self._old_jax_platforms = None
+    self._old_jax_backend_target = None
 
   def __repr__(self):
     return (
@@ -198,12 +200,21 @@ def __enter__(self):
           gke_utils.enable_port_forwarding(proxy_pod, PROXY_SERVER_PORT)
       )
 
+      self._old_jax_platforms = os.environ.get(_JAX_PLATFORMS_KEY)
+      self._old_jax_backend_target = os.environ.get(_JAX_BACKEND_TARGET_KEY)
+
       # Update the JAX backend to use the proxy.
+      os.environ[_JAX_PLATFORMS_KEY] = _JAX_PLATFORM_PROXY
+      os.environ[
+          _JAX_BACKEND_TARGET_KEY
+      ] = f"{_JAX_BACKEND_TARGET_HOSTNAME}:{self._proxy_port}"
+
       jax.config.update(_JAX_PLATFORMS_KEY, _JAX_PLATFORM_PROXY)
       jax.config.update(
           _JAX_BACKEND_TARGET_KEY,
           f"{_JAX_BACKEND_TARGET_HOSTNAME}:{self._proxy_port}",
       )
+
       pathwaysutils.initialize()
       _logger.info(
           "Interactive supercomputing proxy client ready for cluster '%s'.",
@@ -221,6 +232,17 @@ def __exit__(self, exc_type, exc_value, traceback):
     _logger.info("Exiting ISCPathways context.")
     self._cleanup()
 
+  def _restore_env_var(self, key: str, original_value: str | None):
+    """Restores an environment variable to its original value or unsets it."""
+    if original_value is None:
+      _logger.info("Unsetting environment variable: %s", key)
+      os.environ.pop(key, None)
+    else:
+      _logger.info(
+          "Restoring environment variable '%s' to '%s'", key, original_value
+      )
+      os.environ[key] = original_value
+
   def _cleanup(self):
     """Cleans up resources created by the ISCPathways context."""
     # 1. Clear JAX caches and run garbage collection.
@@ -248,6 +270,12 @@ def _cleanup(self):
     gke_utils.delete_gke_job(self._proxy_job_name)
     _logger.info("Pathways proxy GKE job deletion complete.")
 
+    # 4. Restore environment variables.
+    _logger.info("Restoring environment variables.")
+    self._restore_env_var(_JAX_PLATFORMS_KEY, self._old_jax_platforms)
+    self._restore_env_var(_JAX_BACKEND_TARGET_KEY, self._old_jax_backend_target)
+    _logger.info("Environment variables restored.")
+
 
 @contextlib.contextmanager
 def connect(