import & mypy fix

ankona · Aug 29, 2024 · add944f · add944f
1 parent 8314caa
commit add944f
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 17 deletions.
diff --git a/smartsim/_core/mli/infrastructure/storage/backbonefeaturestore.py b/smartsim/_core/mli/infrastructure/storage/backbonefeaturestore.py
@@ -108,7 +108,7 @@ def backend_channel(self) -> t.Optional[str]:
 
         :returns: a stringified channel descriptor"""
         if self.MLI_NOTIFY_CONSUMERS in self:
-            return self[self.MLI_NOTIFY_CONSUMERS]
+            return str(self[self.MLI_NOTIFY_CONSUMERS])
         return None
 
     @backend_channel.setter
@@ -125,7 +125,7 @@ def worker_queue(self) -> t.Optional[str]:
 
         :returns: a stringified channel descriptor"""
         if self.MLI_WORKER_QUEUE in self:
-            return self[self.MLI_WORKER_QUEUE]
+            return str(self[self.MLI_WORKER_QUEUE])
         return None
 
     @worker_queue.setter

diff --git a/tests/dragon/test_featurestore.py b/tests/dragon/test_featurestore.py
@@ -34,9 +34,6 @@
 
 import pytest
 
-from smartsim._core.mli.comm.channel.dragonfli import DragonFLIChannel
-from smartsim.error.errors import SmartSimError
-
 dragon = pytest.importorskip("dragon")
 
 from smartsim._core.mli.comm.channel.dragonchannel import DragonCommChannel

diff --git a/tests/dragon/test_protoclient.py b/tests/dragon/test_protoclient.py
@@ -30,6 +30,8 @@
 
 import pytest
 
+dragon = pytest.importorskip("dragon")
+
 from smartsim._core.mli.comm.channel.dragonfli import DragonFLIChannel
 from smartsim._core.mli.infrastructure.storage.backbonefeaturestore import (
     BackboneFeatureStore,
@@ -41,8 +43,6 @@
 from smartsim.error.errors import SmartSimError
 from smartsim.log import get_logger
 
-dragon = pytest.importorskip("dragon")
-
 # isort: off
 from dragon import fli
 from dragon.channels import Channel

diff --git a/tests/dragon/test_worker_manager.py b/tests/dragon/test_worker_manager.py
@@ -31,18 +31,16 @@
 
 import pytest
 
-from smartsim._core.mli.comm.channel.dragonchannel import DragonCommChannel
-from smartsim._core.mli.infrastructure.storage.backbonefeaturestore import (
-    BackboneFeatureStore,
-)
-from smartsim._core.mli.mli_schemas.tensor.tensor_capnp import OutputDescriptor
-
 torch = pytest.importorskip("torch")
 dragon = pytest.importorskip("dragon")
 
-import base64
 import multiprocessing as mp
 
+from smartsim._core.mli.infrastructure.storage.backbonefeaturestore import (
+    BackboneFeatureStore,
+)
+from smartsim._core.mli.mli_schemas.tensor.tensor_capnp import OutputDescriptor
+
 try:
     mp.set_start_method("dragon")
 except Exception:
@@ -56,7 +54,6 @@
 from dragon.data.ddict.ddict import DDict
 from dragon.mpbridge.queues import DragonQueue
 
-from smartsim._core.mli.comm.channel.channel import CommChannelBase
 from smartsim._core.mli.comm.channel.dragonfli import DragonFLIChannel
 from smartsim._core.mli.infrastructure.control.workermanager import (
     EnvironmentConfigLoader,
@@ -65,12 +62,10 @@
 from smartsim._core.mli.infrastructure.storage.dragonfeaturestore import (
     DragonFeatureStore,
 )
-from smartsim._core.mli.infrastructure.storage.featurestore import FeatureStore
 from smartsim._core.mli.infrastructure.worker.torch_worker import TorchWorker
 from smartsim._core.mli.message_handler import MessageHandler
 from smartsim.log import get_logger
 
-from .featurestore import FileSystemFeatureStore
 from .utils.channel import FileSystemCommChannel
 
 logger = get_logger(__name__)