gimlet-api 0.0.6__py3-none-any.whl → 0.0.8__py3-none-any.whl

{gimlet_api-0.0.6.dist-info → gimlet_api-0.0.8.dist-info}/METADATA

@@ -10,9 +10,10 @@ Requires-Python: >=3
 Requires-Dist: protobuf
 Requires-Dist: grpcio
 Requires-Dist: torch>=2.3.0
-Requires-Dist: torch_mlir_gml==0.0.2
+Requires-Dist: torch-mlir-gml
 Requires-Dist: numpy<2.0.0
 Requires-Dist: transformers>=4.43.3
-Version: 0.0.6
+Requires-Dist: safetensors-mlir
+Version: 0.0.8
 
 UNKNOWN

{gimlet_api-0.0.6.dist-info → gimlet_api-0.0.8.dist-info}/RECORD

@@ -1,13 +1,14 @@
 gml/__init__.py,sha256=H3WQZ_RaN7VNeb__qeHEbKLEwkaG7gpL5FQ8s1IotUA,773
 gml/_utils.py,sha256=mSCWHhCdzcUvHqmJIB2FS215K1LMgJCWcZ6e6FWK3hQ,1184
-gml/client.py,sha256=5QDKljltBeBTCd2hH38--fTSP0bVVcAvSnWsA9YEFQc,13819
-gml/compile.py,sha256=hR8u3LaMiIW8d12FHrvtmtgzUNQq48DYxe8bW-wJ_VY,6054
-gml/device.py,sha256=VUZc6m8QalJ7G9KBKjCY4cIcv2VBd6zAT3ysnh_m1Z0,2585
-gml/hf.py,sha256=GRvEEl9zSIv0iWN91Z6ykFYZ2VdNAVABjZrrzYWUFw4,17792
-gml/model.py,sha256=nXUV6-L4TIkQHCWUpWyG7QJ6YKTZb7eauW9F4pzVTII,6566
+gml/asset_manager.py,sha256=VnbqUZHPOgPrAh6ri9C0EuNhS8tAHIrbUyJPAJuD9po,2053
+gml/client.py,sha256=YFanzPfP619xqBgdfyN_3-Am-jI9eYGvZT8CxhPLTBg,13860
+gml/compile.py,sha256=Ih43r_zU07p91w9aiA0lrPJfmACpAWg0x_HFddMSy7Q,8346
+gml/device.py,sha256=5ocZU_jGUwMfC6PUyAU328Me61BSLwJp8euJCL3mdzo,2550
+gml/hf.py,sha256=e9tw6UGJ1lEZcCplLKo_LgxwTIDWD84DXtQOWZrTR9A,27698
+gml/model.py,sha256=xESdD7tlqn93ym67Lyyk7TZdM3wUqyn7qWdP2AbgdkI,7261
 gml/model_utils.py,sha256=vZvE5cHZIDkUkeZ4Pk4hhV-zOYMiREluv4b8kdqQ3Ig,1375
-gml/pipelines.py,sha256=Bha8J3b5uW8COIejiH12NNF0Tc0XDBt2B3Dez5Jxt4s,5314
-gml/preprocessing.py,sha256=STQDSA1_jXPTenJotNtsNMXOc9h1x_wJyQ100LXS6-g,3209
+gml/pipelines.py,sha256=d5Vm4eW2RfFZ1SpRebOGEic5sfbLIpMyH4NOy8wdVyI,7319
+gml/preprocessing.py,sha256=MaKkEW4ZP9fjpkJQfpc0X3rCUuSuSmJnGMClHamKmZU,3210
 gml/proto/gogoproto/gogo_pb2.py,sha256=WVMIAR8K--mCUkTPM7mEeeXGpQlRRtt_kco10iP3CZs,15728
 gml/proto/mediapipe/framework/calculator_contract_test_pb2.py,sha256=hNjyZCBz3RYa6rN4xR3FOCZKA24gq_LsJ3EMegl5wK4,2031
 gml/proto/mediapipe/framework/calculator_options_pb2.py,sha256=Nq1BQRtLdsIgfkw7ymD3eg2p2_RSlZhiHS7YbDhNHR0,1563
@@ -24,14 +25,16 @@ gml/proto/mediapipe/framework/thread_pool_executor_pb2.py,sha256=9TJ66fqSo1BiJmE
 gml/proto/opentelemetry/proto/common/v1/common_pb2.py,sha256=wQjeDti-C8JiNwRn-z5M5p-Fqxm-SmnbPaoitJcSK-4,2860
 gml/proto/opentelemetry/proto/metrics/v1/metrics_pb2.py,sha256=t2Far6oVcUFQIimzgAkZ8vQd0asMIlvECp4osC0ujgg,9735
 gml/proto/opentelemetry/proto/resource/v1/resource_pb2.py,sha256=cbNmE12Nm3PjW4NXU7-Z-9m_0Zs3Ab8R1xLkDnvclCg,1730
-gml/proto/src/api/corepb/v1/controlplane_pb2.py,sha256=FPNx5fRXj-bnN5mkDUXVz17M33vuHV_hmxH0ggkAUVs,5536
-gml/proto/src/api/corepb/v1/cp_edge_pb2.py,sha256=u41Sohshi6gBfeZO5VnQzfRStFADFzT1Um5mDY9chcg,15309
-gml/proto/src/api/corepb/v1/mediastream_pb2.py,sha256=1qA-ElTgWeGv3oevYlIjK1TIRSgWbR1TTWxA6Q3SOXk,7224
-gml/proto/src/api/corepb/v1/model_exec_pb2.py,sha256=1DM58lSFgfoHk0ui3ZTjDfifgp4dhE7nHvhMwmInpsA,27103
+gml/proto/src/api/corepb/v1/controlplane_pb2.py,sha256=dLZM7rnkWjxHOvayCkK4klFe09GMWLfwPt5MLCkZFzQ,8963
+gml/proto/src/api/corepb/v1/cp_edge_pb2.py,sha256=oIpxq13C1ynK3alzDNZTOL5URxz5qzbDLD9NOM5xxjE,14511
+gml/proto/src/api/corepb/v1/device_info_pb2.py,sha256=-z-FOpAOm3NkcNyRFsENpdW_pqYO1JpmIPtlbWNpH_g,4666
+gml/proto/src/api/corepb/v1/gem_config_pb2.py,sha256=yyEqUqq3-YiX-ByAhbTbZfdh09KuNzEtIYhgk_noJVM,3367
+gml/proto/src/api/corepb/v1/mediastream_pb2.py,sha256=Un9OwDUmWdqv92QP66K-WVOAzxP_4hMoz33JI4W1G5Y,7868
+gml/proto/src/api/corepb/v1/model_exec_pb2.py,sha256=ydwuRVWXNV0ceZ3WVvBIh74rwLWCJHYm-FgzJWhUNE4,28976
 gml/proto/src/common/typespb/jwt_pb2.py,sha256=lxy-bqbyg96i9n_xr2JbkuWX-ldnoJavXPMnApzVSio,5580
 gml/proto/src/common/typespb/status_pb2.py,sha256=IbBJnbsAlvsuTtyT285ZuW6k5VaPfl5kRSOnBxD_H8M,2109
 gml/proto/src/common/typespb/uuid_pb2.py,sha256=5Fm3jYpCPX7sMrP6RhRYsF0SnuZNIBEQJk9f0jwZ2Rw,1188
-gml/proto/src/controlplane/compiler/cpb/v1/cpb_pb2.py,sha256=tkJFPWpndKZy19TFuLKlBfWW1fUQPj0lJLiQ9HfugZU,3213
+gml/proto/src/controlplane/compiler/cpb/v1/cpb_pb2.py,sha256=-AfLzprMY7wKKGrNgNYFzSv7OlV3YdYolH-KtrK130s,2839
 gml/proto/src/controlplane/compiler/cpb/v1/cpb_pb2_grpc.py,sha256=l-gTK9nYpTlVb7QGAckSQXlHhkRdKe2-nrxXc8NQavY,2912
 gml/proto/src/controlplane/directory/directorypb/v1/directory_pb2.py,sha256=KgoUT8ccF-yJPe1r4otQjAPQoKBaQzdBlHoIUSkk0yE,11445
 gml/proto/src/controlplane/directory/directorypb/v1/directory_pb2_grpc.py,sha256=p3OpT8-hfNHu4-29qr-ZahRwO-LoCYM9Q4jomAHTXGA,24572
@@ -41,7 +44,8 @@ gml/proto/src/controlplane/logicalpipeline/lppb/v1/lppb_pb2.py,sha256=wvLQvoh2UA
 gml/proto/src/controlplane/logicalpipeline/lppb/v1/lppb_pb2_grpc.py,sha256=-snjW7n6JveUzJVPFcm25XlL19kowPSKgd61l_jPnHA,9541
 gml/proto/src/controlplane/model/mpb/v1/mpb_pb2.py,sha256=RVedXkNYu2iF5OHiXoYyRw9AGRCUWG7qNyY-5QY71Go,3762
 gml/proto/src/controlplane/model/mpb/v1/mpb_pb2_grpc.py,sha256=KSdb6V04qUHDsb1R2o3wixwTyZgrhwnPYobjnRgWX4I,4735
-gml/tensor.py,sha256=753IsMFYZD7p_f0cQPt4nTIBo5p5S5ELqwCuoHORdMk,14823
-gimlet_api-0.0.6.dist-info/WHEEL,sha256=sobxWSyDDkdg_rinUth-jxhXHqoNqlmNMJY3aTZn2Us,91
-gimlet_api-0.0.6.dist-info/METADATA,sha256=mF73_t-Tn5NPVxLnJBOTCQkKvb0weobtQLmOqSkc4B0,506
-gimlet_api-0.0.6.dist-info/RECORD,,
+gml/register_submodules.py,sha256=U8IwjVygX2vxNi_aK6ljHOD4mmrOhbyVczvy4wwulqU,5027
+gml/tensor.py,sha256=Bv2wshr44ugfdRjWj7JUS8b6_GLBQLZKjGYTojhxm9w,14824
+gimlet_api-0.0.8.dist-info/WHEEL,sha256=sobxWSyDDkdg_rinUth-jxhXHqoNqlmNMJY3aTZn2Us,91
+gimlet_api-0.0.8.dist-info/METADATA,sha256=P4s8-0QNDA4rY8k0blk87Rh3qHfZr-R5A5yKYkwbl2E,531
+gimlet_api-0.0.8.dist-info/RECORD,,

gml/asset_manager.py

@@ -0,0 +1,75 @@
+# Copyright 2023- Gimlet Labs, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import abc
+import tempfile
+from pathlib import Path
+from typing import Dict
+
+
+class AssetManager:
+    @abc.abstractmethod
+    def add_asset(self, name: str) -> Path:
+        pass
+
+    @abc.abstractmethod
+    def assets(self) -> Dict[str, Path]:
+        pass
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc, value, tb):
+        pass
+
+
+class DirectoryAssetManager(AssetManager):
+    def __init__(self, path: str | Path):
+        self.path = Path(path)
+        self._asset_paths: Dict[str, Path] = dict()
+
+    def add_asset(self, name: str) -> Path:
+        path = self.path / name
+        self._asset_paths[name] = path
+        return path
+
+    def assets(self) -> Dict[str, Path]:
+        return self._asset_paths
+
+
+class TempFileAssetManager(AssetManager):
+    def __init__(self):
+        self._assets = dict()
+        self._asset_paths = dict()
+
+    def add_asset(self, name: str) -> Path:
+        tmp = tempfile.NamedTemporaryFile(mode="w")
+        self._assets[name] = tmp
+        file = tmp.__enter__()
+        self._asset_paths[name] = Path(file.name)
+        return self._asset_paths[name]
+
+    def assets(self) -> Dict[str, Path]:
+        return self._asset_paths
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc, value, tb):
+        for tmp in self._assets.values():
+            tmp.__exit__(exc, value, tb)
+        self._assets.clear()
+        self._asset_paths.clear()

gml/client.py

@@ -19,6 +19,8 @@ import uuid
 from pathlib import Path
 from typing import BinaryIO, List, Optional, TextIO, Union
 
+import grpc
+
 import gml.proto.src.api.corepb.v1.model_exec_pb2 as modelexecpb
 import gml.proto.src.common.typespb.uuid_pb2 as uuidpb
 import gml.proto.src.controlplane.compiler.cpb.v1.cpb_pb2 as cpb
@@ -31,7 +33,6 @@ import gml.proto.src.controlplane.logicalpipeline.lppb.v1.lppb_pb2 as lppb
 import gml.proto.src.controlplane.logicalpipeline.lppb.v1.lppb_pb2_grpc as lppb_grpc
 import gml.proto.src.controlplane.model.mpb.v1.mpb_pb2 as mpb
 import gml.proto.src.controlplane.model.mpb.v1.mpb_pb2_grpc as mpb_grpc
-import grpc
 from gml._utils import chunk_file, sha256sum
 from gml.device import DeviceCapabilities
 from gml.model import Model
@@ -252,7 +253,7 @@ class Client:
                 raise Exception("file status is deleted or unknown, cannot re-upload")
         return file_info
 
-    def _get_model_if_exists(self, name: str) -> Optional[cpb.Model]:
+    def _get_model_if_exists(self, name: str) -> Optional[modelexecpb.Model]:
         req = mpb.GetModelRequest(
             name=name,
             org_id=self._get_org_id(),
@@ -260,13 +261,13 @@ class Client:
         stub = self._ms_stub()
         try:
             resp = stub.GetModel(req, metadata=self._get_request_metadata())
-            return cpb.Model(id=resp.id, info=resp.model_info)
+            return modelexecpb.Model(id=resp.id, info=resp.model_info)
         except grpc.RpcError as e:
             if e.code() != grpc.StatusCode.NOT_FOUND:
                 raise e
             return None
 
-    def _create_model(self, model_info: modelexecpb.ModelInfo) -> cpb.Model:
+    def _create_model(self, model_info: modelexecpb.ModelInfo) -> modelexecpb.Model:
         req = mpb.CreateModelRequest(
             org_id=self._get_org_id(),
             name=model_info.name,
@@ -276,9 +277,9 @@ class Client:
         resp = stub.CreateModel(
             req, metadata=self._get_request_metadata(idempotent=True)
         )
-        return cpb.Model(id=resp.id, info=model_info)
+        return modelexecpb.Model(id=resp.id, info=model_info)
 
-    def create_model(self, model: Model) -> cpb.Model:
+    def create_model(self, model: Model) -> modelexecpb.Model:
         existing_model = self._get_model_if_exists(model.name)
         if existing_model is not None:
             print(

gml/compile.py

@@ -16,19 +16,33 @@
 
 import contextlib
 import functools
-from typing import Dict, List, Optional, Sequence, Union
+from typing import Any, Dict, List, Optional, Sequence, Union
 
-import gml.proto.src.api.corepb.v1.model_exec_pb2 as modelexecpb
+import safetensors_mlir
 import torch
-import torch_mlir  # noqa
-
-try:
-    import torch_mlir.fx  # noqa
-    from torch.export import export  # noqa
-
-    has_fx_importer_torch_export = True
-except ImportError:
-    has_fx_importer_torch_export = False
+import torch_mlir
+from mlir.ir import (
+    BF16Type,
+    ComplexType,
+    Context,
+    F16Type,
+    F32Type,
+    F64Type,
+    IntegerType,
+    Operation,
+    RankedTensorType,
+    Value,
+)
+from safetensors.torch import save_file
+from torch._decomp import remove_decompositions
+from torch.export._trace import _export
+from torch_mlir.dialects import torch as torch_d
+from torch_mlir.extras.fx_decomp_util import get_decomposition_table
+from torch_mlir.extras.fx_importer import FxImporter, FxImporterHooks, InputInfo
+from torch_mlir.fx import export_and_import
+
+from gml.asset_manager import AssetManager
+from gml.register_submodules import submodule_registration_workarounds
 
 
 def _default_decomposition_denylist():
@@ -61,33 +75,128 @@ def _patch_aot_export_module():
     torch._functorch.aot_autograd.aot_export_module = orig
 
 
-def to_torch_mlir_w_torch_export(
+_torch_dtype_to_builtin_element_type = {
+    torch.float16: lambda: F16Type.get(),
+    torch.bfloat16: lambda: BF16Type.get(),
+    torch.float32: lambda: F32Type.get(),
+    torch.float64: lambda: F64Type.get(),
+    torch.uint8: lambda: IntegerType.get_unsigned(8),
+    torch.int8: lambda: IntegerType.get_signless(8),
+    torch.int16: lambda: IntegerType.get_signless(16),
+    torch.int32: lambda: IntegerType.get_signless(32),
+    torch.int64: lambda: IntegerType.get_signless(64),
+    torch.bool: lambda: IntegerType.get_signless(1),
+    torch.qint8: lambda: IntegerType.get_signless(8),
+    torch.quint8: lambda: IntegerType.get_unsigned(8),
+    torch.complex32: lambda: ComplexType.get(F16Type.get()),
+    torch.complex64: lambda: ComplexType.get(F32Type.get()),
+    torch.complex128: lambda: ComplexType.get(F64Type.get()),
+}
+
+
+def _get_unique_(tensors, name):
+    index = 0
+    name = "{}_{}".format(name, index)
+    while name in tensors:
+        index += 1
+        name = "{}_{}".format(name, index)
+    return name
+
+
+class TensorSet:
+    def __init__(self):
+        self._tensors: Dict[str, torch.Tensor] = dict()
+
+    def add(self, tensor: torch.Tensor) -> str:
+        shape_desc = "_".join([str(d) for d in tensor.shape])
+        base_name = f"torch_tensor_{shape_desc}_{str(tensor.dtype)}"
+
+        index = 0
+        name = "{}_{}".format(base_name, index)
+        while name in self._tensors and not torch.equal(tensor, self._tensors[name]):
+            index += 1
+            name = "{}_{}".format(base_name, index)
+
+        self._tensors[name] = tensor
+        return name
+
+    def tensors(self) -> Dict[str, torch.Tensor]:
+        return self._tensors
+
+
+class SafetensorImporterHooks(FxImporterHooks):
+    def __init__(self, asset_manager: AssetManager):
+        self._asset_mgr = asset_manager
+        # TODO(james): shard weights into multiple shards.
+        self.asset_name = "weights.shard0"
+        self._tensors = TensorSet()
+
+    def resolve_literal(
+        self,
+        gni: "torch_mlir.extras.fx_importer.GraphNodeImporter",
+        literal: Any,
+        info: Optional[InputInfo],
+    ) -> Optional[Value]:
+        if not isinstance(literal, torch.Tensor):
+            return None
+        tensor = literal
+        ctx = gni._c
+
+        tensor_name = self._tensors.add(tensor)
+
+        file_attr = safetensors_mlir.FileAttr.get(ctx, self.asset_name)
+
+        if tensor.dtype not in _torch_dtype_to_builtin_element_type:
+            raise ValueError("unsupported torch dtype: {}".format(tensor.dtype))
+        elem_type = _torch_dtype_to_builtin_element_type[tensor.dtype]()
+        tensor_type = RankedTensorType.get(tuple(tensor.size()), elem_type)
+
+        tensor_attr = safetensors_mlir.TensorAttr.get(
+            tensor_type, file_attr, tensor_name
+        )
+        builtin_tensor = safetensors_mlir.tensor_ref(tensor_type, tensor_attr)
+
+        vtensor_type = gni._cc.tensor_to_vtensor_type(tensor)
+        return Operation.create(
+            name="torch_c.from_builtin_tensor",
+            results=[vtensor_type],
+            operands=[builtin_tensor],
+        ).result
+
+    def save_tensors(self):
+        file_path = self._asset_mgr.add_asset(self.asset_name)
+        tensors = self._tensors.tensors()
+        for k in tensors:
+            tensors[k] = tensors[k].contiguous()
+        save_file(tensors, file_path)
+
+
+def to_torch_mlir(
     model: torch.nn.Module,
     example_inputs: Sequence[torch.Tensor],
     dynamic_shapes: Optional[
         Sequence[Dict[int, Union[str, "torch.export.dynamic_shapes._Dim"]]]
     ] = None,
     decomposition_denylist: Optional[List[torch._ops.OperatorBase]] = None,
+    weight_manager: Optional[AssetManager] = None,
 ):
-    from torch._decomp import remove_decompositions
-    from torch.export._trace import _export
-    from torch_mlir.extras.fx_decomp_util import get_decomposition_table
-    from torch_mlir.fx import export_and_import
-
     if dynamic_shapes is not None:
         for shape in dynamic_shapes:
             if not isinstance(shape, dict):
                 continue
             for idx in shape:
-                if isinstance(shape[idx], torch.export.dynamic_shapes._Dim):
+                # Assign the value so that pyright understands the type.
+                value = shape[idx]
+                if isinstance(value, torch.export.dynamic_shapes._Dim):
                     continue
-                shape[idx] = torch.export.Dim(shape[idx])
-
+                shape[idx] = torch.export.Dim(value)
     if decomposition_denylist is None:
         decomposition_denylist = _default_decomposition_denylist()
 
     model = model.eval().to("cpu")
 
+    submodule_registration_workarounds(model)
+
     try:
         # Running the model a few times on the inputs, leads to more consistent compiled results.
         for _ in range(2):
@@ -105,76 +214,31 @@ def to_torch_mlir_w_torch_export(
     )
     decomp_table = get_decomposition_table()
     remove_decompositions(decomp_table, decomposition_denylist)
+    hooks = None
+    if weight_manager is not None:
+        hooks = SafetensorImporterHooks(weight_manager)
+
+    context = Context()
+    torch_d.register_dialect(context)
+    safetensors_mlir.register_dialect(context)
+    fx_importer = FxImporter(context=context, hooks=hooks)
+
     with _patch_aot_export_module():
-        return export_and_import(
+        module = export_and_import(
             prog,
             *example_inputs,
             decomposition_table=decomp_table,
+            fx_importer=fx_importer,
         )
 
+    if hooks is not None:
+        hooks.save_tensors()
 
-def to_torch_mlir_fallback(model, example_inputs):
-    from torch.fx.experimental.proxy_tensor import make_fx
-    from torch_mlir import ExampleArgs, OutputType
-    from torch_mlir import compile as torch_mlir_compile
-    from torch_mlir.dynamo import _get_decomposition_table
-
-    example_args = ExampleArgs.get(example_inputs)
-    args = example_args._get_for_tracing(use_tracing=True, ignore_traced_shapes=True)[
-        "forward"
-    ]
     try:
-        # Running the model a few times on the inputs, leads to more consistent compiled results.
-        for _ in range(2):
-            _ = model(*args)
-    except Exception:
-        # Ignore errors running the model. This can happen when the model has data dependent branches.
-        pass
-
-    try:
-        compiled = torch_mlir_compile(
-            model,
-            example_inputs,
-            use_tracing=False,
-            ignore_traced_shapes=False,
-            output_type=OutputType.RAW,
-            use_make_fx=False,
-        )
-        return compiled
-    except Exception:
-        pass
-
-    # If the module can't be exported directly, we try to create an FX graph and then export it.
-    model = make_fx(
-        model, pre_dispatch=True, decomposition_table=_get_decomposition_table()
-    )(*args)
-    compiled = torch_mlir_compile(
-        model,
-        example_inputs,
-        use_tracing=False,
-        ignore_traced_shapes=False,
-        output_type=OutputType.RAW,
-        use_make_fx=False,
-    )
-
-    return compiled
-
-
-def to_torch_mlir(
-    model,
-    example_inputs,
-    dynamic_shapes: Optional[
-        Sequence[Dict[int, Union[str, "torch.export.dynamic_shapes._Dim"]]]
-    ] = None,
-):
-    if has_fx_importer_torch_export:
-        return to_torch_mlir_w_torch_export(model, example_inputs, dynamic_shapes)
-    else:
-        return to_torch_mlir_fallback(model, example_inputs)
-
+        module.operation.verify()
+    except Exception as exc:
+        raise Exception(
+            "failed to verify converted torch model MLIR module: {}".format(module)
+        ) from exc
 
-def torch_mlir_output_kind():
-    if has_fx_importer_torch_export:
-        return modelexecpb.ModelInfo.MODEL_KIND_TORCH
-    else:
-        return modelexecpb.ModelInfo.MODEL_KIND_TORCHSCRIPT
+    return module

gml/device.py

@@ -17,6 +17,7 @@
 from typing import List
 
 import gml.proto.src.api.corepb.v1.cp_edge_pb2 as cpedgepb
+import gml.proto.src.api.corepb.v1.device_info_pb2 as deviceinfopb
 
 
 class DeviceCapabilities:
@@ -27,48 +28,46 @@ class DeviceCapabilities:
     def to_proto(self) -> cpedgepb.DeviceCapabilities:
         return cpedgepb.DeviceCapabilities(
             model_runtimes=[
-                cpedgepb.DeviceCapabilities.ModelRuntimeInfo(
+                deviceinfopb.ModelRuntimeInfo(
                     type=_runtime_str_to_runtime_protos(runtime)
                 )
                 for runtime in self.runtimes
             ],
             cameras=[
-                cpedgepb.DeviceCapabilities.CameraInfo(
+                deviceinfopb.CameraInfo(
                     driver=_camera_driver_str_to_camera_driver_protos(camera),
                     camera_id=str(idx),
                 )
                 for idx, camera in enumerate(self.cameras)
             ],
+            camera_drivers=[
+                deviceinfopb.CameraDriverInfo(
+                    driver=_camera_driver_str_to_camera_driver_protos(camera)
+                )
+                for camera in self.cameras
+            ],
         )
 
 
 def _runtime_str_to_runtime_protos(
     runtime: str,
-) -> cpedgepb.DeviceCapabilities.ModelRuntimeInfo.ModelRuntimeType:
+) -> deviceinfopb.ModelRuntimeType:
     match runtime.lower():
         case "tensorrt":
-            return (
-                cpedgepb.DeviceCapabilities.ModelRuntimeInfo.ModelRuntimeType.MODEL_RUNTIME_TYPE_TENSORRT
-            )
+            return deviceinfopb.ModelRuntimeType.MODEL_RUNTIME_TYPE_TENSORRT
         case "openvino":
-            return (
-                cpedgepb.DeviceCapabilities.ModelRuntimeInfo.ModelRuntimeType.MODEL_RUNTIME_TYPE_OPENVINO
-            )
+            return deviceinfopb.ModelRuntimeType.MODEL_RUNTIME_TYPE_OPENVINO
         case _:
             raise ValueError("invalid runtime: {}".format(runtime))
 
 
 def _camera_driver_str_to_camera_driver_protos(
     driver: str,
-) -> cpedgepb.DeviceCapabilities.CameraInfo.CameraDriver:
+) -> deviceinfopb.CameraDriver:
     match driver.lower():
         case "argus":
-            return (
-                cpedgepb.DeviceCapabilities.CameraInfo.CameraDriver.CAMERA_DRIVER_ARGUS
-            )
+            return deviceinfopb.CameraDriver.CAMERA_DRIVER_ARGUS
         case "v4l2":
-            return (
-                cpedgepb.DeviceCapabilities.CameraInfo.CameraDriver.CAMERA_DRIVER_V4L2
-            )
+            return deviceinfopb.CameraDriver.CAMERA_DRIVER_V4L2
         case _:
             raise ValueError("invalid driver: {}".format(driver))