From 4e5c3f57e91c1efc84a0a86c53b296fa5e61ef31 Mon Sep 17 00:00:00 2001 From: Alyssa Cote Date: Mon, 12 Aug 2024 14:38:01 -0700 Subject: [PATCH 1/2] mli env vars updated --- ex/high_throughput_inference/mock_app.py | 2 +- .../standalone_workermanager.py | 4 ++-- smartsim/_core/launcher/dragon/dragonBackend.py | 2 +- .../_core/mli/infrastructure/environmentloader.py | 8 ++++---- tests/dragon/test_environment_loader.py | 12 ++++++++---- tests/dragon/test_error_handling.py | 12 ++++++++---- tests/dragon/test_worker_manager.py | 2 +- 7 files changed, 25 insertions(+), 17 deletions(-) diff --git a/ex/high_throughput_inference/mock_app.py b/ex/high_throughput_inference/mock_app.py index 7221ee36f..44db70b71 100644 --- a/ex/high_throughput_inference/mock_app.py +++ b/ex/high_throughput_inference/mock_app.py @@ -56,7 +56,7 @@ class ProtoClient: def __init__(self, timing_on: bool): connect_to_infrastructure() - ddict_str = os.environ["SS_INFRA_BACKBONE"] + ddict_str = os.environ["_SMARTSIM_INFRA_BACKBONE"] self._ddict = DDict.attach(ddict_str) self._backbone_descriptor = DragonFeatureStore(self._ddict).descriptor to_worker_fli_str = None diff --git a/ex/high_throughput_inference/standalone_workermanager.py b/ex/high_throughput_inference/standalone_workermanager.py index 2b5ba7df4..982cb6cc3 100644 --- a/ex/high_throughput_inference/standalone_workermanager.py +++ b/ex/high_throughput_inference/standalone_workermanager.py @@ -69,7 +69,7 @@ args = parser.parse_args() connect_to_infrastructure() - ddict_str = os.environ["SS_INFRA_BACKBONE"] + ddict_str = os.environ["_SMARTSIM_INFRA_BACKBONE"] ddict = DDict.attach(ddict_str) to_worker_channel = Channel.make_process_local() @@ -81,7 +81,7 @@ torch_worker = cloudpickle.loads(worker_type_name)() descriptor = base64.b64encode(to_worker_fli_serialized).decode("utf-8") - os.environ["SS_REQUEST_QUEUE"] = descriptor + os.environ["_SMARTSIM_REQUEST_QUEUE"] = descriptor config_loader = EnvironmentConfigLoader( featurestore_factory=DragonFeatureStore.from_descriptor, diff --git a/smartsim/_core/launcher/dragon/dragonBackend.py b/smartsim/_core/launcher/dragon/dragonBackend.py index 4fe6d55ad..daf18e2cb 100644 --- a/smartsim/_core/launcher/dragon/dragonBackend.py +++ b/smartsim/_core/launcher/dragon/dragonBackend.py @@ -515,7 +515,7 @@ def _start_steps(self) -> None: env={ **request.current_env, **request.env, - "SS_INFRA_BACKBONE": self.infra_ddict, + "_SMARTSIM_INFRA_BACKBONE": self.infra_ddict, }, stdout=dragon_process.Popen.PIPE, stderr=dragon_process.Popen.PIPE, diff --git a/smartsim/_core/mli/infrastructure/environmentloader.py b/smartsim/_core/mli/infrastructure/environmentloader.py index b4b9e565c..99202ef2e 100644 --- a/smartsim/_core/mli/infrastructure/environmentloader.py +++ b/smartsim/_core/mli/infrastructure/environmentloader.py @@ -72,8 +72,8 @@ def get_backbone(self) -> t.Optional[FeatureStore]: an environment variable. The backbone is a standalone, system-created feature store used to share internal information among MLI components - :returns: The attached feature store via SS_INFRA_BACKBONE""" - descriptor = os.getenv("SS_INFRA_BACKBONE", "") + :returns: The attached feature store via _SMARTSIM_INFRA_BACKBONE""" + descriptor = os.getenv("_SMARTSIM_INFRA_BACKBONE", "") if not descriptor: logger.warning("No backbone descriptor is configured") @@ -90,8 +90,8 @@ def get_queue(self) -> t.Optional[CommChannelBase]: """Attach to a queue-like communication channel using the descriptor found in an environment variable. - :returns: The attached queue specified via `SS_REQUEST_QUEUE`""" - descriptor = os.getenv("SS_REQUEST_QUEUE", "") + :returns: The attached queue specified via `_SMARTSIM_REQUEST_QUEUE`""" + descriptor = os.getenv("_SMARTSIM_REQUEST_QUEUE", "") if not descriptor: logger.warning("No queue descriptor is configured") diff --git a/tests/dragon/test_environment_loader.py b/tests/dragon/test_environment_loader.py index 6ae5d2b30..8f2716488 100644 --- a/tests/dragon/test_environment_loader.py +++ b/tests/dragon/test_environment_loader.py @@ -55,7 +55,9 @@ def test_environment_loader_attach_fli(content: bytes, monkeypatch: pytest.Monke """A descriptor can be stored, loaded, and reattached""" chan = Channel.make_process_local() queue = FLInterface(main_ch=chan) - monkeypatch.setenv("SS_REQUEST_QUEUE", du.B64.bytes_to_str(queue.serialize())) + monkeypatch.setenv( + "_SMARTSIM_REQUEST_QUEUE", du.B64.bytes_to_str(queue.serialize()) + ) config = EnvironmentConfigLoader( featurestore_factory=DragonFeatureStore.from_descriptor, @@ -76,7 +78,9 @@ def test_environment_loader_serialize_fli(monkeypatch: pytest.MonkeyPatch): queue are the same""" chan = Channel.make_process_local() queue = FLInterface(main_ch=chan) - monkeypatch.setenv("SS_REQUEST_QUEUE", du.B64.bytes_to_str(queue.serialize())) + monkeypatch.setenv( + "_SMARTSIM_REQUEST_QUEUE", du.B64.bytes_to_str(queue.serialize()) + ) config = EnvironmentConfigLoader( featurestore_factory=DragonFeatureStore.from_descriptor, @@ -89,7 +93,7 @@ def test_environment_loader_serialize_fli(monkeypatch: pytest.MonkeyPatch): def test_environment_loader_flifails(monkeypatch: pytest.MonkeyPatch): """An incorrect serialized descriptor will fails to attach""" - monkeypatch.setenv("SS_REQUEST_QUEUE", "randomstring") + monkeypatch.setenv("_SMARTSIM_REQUEST_QUEUE", "randomstring") config = EnvironmentConfigLoader( featurestore_factory=DragonFeatureStore.from_descriptor, callback_factory=None, @@ -104,7 +108,7 @@ def test_environment_loader_backbone_load_dfs(monkeypatch: pytest.MonkeyPatch): """Verify the dragon feature store is loaded correctly by the EnvironmentConfigLoader to demonstrate featurestore_factory correctness""" feature_store = DragonFeatureStore(DDict()) - monkeypatch.setenv("SS_INFRA_BACKBONE", feature_store.descriptor) + monkeypatch.setenv("_SMARTSIM_INFRA_BACKBONE", feature_store.descriptor) config = EnvironmentConfigLoader( featurestore_factory=DragonFeatureStore.from_descriptor, diff --git a/tests/dragon/test_error_handling.py b/tests/dragon/test_error_handling.py index 208ab1e5e..5603269b2 100644 --- a/tests/dragon/test_error_handling.py +++ b/tests/dragon/test_error_handling.py @@ -89,9 +89,11 @@ def setup_worker_manager_model_bytes( chan = Channel.make_process_local() queue = FLInterface(main_ch=chan) - monkeypatch.setenv("SS_REQUEST_QUEUE", du.B64.bytes_to_str(queue.serialize())) + monkeypatch.setenv( + "_SMARTSIM_REQUEST_QUEUE", du.B64.bytes_to_str(queue.serialize()) + ) # Put backbone descriptor into env var for the `EnvironmentConfigLoader` - monkeypatch.setenv("SS_INFRA_BACKBONE", backbone_descriptor) + monkeypatch.setenv("_SMARTSIM_INFRA_BACKBONE", backbone_descriptor) worker_manager = WorkerManager( EnvironmentConfigLoader( @@ -127,9 +129,11 @@ def setup_worker_manager_model_key( chan = Channel.make_process_local() queue = FLInterface(main_ch=chan) - monkeypatch.setenv("SS_REQUEST_QUEUE", du.B64.bytes_to_str(queue.serialize())) + monkeypatch.setenv( + "_SMARTSIM_REQUEST_QUEUE", du.B64.bytes_to_str(queue.serialize()) + ) # Put backbone descriptor into env var for the `EnvironmentConfigLoader` - monkeypatch.setenv("SS_INFRA_BACKBONE", backbone_descriptor) + monkeypatch.setenv("_SMARTSIM_INFRA_BACKBONE", backbone_descriptor) worker_manager = WorkerManager( EnvironmentConfigLoader( diff --git a/tests/dragon/test_worker_manager.py b/tests/dragon/test_worker_manager.py index 864e14993..c8332c260 100644 --- a/tests/dragon/test_worker_manager.py +++ b/tests/dragon/test_worker_manager.py @@ -167,7 +167,7 @@ def test_worker_manager(prepare_environment: pathlib.Path) -> None: # NOTE: env vars should be set prior to instantiating EnvironmentConfigLoader # or test environment may be unable to send messages w/queue descriptor = base64.b64encode(to_worker_fli_serialized).decode("utf-8") - os.environ["SS_REQUEST_QUEUE"] = descriptor + os.environ["_SMARTSIM_REQUEST_QUEUE"] = descriptor config_loader = EnvironmentConfigLoader( featurestore_factory=DragonFeatureStore.from_descriptor, From 196364c59a2ba6b1bed8df7cfcc4ad903ded770f Mon Sep 17 00:00:00 2001 From: Alyssa Cote Date: Mon, 12 Aug 2024 14:43:30 -0700 Subject: [PATCH 2/2] changelog --- doc/changelog.md | 1 + 1 file changed, 1 insertion(+) diff --git a/doc/changelog.md b/doc/changelog.md index 80dd23cf4..18a0fed70 100644 --- a/doc/changelog.md +++ b/doc/changelog.md @@ -13,6 +13,7 @@ Jump to: Description +- Update MLI environment variables using new naming convention - Reduce a copy by using torch.from_numpy instead of torch.tensor - Enable dynamic feature store selection - Fix dragon package installation bug