gimlet-api 0.0.5__py3-none-any.whl → 0.0.7__py3-none-any.whl

{gimlet_api-0.0.5.dist-info → gimlet_api-0.0.7.dist-info}/METADATA

@@ -9,8 +9,11 @@ Classifier: Typing :: Typed
 Requires-Python: >=3
 Requires-Dist: protobuf
 Requires-Dist: grpcio
-Requires-Dist: torch
-Requires-Dist: torch_mlir_gml
-Version: 0.0.5
+Requires-Dist: torch>=2.3.0
+Requires-Dist: torch-mlir-gml
+Requires-Dist: numpy<2.0.0
+Requires-Dist: transformers>=4.43.3
+Requires-Dist: safetensors-mlir
+Version: 0.0.7
 
 UNKNOWN

{gimlet_api-0.0.5.dist-info → gimlet_api-0.0.7.dist-info}/RECORD

@@ -1,11 +1,14 @@
-gml/__init__.py,sha256=e81afOgzDBpBNVb3pv_9Y0waGPuX95Y73ofWsq-UVak,718
+gml/__init__.py,sha256=H3WQZ_RaN7VNeb__qeHEbKLEwkaG7gpL5FQ8s1IotUA,773
 gml/_utils.py,sha256=mSCWHhCdzcUvHqmJIB2FS215K1LMgJCWcZ6e6FWK3hQ,1184
-gml/client.py,sha256=4wJ5JeyhMc7UTJazD0GvcHH90Oqiu-x7UY-Q1Zl2Wxo,10640
-gml/compile.py,sha256=pS_mjATIvAH-zxSbIgZeizmpi08571zpOkn5bWRyzM0,2087
-gml/model.py,sha256=Nwf44iTZjMOwFL05JKHaEvTd1QAkJdtgaAuFs6RVUp8,5256
-gml/model_utils.py,sha256=67w32zLI7rjlNsKW6nvCUp-HUjYNMCKMDhjymgJ6EOU,1275
-gml/pipelines.py,sha256=RObsuECms-dGFAMskxlknxfV5-9d5sEOll7McUP5reo,3737
-gml/preprocessing.py,sha256=w-D_GzRndY-Hec8O5UdMSVDXMWtohbwPRxizUjhHsFA,3004
+gml/asset_manager.py,sha256=VnbqUZHPOgPrAh6ri9C0EuNhS8tAHIrbUyJPAJuD9po,2053
+gml/client.py,sha256=5QDKljltBeBTCd2hH38--fTSP0bVVcAvSnWsA9YEFQc,13819
+gml/compile.py,sha256=K4WdC01WkyLlbcoSJzdF2LaVmOgxIkIdc3YNjRFRw9s,10849
+gml/device.py,sha256=sMILurG02aDjL8wrdBW3ftC44WoAPUeZ4Y0yQ0DtaBk,2665
+gml/hf.py,sha256=hi0Af0Q3FM7VvfLB1PkrNai1j7siH6Ouwc_sXf8QE8c,17900
+gml/model.py,sha256=cHFjIEplWCDeSCSl_IPHYzNyv-KOenx-OAsEd-5TpTs,7260
+gml/model_utils.py,sha256=vZvE5cHZIDkUkeZ4Pk4hhV-zOYMiREluv4b8kdqQ3Ig,1375
+gml/pipelines.py,sha256=6tujvMpAACwmEmUfZGFuRn0N8zKvjSVIsIX7FKcPvEU,4301
+gml/preprocessing.py,sha256=STQDSA1_jXPTenJotNtsNMXOc9h1x_wJyQ100LXS6-g,3209
 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
@@ -22,22 +25,26 @@ 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=QPd5gMZvt97ukeHbrgqwroJH4oKTNtjEiAlgfk-FH-k,13496
-gml/proto/src/api/corepb/v1/mediastream_pb2.py,sha256=yNTIivI9LChFifs90eO5dKwVmAR3WGPYQWXHOuXGUD4,6528
-gml/proto/src/api/corepb/v1/model_exec_pb2.py,sha256=U8ujw_7ijKrULiCqxJ4WY9VWiUtIemfIy2e6h1NPrto,25106
+gml/proto/src/api/corepb/v1/controlplane_pb2.py,sha256=ChS3EbDhgtywVrI1fU-CqqynhRpRU-YrdK7gFCaN46w,6008
+gml/proto/src/api/corepb/v1/cp_edge_pb2.py,sha256=O0YhmiQrwy-bFuqYt5Vno2M8m99-qlwyhkfeLBB14WQ,14903
+gml/proto/src/api/corepb/v1/device_info_pb2.py,sha256=j4gvzhM1MsDLMHCAwah1X5xxpfG_5nODs7K83mTY0zI,3425
+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=_WV7zav0uaPHzP-yvjRtUtrwexWiz4eqVclIZmhqdcY,7193
+gml/proto/src/api/corepb/v1/model_exec_pb2.py,sha256=1DM58lSFgfoHk0ui3ZTjDfifgp4dhE7nHvhMwmInpsA,27103
 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_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=AvPGThkFumcElIYY9fjqWdTX9J44RCEJbaYi1F7gn4c,5661
-gml/proto/src/controlplane/logicalpipeline/lppb/v1/lppb_pb2_grpc.py,sha256=q1PugN3Jm_4v5hVWADJLCIkIEC2_beKEqEH4vb_SpH8,7396
+gml/proto/src/controlplane/logicalpipeline/lppb/v1/lppb_pb2.py,sha256=wvLQvoh2UA5qCcMALT6PS47LYmmVdBz9U47WFLs5Ayg,6330
+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=8Ne95YFAkid_jjbtuXFtn_Eu0Hn9u5IaBRuwR7BpzxI,11827
-gimlet_api-0.0.5.dist-info/WHEEL,sha256=sobxWSyDDkdg_rinUth-jxhXHqoNqlmNMJY3aTZn2Us,91
-gimlet_api-0.0.5.dist-info/METADATA,sha256=t2CgErN9FNsBNkrCtmXNZcFlTnWzB3cA1vtL4jA1Y-c,429
-gimlet_api-0.0.5.dist-info/RECORD,,
+gml/tensor.py,sha256=753IsMFYZD7p_f0cQPt4nTIBo5p5S5ELqwCuoHORdMk,14823
+gimlet_api-0.0.7.dist-info/WHEEL,sha256=sobxWSyDDkdg_rinUth-jxhXHqoNqlmNMJY3aTZn2Us,91
+gimlet_api-0.0.7.dist-info/METADATA,sha256=IhN5ODabyw5hfMJDD-LCeGoiZSmqn1kWCKDNay8or5s,531
+gimlet_api-0.0.7.dist-info/RECORD,,

gml/__init__.py

@@ -15,4 +15,5 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from gml.client import Client  # noqa
+from gml.hf import import_huggingface_pipeline  # noqa
 from gml.model import ModelFromFiles, TorchModel  # noqa

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

@@ -21,6 +21,8 @@ from typing import BinaryIO, List, Optional, TextIO, Union
 
 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
+import gml.proto.src.controlplane.compiler.cpb.v1.cpb_pb2_grpc as cpb_grpc
 import gml.proto.src.controlplane.directory.directorypb.v1.directory_pb2 as directorypb
 import gml.proto.src.controlplane.directory.directorypb.v1.directory_pb2_grpc as directorypb_grpc
 import gml.proto.src.controlplane.filetransfer.ftpb.v1.ftpb_pb2 as ftpb
@@ -31,6 +33,7 @@ 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
 from gml.pipelines import Pipeline
 
@@ -65,6 +68,10 @@ class FileAlreadyExists(Exception):
     pass
 
 
+class ModelAlreadyExists(Exception):
+    pass
+
+
 class OrgNotSet(Exception):
     pass
 
@@ -106,6 +113,7 @@ class Client:
 
         self._org_id_cache: Optional[uuidpb.UUID] = None
         self._fts_stub_cache: Optional[ftpb_grpc.FileTransferServiceStub] = None
+        self._cs_stub_cache: Optional[cpb_grpc.CompilerServiceStub] = None
         self._lps_stub_cache: Optional[lppb_grpc.LogicalPipelineServiceStub] = None
         self._ods_stub_cache: Optional[directorypb_grpc.OrgDirectoryServiceStub] = None
         self._ms_stub_cache: Optional[mpb_grpc.ModelServiceStub] = None
@@ -137,6 +145,13 @@ class Client:
             )
         return self._ods_stub_cache
 
+    def _cs_stub(self):
+        if self._cs_stub_cache is None:
+            self._cs_stub_cache = cpb_grpc.CompilerServiceStub(
+                self._channel_factory.get_grpc_channel()
+            )
+        return self._cs_stub_cache
+
     def _ms_stub(self):
         if self._ms_stub_cache is None:
             self._ms_stub_cache = mpb_grpc.ModelServiceStub(
@@ -185,7 +200,7 @@ class Client:
         file_id: uuidpb.UUID,
         sha256: str,
         file: TextIO | BinaryIO,
-        chunk_size=64 * 1024,
+        chunk_size=1024 * 1024,
     ):
         def chunked_requests():
             file.seek(0)
@@ -206,7 +221,7 @@ class Client:
         name: str,
         file: TextIO | BinaryIO,
         sha256: Optional[str] = None,
-        chunk_size=64 * 1024,
+        chunk_size=1024 * 1024,
     ) -> ftpb.FileInfo:
         file_info = self._create_file(name)
 
@@ -237,7 +252,21 @@ class Client:
                 raise Exception("file status is deleted or unknown, cannot re-upload")
         return file_info
 
-    def _create_model(self, model_info: modelexecpb.ModelInfo):
+    def _get_model_if_exists(self, name: str) -> Optional[cpb.Model]:
+        req = mpb.GetModelRequest(
+            name=name,
+            org_id=self._get_org_id(),
+        )
+        stub = self._ms_stub()
+        try:
+            resp = stub.GetModel(req, metadata=self._get_request_metadata())
+            return cpb.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:
         req = mpb.CreateModelRequest(
             org_id=self._get_org_id(),
             name=model_info.name,
@@ -247,26 +276,36 @@ class Client:
         resp = stub.CreateModel(
             req, metadata=self._get_request_metadata(idempotent=True)
         )
-        return resp.id
+        return cpb.Model(id=resp.id, info=model_info)
+
+    def create_model(self, model: Model) -> cpb.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
+                )
+            )
+            return existing_model
 
-    def create_model(self, model: Model):
         model_info = model.to_proto()
-        for asset_name, file in model.collect_assets().items():
-            if isinstance(file, Path) or isinstance(file, str):
-                file = open(file, "rb")
+        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)
+                sha256 = sha256sum(file)
 
-            upload_name = model.name
-            if asset_name:
-                upload_name += ":" + asset_name
-            print(f"Uploading {upload_name}...")
+                upload_name = model.name
+                if asset_name:
+                    upload_name += ":" + asset_name
+                print(f"Uploading {upload_name}...")
 
-            file_info = self._upload_file_if_not_exists(sha256, file, sha256)
+                file_info = self._upload_file_if_not_exists(sha256, file, sha256)
 
-            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)
 
@@ -285,9 +324,7 @@ class Client:
             if isinstance(pipeline, Pipeline):
                 if self._org_name is None:
                     raise ValueError("must set `org` to upload a pipeline")
-                yaml = pipeline.to_yaml(
-                    [model.name for model in models], self._org_name
-                )
+                yaml = pipeline.to_yaml(models, self._org_name)
             else:
                 yaml = pipeline
         else:
@@ -306,3 +343,52 @@ class Client:
             req, metadata=self._get_request_metadata(idempotent=True)
         )
         return resp.id
+
+    def check_compile(
+        self,
+        *,
+        models: List[Model],
+        pipeline_file: Optional[Path] = None,
+        pipeline: Optional[Union[str, Pipeline]] = None,
+        runtimes=List[str],
+        cameras=List[str],
+    ):
+        model_with_assets = []
+        for model in models:
+            created_model = self.create_model(model)
+            model_with_assets.append(created_model)
+
+        if pipeline_file is not None:
+            with open(pipeline_file, "r") as f:
+                yaml = f.read()
+        elif isinstance(pipeline, Pipeline):
+            if self._org_name is None:
+                raise ValueError("must set `org` to compile a Pipeline object")
+            yaml = pipeline.to_yaml(models, self._org_name)
+        elif pipeline is None:
+            raise ValueError("must specify one of 'pipeline_file' or 'pipeline'")
+        else:
+            yaml = pipeline
+
+        stub = self._lps_stub()
+        req = lppb.ParseLogicalPipelineYAMLRequest(
+            yaml=yaml,
+        )
+        resp: lppb.ParseLogicalPipelineYAMLResponse = stub.ParseLogicalPipelineYAML(
+            req, metadata=self._get_request_metadata()
+        )
+        pipeline = resp.logical_pipeline
+
+        capabilities = DeviceCapabilities(runtimes=runtimes, cameras=cameras)
+
+        stub = self._cs_stub()
+        req = cpb.CompileRequest(
+            logical_pipeline=pipeline,
+            device_capabilities=capabilities.to_proto(),
+            models=model_with_assets,
+        )
+        try:
+            stub.Compile(req, metadata=self._get_request_metadata())
+            print("Compilation successful")
+        except grpc.RpcError as e:
+            print("Compilation failed, code {0}: {1}".format(e.code(), e.details()))