Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions doc/changelog.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion ex/high_throughput_inference/mock_app.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 2 additions & 2 deletions ex/high_throughput_inference/standalone_workermanager.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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,
Expand Down
2 changes: 1 addition & 1 deletion smartsim/_core/launcher/dragon/dragonBackend.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
8 changes: 4 additions & 4 deletions smartsim/_core/mli/infrastructure/environmentloader.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand All @@ -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")
Expand Down
12 changes: 8 additions & 4 deletions tests/dragon/test_environment_loader.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand Down
12 changes: 8 additions & 4 deletions tests/dragon/test_error_handling.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down
2 changes: 1 addition & 1 deletion tests/dragon/test_worker_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down