Add FeatureStore descriptor to tensor & model keys

doc/changelog.md

@@ -13,6 +13,7 @@ Jump to:
 
 Description
 
+- Enable dynamic feature store selection
 - Fix dragon package installation bug
 - Adjust schemas for better performance
 - Add TorchWorker first implementation and mock inference app example

ex/high_throughput_inference/mli_driver.py

@@ -1,5 +1,4 @@
-
-
+import argparse
 import os
 import base64
 import cloudpickle
@@ -26,11 +25,23 @@
 
 torch_worker_str = base64.b64encode(cloudpickle.dumps(TorchWorker)).decode("ascii")
 
-worker_manager_rs = exp.create_run_settings(sys.executable, [worker_manager_script_name, "--device", device, "--worker_class", torch_worker_str])
+worker_manager_rs = exp.create_run_settings(
+    sys.executable,
+    [
+        worker_manager_script_name,
+        "--device",
+        device,
+        "--worker_class",
+        torch_worker_str,
+    ],
+)
 worker_manager = exp.create_model("worker_manager", run_settings=worker_manager_rs)
 worker_manager.attach_generator_files(to_copy=[worker_manager_script_name])
 
-app_rs = exp.create_run_settings(sys.executable, exe_args = [app_script_name, "--device", device])
+app_rs = exp.create_run_settings(
+    sys.executable,
+    exe_args=[app_script_name, "--device", device],
+)
 app = exp.create_model("app", run_settings=app_rs)
 app.attach_generator_files(to_copy=[app_script_name], to_symlink=[model_name])
 
@@ -47,4 +58,4 @@
         break
     time.sleep(5)
 
-print("Exiting.")
+print("Exiting.")

ex/high_throughput_inference/mock_app.py

@@ -44,16 +44,21 @@
 import numbers
 
 from collections import OrderedDict
+from smartsim._core.mli.infrastructure.storage.dragonfeaturestore import (
+    DragonFeatureStore,
+)
 from smartsim._core.mli.message_handler import MessageHandler
 from smartsim.log import get_logger
 
 logger = get_logger("App")
 
+
 class ProtoClient:
     def __init__(self, timing_on: bool):
         connect_to_infrastructure()
-        ddict_str = os.environ["SS_DRG_DDICT"]
+        ddict_str = os.environ["SS_INFRA_BACKBONE"]
         self._ddict = DDict.attach(ddict_str)
+        self._backbone_descriptor = DragonFeatureStore(self._ddict).descriptor
         to_worker_fli_str = None
         while to_worker_fli_str is None:
             try:
@@ -88,39 +93,45 @@ def start_timings(self, batch_size: int):
     def end_timings(self):
         if self._timing_on:
             self._add_label_to_timings("total_time")
-            self._timings["total_time"].append(self._format_number(time.perf_counter()-self._start))
+            self._timings["total_time"].append(
+                self._format_number(time.perf_counter() - self._start)
+            )
 
     def measure_time(self, label: str):
         if self._timing_on:
             self._add_label_to_timings(label)
-            self._timings[label].append(self._format_number(time.perf_counter()-self._interm))
+            self._timings[label].append(
+                self._format_number(time.perf_counter() - self._interm)
+            )
             self._interm = time.perf_counter()
 
     def print_timings(self, to_file: bool = False):
         print(" ".join(self._timings.keys()))
-        value_array = numpy.array([value for  value in self._timings.values()], dtype=float)
+        value_array = numpy.array(
+            [value for value in self._timings.values()], dtype=float
+        )
         value_array = numpy.transpose(value_array)
         for i in range(value_array.shape[0]):
             print(" ".join(self._format_number(value) for value in value_array[i]))
         if to_file:
             numpy.save("timings.npy", value_array)
             numpy.savetxt("timings.txt", value_array)
 
-
     def run_model(self, model: bytes | str, batch: torch.Tensor):
         tensors = [batch.numpy()]
         self.start_timings(batch.shape[0])
         built_tensor_desc = MessageHandler.build_tensor_descriptor(
-            "c", "float32", list(batch.shape))
+            "c", "float32", list(batch.shape)
+        )
         self.measure_time("build_tensor_descriptor")
         built_model = None
         if isinstance(model, str):
-            model_arg = MessageHandler.build_model_key(model)
+            model_arg = MessageHandler.build_model_key(model, self._backbone_descriptor)
         else:
             model_arg = MessageHandler.build_model(model, "resnet-50", "1.0")
         request = MessageHandler.build_request(
             reply_channel=self._from_worker_ch_serialized,
-            model= model_arg,
+            model=model_arg,
             inputs=[built_tensor_desc],
             outputs=[],
             output_descriptors=[],
@@ -129,10 +140,12 @@ def run_model(self, model: bytes | str, batch: torch.Tensor):
         self.measure_time("build_request")
         request_bytes = MessageHandler.serialize_request(request)
         self.measure_time("serialize_request")
-        with self._to_worker_fli.sendh(timeout=None, stream_channel=self._to_worker_ch) as to_sendh:
+        with self._to_worker_fli.sendh(
+            timeout=None, stream_channel=self._to_worker_ch
+        ) as to_sendh:
             to_sendh.send_bytes(request_bytes)
             for t in tensors:
-                to_sendh.send_bytes(t.tobytes()) #TODO NOT FAST ENOUGH!!!
+                to_sendh.send_bytes(t.tobytes())  # TODO NOT FAST ENOUGH!!!
                 # to_sendh.send_bytes(bytes(t.data))
         logger.info(f"Message size: {len(request_bytes)} bytes")
 
@@ -159,7 +172,7 @@ def set_model(self, key: str, model: bytes):
         self._ddict[key] = model
 
 
-class ResNetWrapper():
+class ResNetWrapper:
     def __init__(self, name: str, model: str):
         self._model = torch.jit.load(model)
         self._name = name
@@ -168,7 +181,7 @@ def __init__(self, name: str, model: str):
         torch.jit.save(scripted, buffer)
         self._serialized_model = buffer.getvalue()
 
-    def get_batch(self, batch_size: int=32):
+    def get_batch(self, batch_size: int = 32):
         return torch.randn((batch_size, 3, 224, 224), dtype=torch.float32)
 
     @property
@@ -179,6 +192,7 @@ def model(self):
     def name(self):
         return self._name
 
+
 if __name__ == "__main__":
 
     parser = argparse.ArgumentParser("Mock application")
@@ -194,8 +208,8 @@ def name(self):
 
     for batch_size in [1, 2, 4, 8, 16, 32, 64, 128]:
         logger.info(f"Batch size: {batch_size}")
-        for iteration_number in range(total_iterations + int(batch_size==1)):
+        for iteration_number in range(total_iterations + int(batch_size == 1)):
             logger.info(f"Iteration: {iteration_number}")
             client.run_model(resnet.name, resnet.get_batch(batch_size))
 
-    client.print_timings(to_file=True)
+    client.print_timings(to_file=True)

ex/high_throughput_inference/standalone_workermanager.py

@@ -31,17 +31,19 @@
 from dragon.data.ddict.ddict import DDict
 from dragon.utils import b64decode, b64encode
 from dragon.globalservices.api_setup import connect_to_infrastructure
+
 # isort: on
 import argparse
 import base64
 import cloudpickle
-import pickle
+import optparse
 import os
 
 from smartsim._core.mli.comm.channel.dragonchannel import DragonCommChannel
-from smartsim._core.mli.infrastructure.storage.dragonfeaturestore import DragonFeatureStore
 from smartsim._core.mli.comm.channel.dragonfli import DragonFLIChannel
-from smartsim._core.mli.infrastructure.worker.torch_worker import TorchWorker
+from smartsim._core.mli.infrastructure.storage.dragonfeaturestore import (
+    DragonFeatureStore,
+)
 from smartsim._core.mli.infrastructure.control.workermanager import WorkerManager
 from smartsim._core.mli.infrastructure.environmentloader import EnvironmentConfigLoader
 
@@ -67,30 +69,31 @@
 
     args = parser.parse_args()
     connect_to_infrastructure()
-    ddict_str = os.environ["SS_DRG_DDICT"]
+    ddict_str = os.environ["SS_INFRA_BACKBONE"]
     ddict = DDict.attach(ddict_str)
 
     to_worker_channel = Channel.make_process_local()
     to_worker_fli = fli.FLInterface(main_ch=to_worker_channel, manager_ch=None)
     to_worker_fli_serialized = to_worker_fli.serialize()
     ddict["to_worker_fli"] = to_worker_fli_serialized
 
-    torch_worker = cloudpickle.loads(base64.b64decode(args.worker_class.encode('ascii')))()
-
-    dfs = DragonFeatureStore(ddict)
-    comm_channel = DragonFLIChannel(to_worker_fli_serialized)
+    worker_type_name = base64.b64decode(args.worker_class.encode("ascii"))
+    torch_worker = cloudpickle.loads(worker_type_name)()
 
-    os.environ["SSFeatureStore"] = base64.b64encode(pickle.dumps(dfs)).decode("utf-8")
-    os.environ["SSQueue"] = base64.b64encode(to_worker_fli_serialized).decode("utf-8")
+    descriptor = base64.b64encode(to_worker_fli_serialized).decode("utf-8")
+    os.environ["SS_REQUEST_QUEUE"] = descriptor
 
-    config_loader = EnvironmentConfigLoader()
+    config_loader = EnvironmentConfigLoader(
+        featurestore_factory=DragonFeatureStore.from_descriptor,
+        callback_factory=DragonCommChannel,
+        queue_factory=DragonFLIChannel.from_descriptor,
+    )
 
     worker_manager = WorkerManager(
         config_loader=config_loader,
         worker=torch_worker,
         as_service=True,
         cooldown=10,
-        comm_channel_type=DragonCommChannel,
-        device = args.device,
+        device=args.device,
     )
     worker_manager.execute()

smartsim/_core/_cli/scripts/dragon_install.py

@@ -1,7 +1,9 @@
 import os
 import pathlib
+import shutil
 import sys
 import typing as t
+from urllib.request import urlretrieve
 
 from github import Github
 from github.GitReleaseAsset import GitReleaseAsset
@@ -160,13 +162,26 @@ def retrieve_asset(working_dir: pathlib.Path, asset: GitReleaseAsset) -> pathlib
 
     # if we've previously downloaded the release and still have
     # wheels laying around, use that cached version instead
-    if download_dir.exists() and list(download_dir.rglob("*.whl")):
+    if download_dir.exists() or list(download_dir.rglob("*.whl")):
         return download_dir
 
-    archive = WebTGZ(asset.browser_download_url)
+    download_dir.mkdir(parents=True, exist_ok=True)
+
+    # grab a copy of the complete asset
+    asset_path = download_dir / str(asset.name)
+    download_url = asset.browser_download_url
+
+    try:
+        urlretrieve(download_url, str(asset_path))
+        logger.debug(f"Retrieved asset {asset.name} from {download_url}")
+    except Exception:
+        logger.exception(f"Unable to download asset from: {download_url}")
+
+    # extract the asset
+    archive = WebTGZ(download_url)
     archive.extract(download_dir)
 
-    logger.debug(f"Retrieved {asset.browser_download_url} to {download_dir}")
+    logger.debug(f"Extracted {download_url} to {download_dir}")
     return download_dir
 
 

smartsim/_core/launcher/dragon/dragonBackend.py

@@ -457,33 +457,27 @@ def create_run_policy(
         if isinstance(request, DragonRunRequest):
             run_request: DragonRunRequest = request
 
-            affinity = dragon_policy.Policy.Affinity.DEFAULT
             cpu_affinity: t.List[int] = []
             gpu_affinity: t.List[int] = []
 
             # Customize policy only if the client requested it, otherwise use default
             if run_request.policy is not None:
                 # Affinities are not mutually exclusive. If specified, both are used
                 if run_request.policy.cpu_affinity:
-                    affinity = dragon_policy.Policy.Affinity.SPECIFIC
                     cpu_affinity = run_request.policy.cpu_affinity
 
                 if run_request.policy.gpu_affinity:
-                    affinity = dragon_policy.Policy.Affinity.SPECIFIC
                     gpu_affinity = run_request.policy.gpu_affinity
             logger.debug(
-                f"Affinity strategy: {affinity}, "
                 f"CPU affinity mask: {cpu_affinity}, "
                 f"GPU affinity mask: {gpu_affinity}"
             )
-            if affinity != dragon_policy.Policy.Affinity.DEFAULT:
-                return dragon_policy.Policy(
-                    placement=dragon_policy.Policy.Placement.HOST_NAME,
-                    host_name=node_name,
-                    affinity=affinity,
-                    cpu_affinity=cpu_affinity,
-                    gpu_affinity=gpu_affinity,
-                )
+            return dragon_policy.Policy(
+                placement=dragon_policy.Policy.Placement.HOST_NAME,
+                host_name=node_name,
+                cpu_affinity=cpu_affinity,
+                gpu_affinity=gpu_affinity,
+            )
 
         return dragon_policy.Policy(
             placement=dragon_policy.Policy.Placement.HOST_NAME,
@@ -521,7 +515,7 @@ def _start_steps(self) -> None:
                         env={
                             **request.current_env,
                             **request.env,
-                            "SS_DRG_DDICT": self.infra_ddict,
+                            "SS_INFRA_BACKBONE": self.infra_ddict,
                         },
                         stdout=dragon_process.Popen.PIPE,
                         stderr=dragon_process.Popen.PIPE,

smartsim/_core/mli/comm/channel/channel.py

@@ -42,11 +42,13 @@ def __init__(self, descriptor: t.Union[str, bytes]) -> None:
     @abstractmethod
     def send(self, value: bytes) -> None:
         """Send a message through the underlying communication channel
+
         :param value: The value to send"""
 
     @abstractmethod
     def recv(self) -> t.List[bytes]:
         """Receieve a message through the underlying communication channel
+
         :returns: the received message"""
 
     @property

smartsim/_core/mli/comm/channel/dragonchannel.py

@@ -24,6 +24,7 @@
 # OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
+import base64
 import sys
 import typing as t
 
@@ -55,7 +56,23 @@ def send(self, value: bytes) -> None:
 
     def recv(self) -> t.List[bytes]:
         """Receieve a message through the underlying communication channel
+
         :returns: the received message"""
         with self._channel.recvh(timeout=None) as recvh:
             message_bytes: bytes = recvh.recv_bytes(timeout=None)
             return [message_bytes]
+
+    @classmethod
+    def from_descriptor(
+        cls,
+        descriptor: str,
+    ) -> "DragonCommChannel":
+        """A factory method that creates an instance from a descriptor string
+
+        :param descriptor: The descriptor that uniquely identifies the resource
+        :returns: An attached DragonCommChannel"""
+        try:
+            return DragonCommChannel(base64.b64decode(descriptor))
+        except:
+            logger.error(f"Failed to create dragon comm channel: {descriptor}")
+            raise