From a57a841469cb360663c7231e3d43ed37fae6b8a6 Mon Sep 17 00:00:00 2001
From: Ben Clifford <benc@hawaga.org.uk>
Date: Thu, 24 Oct 2024 15:11:44 +0000
Subject: [PATCH] GCE now accepts user_endpoint_config via resource
 specification

---
 mypy.ini                          |  3 +++
 parsl/executors/globus_compute.py | 18 ++++++++++++++----
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/mypy.ini b/mypy.ini
index 4b64a12de2..e46e11fd63 100644
--- a/mypy.ini
+++ b/mypy.ini
@@ -177,6 +177,9 @@ ignore_missing_imports = True
 #[mypy-multiprocessing.synchronization.*]
 #ignore_missing_imports = True
 
+[mypy-globus_compute_sdk.*]
+ignore_missing_imports = True
+
 [mypy-pandas.*]
 ignore_missing_imports = True
 
diff --git a/parsl/executors/globus_compute.py b/parsl/executors/globus_compute.py
index cac59be20c..ac484beee4 100644
--- a/parsl/executors/globus_compute.py
+++ b/parsl/executors/globus_compute.py
@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import uuid
 from concurrent.futures import Future
 from typing import Any, Callable, Dict, Optional, Union
@@ -11,7 +13,6 @@
     _globus_compute_enabled = True
 except ImportError:
     _globus_compute_enabled = False
-    Client: Any  # type: ignore[no-redef]
 
 UUID_LIKE_T = Union[uuid.UUID, str]
 
@@ -54,7 +55,9 @@ def __init__(
         user_endpoint_config:
             User endpoint configuration values as described
             and allowed by endpoint administrators. Must be a JSON-serializable dict
-            or None.
+            or None. Refer docs from `globus-compute
+            <https://globus-compute.readthedocs.io/en/latest/endpoints/endpoints.html#templating-endpoint-configuration>`_
+            for more info.
 
         label:
             a label to name the executor; mainly utilized for
@@ -118,11 +121,16 @@ def submit(self, func: Callable, resource_specification: Dict[str, Any], *args:
 
         func: Callable
             Python function to execute remotely
+
         resource_specification: Dict[str, Any]
-            Resource specification used to run MPI applications on Endpoints configured
-            to use globus compute's MPIEngine
+            Resource specification can be used specify MPI resources required by MPI applications on
+            Endpoints configured to use globus compute's MPIEngine. GCE also accepts *user_endpoint_config*
+            to configure endpoints when the endpoint is a `Multi-User Endpoint
+            <https://globus-compute.readthedocs.io/en/latest/endpoints/endpoints.html#templating-endpoint-configuration>`_
+
         args:
             Args to pass to the function
+
         kwargs:
             kwargs to pass to the function
 
@@ -132,6 +140,8 @@ def submit(self, func: Callable, resource_specification: Dict[str, Any], *args:
         Future
         """
         self._executor.resource_specification = resource_specification or self.resource_specification
+        # Pop user_endpoint_config since it is illegal in resource_spec for globus_compute
+        self._executor.user_endpoint_config = resource_specification.pop('user_endpoint_config', self.user_endpoint_config)
         return self._executor.submit(func, *args, **kwargs)
 
     def shutdown(self, wait=True, *, cancel_futures=False):