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

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

@@ -9,11 +9,13 @@ Classifier: Typing :: Typed
 Requires-Python: >=3
 Requires-Dist: protobuf
 Requires-Dist: grpcio
-Requires-Dist: torch>=2.3.0
+Requires-Dist: torch>=2.6.0
 Requires-Dist: torch-mlir-gml
 Requires-Dist: numpy<2.0.0
+Requires-Dist: rich
 Requires-Dist: transformers>=4.43.3
+Requires-Dist: tokenizers>=0.21.0
 Requires-Dist: safetensors-mlir
-Version: 0.0.8
+Version: 0.0.10
 
 UNKNOWN

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

@@ -1,14 +1,14 @@
 gml/__init__.py,sha256=H3WQZ_RaN7VNeb__qeHEbKLEwkaG7gpL5FQ8s1IotUA,773
 gml/_utils.py,sha256=mSCWHhCdzcUvHqmJIB2FS215K1LMgJCWcZ6e6FWK3hQ,1184
 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/client.py,sha256=AcnG5mniHOfq-He-uCph2-xQ39cZwmXZePaUEed87b8,14378
+gml/compile.py,sha256=3L5fpD8DK45RLiywj1b5NuDlbsxpzRxI87k1GahlMpc,9851
+gml/device.py,sha256=Iw71NnuLcgjY32ZMXHlnlPkosTuHEmL9E98utmNChlM,2650
+gml/hf.py,sha256=Kv2yffy8omTRQDPnoIZocG2EOyfhr7UvLFIvTmRxw0g,36170
+gml/model.py,sha256=8fIYlLRduTsUZfYJr_YVPNxbEVIzr7_yaaTe4T-TZ2Y,8429
 gml/model_utils.py,sha256=vZvE5cHZIDkUkeZ4Pk4hhV-zOYMiREluv4b8kdqQ3Ig,1375
-gml/pipelines.py,sha256=d5Vm4eW2RfFZ1SpRebOGEic5sfbLIpMyH4NOy8wdVyI,7319
-gml/preprocessing.py,sha256=MaKkEW4ZP9fjpkJQfpc0X3rCUuSuSmJnGMClHamKmZU,3210
+gml/pipelines.py,sha256=LKj_lh5I5HzyUUIPG4CImiqBnQPrJsj0CHPKhLiOOGo,8374
+gml/preprocessing.py,sha256=YPcxwBOdx0h0ADzoloYbFw9qUGFbi167E8HA4Zwn7Pk,3928
 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
@@ -23,29 +23,31 @@ gml/proto/mediapipe/framework/stream_handler_pb2.py,sha256=kNo-2Fdua_CeyJInI3q5r
 gml/proto/mediapipe/framework/test_calculators_pb2.py,sha256=tXF25VpGtHGArffRqFmjD6FO7xmuCPd5j9UYON2SVSM,2230
 gml/proto/mediapipe/framework/thread_pool_executor_pb2.py,sha256=9TJ66fqSo1BiJmEAQesK0fnVe55zcJpOqVip6HotgyE,2345
 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/metrics/v1/metrics_pb2.py,sha256=k8oW5tmFlJK2574Ky6kDc0JmNNQCLroRwCCGyxDd7JA,9968
 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=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/api/corepb/v1/compiled_pipeline_pb2.py,sha256=g3MxBqshtwaM9_Nrbvwo995_XWq-maXGP6mDeiEzZKo,7529
+gml/proto/src/api/corepb/v1/controlplane_pb2.py,sha256=DylHEVXr36Deh5p-WK8aRwQF-uGW5mJ2mo8pJ3qg7KA,13213
+gml/proto/src/api/corepb/v1/cp_edge_pb2.py,sha256=H0WgAgv6-qaf7wnnKALmSBpD_czmUNHNYpsnE3Tmcrs,14988
+gml/proto/src/api/corepb/v1/deployed_pipeline_pb2.py,sha256=cZjoJuZ3fpCiw2Ox7bcHCXYqRTebb08n-aodwjE-xKI,3053
+gml/proto/src/api/corepb/v1/device_info_pb2.py,sha256=pTZGPjfglje-Wu_-R4qiwPtewXNJIGq5Kedme9SHiaU,6713
+gml/proto/src/api/corepb/v1/gem_config_pb2.py,sha256=vC0g3k9hDv-LhiV6LwaYCly6x00Xx_YA0i2AZSwCo_I,5396
+gml/proto/src/api/corepb/v1/mediastream_pb2.py,sha256=mgi5-prV7Lz0XJ2wo04jGLSvbnDGtdmduSv_6d6I9oA,8368
+gml/proto/src/api/corepb/v1/model_exec_pb2.py,sha256=_TXJvHSxkX1Il6xEVEiFIfei_ZV4KhdL3cSKaMgIYIw,33548
 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=-AfLzprMY7wKKGrNgNYFzSv7OlV3YdYolH-KtrK130s,2839
+gml/proto/src/controlplane/compiler/cpb/v1/cpb_pb2.py,sha256=4mp1QWV7FOzF_nC3RDKZ9vTA-ezMhukcjBEt1lcjGmM,4933
 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
 gml/proto/src/controlplane/filetransfer/ftpb/v1/ftpb_pb2.py,sha256=r8mbJNTq45_c0amPnTr8OFZasCk7XWu2YS_eu7GfWJg,7050
 gml/proto/src/controlplane/filetransfer/ftpb/v1/ftpb_pb2_grpc.py,sha256=XlE4R2PJaOmzQocx7y6SKJvuqt8tYBGzBuhajvzG0cc,12919
-gml/proto/src/controlplane/logicalpipeline/lppb/v1/lppb_pb2.py,sha256=wvLQvoh2UA5qCcMALT6PS47LYmmVdBz9U47WFLs5Ayg,6330
+gml/proto/src/controlplane/logicalpipeline/lppb/v1/lppb_pb2.py,sha256=2s2p6dURKJLboaR965m2-rGTo_63Bi1cXsA90Hz9u-M,6632
 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/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/tensor.py,sha256=aPLm3I3qkYNDcJmntaUycqqN5rsZmcj8ql0EkupJudY,14977
+gimlet_api-0.0.10.dist-info/WHEEL,sha256=sobxWSyDDkdg_rinUth-jxhXHqoNqlmNMJY3aTZn2Us,91
+gimlet_api-0.0.10.dist-info/METADATA,sha256=i3n2dnjznNFL6XFsj1bL0T544E0FmMVQySLgiBkUW04,586
+gimlet_api-0.0.10.dist-info/RECORD,,

gml/client.py

@@ -18,8 +18,12 @@ import os
 import uuid
 from pathlib import Path
 from typing import BinaryIO, List, Optional, TextIO, Union
+from urllib.parse import quote
 
 import grpc
+from rich.progress import (
+    Console,
+)
 
 import gml.proto.src.api.corepb.v1.model_exec_pb2 as modelexecpb
 import gml.proto.src.common.typespb.uuid_pb2 as uuidpb
@@ -39,6 +43,7 @@ from gml.model import Model
 from gml.pipelines import Pipeline
 
 DEFAULT_CONTROLPLANE_ADDR = "app.gimletlabs.ai"
+console = Console()
 
 
 class _ChannelFactory:
@@ -282,31 +287,28 @@ class Client:
     def create_model(self, model: Model) -> modelexecpb.Model:
         existing_model = self._get_model_if_exists(model.name)
         if existing_model is not None:
-            print(
-                'warning: model "{}" already exists and will not be uploaded.'.format(
-                    model.name
-                )
+            console.print(
+                f'[yellow]warning:[/yellow] model "{model.name}" already exists and will not be uploaded.'
             )
             return existing_model
-
         model_info = model.to_proto()
-        with model.collect_assets() as model_assets:
-            for asset_name, file in model_assets.items():
-                if isinstance(file, Path) or isinstance(file, str):
-                    file = open(file, "rb")
-
-                sha256 = sha256sum(file)
+        with console.status(f'Creating model "{model.name}"...'):
+            with model.collect_assets() as model_assets:
+                for asset_name, file in model_assets.items():
+                    if isinstance(file, Path) or isinstance(file, str):
+                        file = open(file, "rb")
 
-                upload_name = model.name
-                if asset_name:
-                    upload_name += ":" + asset_name
-                print(f"Uploading {upload_name}...")
+                    sha256 = sha256sum(file)
 
-                file_info = self._upload_file_if_not_exists(sha256, file, sha256)
+                    upload_name = model.name
+                    if asset_name:
+                        upload_name += ":" + asset_name
+                    file_info = self._upload_file_if_not_exists(sha256, file, sha256)
+                    console.print(f"Uploaded {upload_name}.")
 
-                model_info.file_assets[asset_name].MergeFrom(file_info.file_id)
+                    model_info.file_assets[asset_name].MergeFrom(file_info.file_id)
 
-                file.close()
+                    file.close()
 
         return self._create_model(model_info)
 
@@ -331,6 +333,8 @@ class Client:
         else:
             raise ValueError("must specify one of 'pipeline_file' or 'pipeline'")
 
+        console.print(f'Uploading pipeline "{name}" to {self._org_name}...')
+
         for model in models:
             self.create_model(model)
 
@@ -343,6 +347,11 @@ class Client:
         resp: lppb.CreateLogicalPipelineResponse = stub.CreateLogicalPipeline(
             req, metadata=self._get_request_metadata(idempotent=True)
         )
+
+        url = f"https://{os.getenv('GML_CONTROLPLANE_ADDR')}/orgs/{quote(self._org_name)}/pipelines/{quote(name)}"
+        console.print(
+            f"[green]Pipeline upload complete![/green]\nView your pipeline at: [cyan]{url}[/cyan]"
+        )
         return resp.id
 
     def check_compile(
@@ -351,8 +360,8 @@ class Client:
         models: List[Model],
         pipeline_file: Optional[Path] = None,
         pipeline: Optional[Union[str, Pipeline]] = None,
-        runtimes=List[str],
-        cameras=List[str],
+        runtimes: List[str],
+        cameras: List[str],
     ):
         model_with_assets = []
         for model in models:

gml/compile.py

@@ -16,10 +16,11 @@
 
 import contextlib
 import functools
-from typing import Any, Dict, List, Optional, Sequence, Union
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Union
 
 import safetensors_mlir
 import torch
+import torch.utils._pytree
 import torch_mlir
 from mlir.ir import (
     BF16Type,
@@ -28,6 +29,7 @@ from mlir.ir import (
     F16Type,
     F32Type,
     F64Type,
+    Float8E4M3FNType,
     IntegerType,
     Operation,
     RankedTensorType,
@@ -40,6 +42,7 @@ 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 transformers import DynamicCache
 
 from gml.asset_manager import AssetManager
 from gml.register_submodules import submodule_registration_workarounds
@@ -53,6 +56,45 @@ def _default_decomposition_denylist():
     ]
 
 
+_registered_dynamic_cache_pytree_node = False
+
+
+def register_dynamic_cache_pytree_node():
+    """
+    Registers flattening/unflattening for transformers.DynamicCache
+    Pytree is a representation of tensor collections used inside torch.export.
+    """
+
+    global _registered_dynamic_cache_pytree_node
+    if _registered_dynamic_cache_pytree_node:
+        return
+    _registered_dynamic_cache_pytree_node = True
+
+    def flatten_cache_with_keys(dynamic_cache: DynamicCache):
+        return [
+            (
+                torch.utils._pytree.MappingKey(i),
+                list(value),
+            )
+            for i, value in enumerate(dynamic_cache.to_legacy_cache())
+        ], None
+
+    def flatten_cache(dynamic_cache: DynamicCache):
+        flattened, ctx = flatten_cache_with_keys(dynamic_cache)
+        return [v for _, v in flattened], ctx
+
+    def unflatten_cache(flattened: Iterable[Any], context: Any):
+        return DynamicCache.from_legacy_cache(flattened)
+
+    torch.utils._pytree.register_pytree_node(
+        DynamicCache,
+        flatten_cache,
+        unflatten_cache,
+        serialized_type_name=f"{DynamicCache.__module__}.{DynamicCache.__name__}",
+        flatten_with_keys_fn=flatten_cache_with_keys,
+    )
+
+
 @contextlib.contextmanager
 def _patch_aot_export_module():
     """This contextmanager prevents PyTorch dispatch from running when calling aot_export_module.
@@ -91,6 +133,8 @@ _torch_dtype_to_builtin_element_type = {
     torch.complex32: lambda: ComplexType.get(F16Type.get()),
     torch.complex64: lambda: ComplexType.get(F32Type.get()),
     torch.complex128: lambda: ComplexType.get(F64Type.get()),
+    # Quantized types.
+    torch.float8_e4m3fn: lambda: Float8E4M3FNType.get(),
 }
 
 
@@ -179,6 +223,7 @@ def to_torch_mlir(
     ] = None,
     decomposition_denylist: Optional[List[torch._ops.OperatorBase]] = None,
     weight_manager: Optional[AssetManager] = None,
+    export_predispatch: bool = False,
 ):
     if dynamic_shapes is not None:
         for shape in dynamic_shapes:
@@ -205,10 +250,11 @@ def to_torch_mlir(
         # Ignore errors running the model. This can happen when the model has data dependent branches.
         pass
 
+    register_dynamic_cache_pytree_node()
     prog = _export(
         model,
         tuple(example_inputs),
-        pre_dispatch=False,
+        pre_dispatch=export_predispatch,
         strict=False,
         dynamic_shapes=dynamic_shapes,
     )

gml/device.py

@@ -57,6 +57,8 @@ def _runtime_str_to_runtime_protos(
             return deviceinfopb.ModelRuntimeType.MODEL_RUNTIME_TYPE_TENSORRT
         case "openvino":
             return deviceinfopb.ModelRuntimeType.MODEL_RUNTIME_TYPE_OPENVINO
+        case "hailort":
+            return deviceinfopb.ModelRuntimeType.MODEL_RUNTIME_TYPE_HAILORT
         case _:
             raise ValueError("invalid runtime: {}".format(runtime))