Merge pull request #150 from proxystore/issue-134

gpauloski · web-flow · commit eb1b58704be6 · 2024-09-06T17:27:42.000-05:00
Add `async_resolve` option to `ProxyTransformer` and make `RunConfig.env_vars` optional
diff --git a/taps/run/config.py b/taps/run/config.py
@@ -59,8 +59,8 @@ class RunConfig(BaseModel):
             '"{executor}" for formatting).'
         ),
     )
-    env_vars: Dict[str, str] = Field(  # noqa: UP006
-        default_factory=dict,
+    env_vars: Optional[Dict[str, str]] = Field(  # noqa: UP006,UP007
+        None,
         description='Environment variables to set during benchmarking.',
     )
 
diff --git a/taps/run/main.py b/taps/run/main.py
@@ -80,7 +80,8 @@ def run(config: Config, run_dir: pathlib.Path) -> None:
     config.write_toml(CONFIG_FILENAME)
     logger.debug(f'Wrote config to {CONFIG_FILENAME}')
 
-    with update_environment(config.run.env_vars):
+    env_vars = config.run.env_vars if config.run.env_vars is not None else {}
+    with update_environment(env_vars):
         with Timer() as app_init_timer:
             app = config.app.get_app()
         logger.log(
diff --git a/taps/run/utils.py b/taps/run/utils.py
@@ -187,8 +187,8 @@ def update_environment(
     os.environ.update(variables)
     if len(variables) > 0:
         logger.debug(
-            f'Updated {len(variables)} environment variable(s) '
-            f'({", ".join(variables.keys())})',
+            f'Updated {len(variables)} environment variable(s): '
+            f'{", ".join(variables.keys())}',
         )
 
     try:
@@ -200,6 +200,6 @@ def update_environment(
         os.environ.update(previous)
         if len(previous) > 0:
             logger.debug(
-                f'Restored {len(previous)} environment variable(s) '
-                f'({",".join(previous.keys())})',
+                f'Restored {len(previous)} environment variable(s): '
+                f'{",".join(previous.keys())}',
             )
diff --git a/taps/transformer/_proxy.py b/taps/transformer/_proxy.py
@@ -1,15 +1,23 @@
 from __future__ import annotations
 
+import sys
 from typing import Any
 from typing import Literal
 from typing import TypeVar
 
+if sys.version_info >= (3, 11):  # pragma: >=3.11 cover
+    from typing import Self
+else:  # pragma: <3.11 cover
+    from typing_extensions import Self
+
 from proxystore.proxy import extract
 from proxystore.proxy import Proxy
 from proxystore.store import get_store
 from proxystore.store import Store
 from proxystore.store.config import ConnectorConfig
+from proxystore.store.utils import resolve_async
 from pydantic import Field
+from pydantic import model_validator
 
 from taps.plugins import register
 from taps.transformer._protocol import TransformerConfig
@@ -29,17 +37,34 @@ class ProxyTransformerConfig(TransformerConfig):
         description='Connector configuration.',
     )
     cache_size: int = Field(16, description='cache size')
+    async_resolve: bool = Field(
+        False,
+        description=(
+            'Asynchronously resolve proxies. Not compatible with '
+            'extract_target=True.'
+        ),
+    )
     extract_target: bool = Field(
         False,
         description=(
-            'Extract the target from the proxy when resolving the identifier.'
+            'Extract the target from the proxy when resolving the identifier. '
+            'Not compatible with async_resolve=True.'
         ),
     )
     populate_target: bool = Field(
         True,
         description='Populate target objects of newly created proxies.',
     )
 
+    @model_validator(mode='after')
+    def _validate_mutex_options(self) -> Self:
+        if self.async_resolve and self.extract_target:
+            raise ValueError(
+                'Options async_resolve and extract_target cannot be '
+                'enabled at the same time.',
+            )
+        return self
+
     def get_transformer(self) -> ProxyTransformer:
         """Create a transformer from the configuration."""
         connector = self.connector.get_connector()
@@ -51,6 +76,7 @@ def get_transformer(self) -> ProxyTransformer:
                 populate_target=self.populate_target,
                 register=True,
             ),
+            async_resolve=self.async_resolve,
             extract_target=self.extract_target,
         )
 
@@ -62,29 +88,43 @@ class ProxyTransformer:
 
     Args:
         store: Store instance to use for proxying objects.
+        async_resolve: Begin asynchronously resolving proxies when the
+            transformer resolves a proxy (which is otherwise a no-op unless
+            `extract_target=True`). Not compatible with `extract_target=True`.
         extract_target: When `True`, resolving an identifier (i.e., a proxy)
             will return the target object. Otherwise, the proxy is returned
-            since a proxy can act as the target object.
+            since a proxy can act as the target object. Not compatible
+            with `async_resolve=True`.
     """
 
     def __init__(
         self,
         store: Store[Any],
         *,
+        async_resolve: bool = False,
         extract_target: bool = False,
     ) -> None:
+        if async_resolve and extract_target:
+            raise ValueError(
+                'Options async_resolve and extract_target cannot be '
+                'enabled at the same time.',
+            )
+
         self.store = store
+        self.async_resolve = async_resolve
         self.extract_target = extract_target
 
     def __repr__(self) -> str:
         ctype = type(self).__name__
         store = f'store={self.store}'
+        async_ = f'async_resolve={self.async_resolve}'
         extract = f'extract_target={self.extract_target}'
-        return f'{ctype}({store}, {extract})'
+        return f'{ctype}({store}, {async_}, {extract})'
 
     def __getstate__(self) -> dict[str, Any]:
         return {
             'config': self.store.config(),
+            'async_resolve': self.async_resolve,
             'extract_target': self.extract_target,
         }
 
@@ -94,6 +134,7 @@ def __setstate__(self, state: dict[str, Any]) -> None:
             self.store = store
         else:
             self.store = Store.from_config(state['config'])
+        self.async_resolve = state['async_resolve']
         self.extract_target = state['extract_target']
 
     def close(self) -> None:
@@ -125,4 +166,8 @@ def resolve(self, identifier: Proxy[T]) -> T | Proxy[T]:
             The resolved object or a proxy of the resolved object depending \
             on the setting of `extract_target`.
         """
-        return extract(identifier) if self.extract_target else identifier
+        if self.extract_target:
+            return extract(identifier)
+        if self.async_resolve:
+            resolve_async(identifier)
+        return identifier
diff --git a/tests/transformer/proxy_test.py b/tests/transformer/proxy_test.py
@@ -10,6 +10,7 @@
 from proxystore.store import Store
 from proxystore.store import unregister_store
 from proxystore.store.config import ConnectorConfig
+from pydantic import ValidationError
 
 from taps.transformer import ProxyTransformer
 from taps.transformer import ProxyTransformerConfig
@@ -26,14 +27,40 @@ def test_file_config(tmp_path: pathlib.Path) -> None:
     transformer.close()
 
 
-@pytest.mark.parametrize('extract', (True, False))
-def test_proxy_transformer(extract: bool) -> None:
+def test_config_validation_error(tmp_path: pathlib.Path) -> None:
+    with pytest.raises(
+        ValidationError,
+        match='Options async_resolve and extract_target cannot be enabled',
+    ):
+        ProxyTransformerConfig(
+            connector=ConnectorConfig(
+                kind='file',
+                options={'store_dir': str(tmp_path)},
+            ),
+            async_resolve=True,
+            extract_target=True,
+        )
+
+
+@pytest.mark.parametrize(
+    ('extract', 'async_'),
+    (
+        (False, False),
+        (True, False),
+        (False, True),
+    ),
+)
+def test_proxy_transformer(extract: bool, async_: bool) -> None:
     with Store(
         'test-proxy-transformer',
         LocalConnector(),
         register=True,
     ) as store:
-        transformer = ProxyTransformer(store, extract_target=extract)
+        transformer = ProxyTransformer(
+            store,
+            async_resolve=async_,
+            extract_target=extract,
+        )
         assert isinstance(repr(transformer), str)
 
         obj = [1, 2, 3]
@@ -46,6 +73,23 @@ def test_proxy_transformer(extract: bool) -> None:
         transformer.close()
 
 
+def test_proxy_transformer_value_error() -> None:
+    with Store(
+        'test-proxy-transformer-value-error',
+        LocalConnector(),
+        register=True,
+    ) as store:
+        with pytest.raises(
+            ValueError,
+            match='Options async_resolve and extract_target cannot be enabled',
+        ):
+            ProxyTransformer(
+                store,
+                async_resolve=True,
+                extract_target=True,
+            )
+
+
 def test_proxy_transformer_pickling() -> None:
     name = 'test-proxy-transformer-pickle'
     with Store(name, LocalConnector(), register=True) as store:

Original file line number	Diff line number	Diff line change
`@@ -59,8 +59,8 @@ class RunConfig(BaseModel):`
`59`	`59`	`'"{executor}" for formatting).'`
`60`	`60`	`),`
`61`	`61`	`)`
`62`		`- env_vars: Dict[str, str] = Field( # noqa: UP006`
`63`		`- default_factory=dict,`
	`62`	`+ env_vars: Optional[Dict[str, str]] = Field( # noqa: UP006,UP007`
	`63`	`+ None,`
`64`	`64`	`description='Environment variables to set during benchmarking.',`
`65`	`65`	`)`
`66`	`66`