gimlet-api 0.0.1__py3-none-any.whl → 0.0.4__py3-none-any.whl

gimlet_api-0.0.4.dist-info/METADATA

@@ -0,0 +1,16 @@
+Metadata-Version: 2.1
+Name: gimlet-api
+Author: Gimlet Labs, Inc.
+Author-email: support@gimletlabs.ai
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Typing :: Typed
+Requires-Python: >=3
+Requires-Dist: protobuf
+Requires-Dist: grpcio
+Requires-Dist: torch
+Requires-Dist: torch_mlir_gml
+Version: 0.0.4
+
+UNKNOWN

gimlet_api-0.0.4.dist-info/RECORD

@@ -0,0 +1,43 @@
+gml/__init__.py,sha256=e81afOgzDBpBNVb3pv_9Y0waGPuX95Y73ofWsq-UVak,718
+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=91I0kCyix1e2ZR7UAdxtyxOIS6pTEfUQ3G62MBkYx5k,5221
+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/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
+gml/proto/mediapipe/framework/calculator_pb2.py,sha256=sP2SNeN0mp3YXtWaGFKMDdC3R_uclrLxnYgxWwTy4p8,9092
+gml/proto/mediapipe/framework/calculator_profile_pb2.py,sha256=uy2yO6DffTL7f4HNuxvpztSZefVE3-SVxvnyxR49Ya4,6335
+gml/proto/mediapipe/framework/mediapipe_options_pb2.py,sha256=84T1x8HgcwmtF0Eq8ZKeNMWwG6_COAr2246hZE611n0,1318
+gml/proto/mediapipe/framework/packet_factory_pb2.py,sha256=AfSB4dAuJbDUNgay93Yi56Z9_DLjJZ-XeD6gNZeuxCc,1916
+gml/proto/mediapipe/framework/packet_generator_pb2.py,sha256=ORDDp5UBLwMPhYiXoAbW8TSLZ6QTaMUeSrmh5jwHPpA,2167
+gml/proto/mediapipe/framework/packet_test_pb2.py,sha256=K8Uo_1VzSg-u5HFgNjFLQptVsakzHNFeWJz3eldto5g,1797
+gml/proto/mediapipe/framework/status_handler_pb2.py,sha256=dgiW2ohm-ho07z1k4TM_Xt0MLvC1ZtDgkhhXDByshMk,1707
+gml/proto/mediapipe/framework/stream_handler_pb2.py,sha256=kNo-2Fdua_CeyJInI3q5r9IoAUanjhk9jh01Z1KXu6Q,2043
+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/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/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/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/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.4.dist-info/WHEEL,sha256=sobxWSyDDkdg_rinUth-jxhXHqoNqlmNMJY3aTZn2Us,91
+gimlet_api-0.0.4.dist-info/METADATA,sha256=Bm_ca7tpLR7foub0QOavwvJ9_2_4FQlFPr6n8zwoJBU,429
+gimlet_api-0.0.4.dist-info/RECORD,,

{gimlet_api-0.0.1.dist-info → gimlet_api-0.0.4.dist-info}/WHEEL

@@ -1,5 +1,4 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.43.0)
+Generator: bazel-wheelmaker 1.0
 Root-Is-Purelib: true
 Tag: py3-none-any
-

gml/__init__.py

@@ -0,0 +1,18 @@
+# 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
+
+from gml.client import Client  # noqa
+from gml.model import ModelFromFiles, TorchModel  # noqa

gml/_utils.py

@@ -0,0 +1,42 @@
+# 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 hashlib
+import os
+from typing import BinaryIO, TextIO
+
+
+def chunk_file(f: TextIO | BinaryIO, chunk_size=64 * 1024):
+    while True:
+        chunk = f.read(chunk_size)
+        if not chunk:
+            break
+        yield chunk
+
+
+def sha256sum(f: BinaryIO, buffer_size=64 * 1024):
+    sha256sum = hashlib.sha256()
+    for chunk in chunk_file(f, buffer_size):
+        sha256sum.update(chunk)
+
+    return sha256sum.hexdigest()
+
+
+def get_file_size(f: TextIO | BinaryIO):
+    f.seek(0, os.SEEK_END)
+    file_size = f.tell()
+    f.seek(0)
+    return file_size

gml/client.py

@@ -0,0 +1,308 @@
+# 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 os
+import uuid
+from pathlib import Path
+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.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
+import gml.proto.src.controlplane.filetransfer.ftpb.v1.ftpb_pb2_grpc as ftpb_grpc
+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.model import Model
+from gml.pipelines import Pipeline
+
+DEFAULT_CONTROLPLANE_ADDR = "app.gimletlabs.ai"
+
+
+class _ChannelFactory:
+    """
+    _ChannelFactory creates grpc channels to a controlplane.
+    """
+
+    def __init__(self, controlplane_addr: str, insecure_no_ssl=False):
+        self.controlplane_addr = controlplane_addr
+        self.insecure_no_ssl = insecure_no_ssl
+
+        self._channel_cache: grpc.Channel = None
+
+    def get_grpc_channel(self) -> grpc.Channel:
+        if self._channel_cache is not None:
+            return self._channel_cache
+        return self._create_grpc_channel()
+
+    def _create_grpc_channel(self) -> grpc.Channel:
+        if self.insecure_no_ssl:
+            return grpc.insecure_channel(self.controlplane_addr)
+
+        creds = grpc.ssl_channel_credentials()
+        return grpc.secure_channel(self.controlplane_addr, creds)
+
+
+class FileAlreadyExists(Exception):
+    pass
+
+
+class OrgNotSet(Exception):
+    pass
+
+
+class APIKeyNotSet(Exception):
+    pass
+
+
+class Client:
+    """
+    Client provides authorized access to a controlplane.
+    """
+
+    def __init__(
+        self,
+        api_key: Optional[str] = None,
+        controlplane_addr: Optional[str] = None,
+        org: Optional[str] = None,
+        insecure_no_ssl: bool = False,
+    ):
+        self._org_name = org
+        self._api_key = api_key
+        if self._api_key is None:
+            self._api_key = os.getenv("GML_API_KEY")
+            if self._api_key is None:
+                raise APIKeyNotSet(
+                    "must provide api_key explicitly or through environment variable GML_API_KEY"
+                )
+
+        self._controlplane_addr = controlplane_addr
+        if self._controlplane_addr is None:
+            self._controlplane_addr = os.getenv("GML_CONTROLPLANE_ADDR")
+            if self._controlplane_addr is None:
+                self._controlplane_addr = DEFAULT_CONTROLPLANE_ADDR
+
+        self._channel_factory = _ChannelFactory(
+            self._controlplane_addr, insecure_no_ssl=insecure_no_ssl
+        )
+
+        self._org_id_cache: Optional[uuidpb.UUID] = None
+        self._fts_stub_cache: Optional[ftpb_grpc.FileTransferServiceStub] = 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
+
+    def _get_request_metadata(self, idempotent=False):
+        md = [("x-api-key", self._api_key)]
+        if idempotent:
+            md.append(("x-idempotency-key", uuid.uuid4().hex))
+        return md
+
+    def _fts_stub(self):
+        if self._fts_stub_cache is None:
+            self._fts_stub_cache = ftpb_grpc.FileTransferServiceStub(
+                self._channel_factory.get_grpc_channel()
+            )
+        return self._fts_stub_cache
+
+    def _lps_stub(self):
+        if self._lps_stub_cache is None:
+            self._lps_stub_cache = lppb_grpc.LogicalPipelineServiceStub(
+                self._channel_factory.get_grpc_channel()
+            )
+        return self._lps_stub_cache
+
+    def _ods_stub(self):
+        if self._ods_stub_cache is None:
+            self._ods_stub_cache = directorypb_grpc.OrgDirectoryServiceStub(
+                self._channel_factory.get_grpc_channel()
+            )
+        return self._ods_stub_cache
+
+    def _ms_stub(self):
+        if self._ms_stub_cache is None:
+            self._ms_stub_cache = mpb_grpc.ModelServiceStub(
+                self._channel_factory.get_grpc_channel()
+            )
+        return self._ms_stub_cache
+
+    def _get_org_id(self):
+        if self._org_name is None:
+            raise OrgNotSet("organization not set for method that is org specific")
+        stub = self._ods_stub()
+        req = directorypb.GetOrgRequest(org_name=self._org_name)
+        resp: directorypb.GetOrgResponse = stub.GetOrg(
+            req, metadata=self._get_request_metadata()
+        )
+        return resp.org_info.id
+
+    def _org_id(self):
+        if self._org_id_cache is None:
+            self._org_id_cache = self._get_org_id()
+        return self._org_id_cache
+
+    def _create_file(self, name: str) -> ftpb.FileInfo:
+        stub = self._fts_stub()
+        try:
+            req = ftpb.CreateFileInfoRequest(name=name)
+            resp: ftpb.CreateFileInfoResponse = stub.CreateFileInfo(
+                req, metadata=self._get_request_metadata()
+            )
+        except grpc.RpcError as e:
+            if e.code() == grpc.StatusCode.ALREADY_EXISTS:
+                raise FileAlreadyExists(
+                    f"A file already exists with name: {name}"
+                ) from e
+            raise e
+        return resp.info
+
+    def _file_info_by_name(self, name: str) -> ftpb.GetFileInfoByNameResponse:
+        stub = self._fts_stub()
+
+        req = ftpb.GetFileInfoByNameRequest(name=name)
+        return stub.GetFileInfoByName(req, metadata=self._get_request_metadata()).info
+
+    def _upload_created_file(
+        self,
+        file_id: uuidpb.UUID,
+        sha256: str,
+        file: TextIO | BinaryIO,
+        chunk_size=64 * 1024,
+    ):
+        def chunked_requests():
+            file.seek(0)
+            for chunk in chunk_file(file, chunk_size):
+                req = ftpb.UploadFileRequest(
+                    file_id=file_id, sha256sum=sha256, chunk=chunk
+                )
+                yield req
+
+        stub = self._fts_stub()
+        resp: ftpb.UploadFileResponse = stub.UploadFile(
+            chunked_requests(), metadata=self._get_request_metadata()
+        )
+        return resp
+
+    def upload_file(
+        self,
+        name: str,
+        file: TextIO | BinaryIO,
+        sha256: Optional[str] = None,
+        chunk_size=64 * 1024,
+    ) -> ftpb.FileInfo:
+        file_info = self._create_file(name)
+
+        if sha256 is None:
+            sha256 = sha256sum(file)
+        self._upload_created_file(file_info.file_id, sha256, file, chunk_size)
+        return self._file_info_by_name(name)
+
+    def _upload_file_if_not_exists(
+        self,
+        name: str,
+        file: TextIO | BinaryIO,
+        sha256: Optional[str] = None,
+    ) -> ftpb.FileInfo:
+        file_info: Optional[ftpb.FileInfo] = None
+        try:
+            file_info = self.upload_file(name, file, sha256)
+        except FileAlreadyExists:
+            file_info = self._file_info_by_name(name)
+
+        match file_info.status:
+            case ftpb.FILE_STATUS_READY:
+                pass
+            case ftpb.FILE_STATUS_CREATED:
+                self._upload_created_file(file_info.file_id, sha256, file)
+                file_info = self._file_info_by_name(name)
+            case _:
+                raise Exception("file status is deleted or unknown, cannot re-upload")
+        return file_info
+
+    def _create_model(self, model_info: modelexecpb.ModelInfo):
+        req = mpb.CreateModelRequest(
+            org_id=self._get_org_id(),
+            name=model_info.name,
+            model_info=model_info,
+        )
+        stub = self._ms_stub()
+        resp = stub.CreateModel(
+            req, metadata=self._get_request_metadata(idempotent=True)
+        )
+        return resp.id
+
+    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")
+
+            sha256 = sha256sum(file)
+
+            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)
+
+            model_info.file_assets[asset_name].MergeFrom(file_info.file_id)
+
+            file.close()
+
+        return self._create_model(model_info)
+
+    def upload_pipeline(
+        self,
+        *,
+        name: str,
+        models: List[Model],
+        pipeline_file: Optional[Path] = None,
+        pipeline: Optional[Union[str, Pipeline]] = None,
+    ) -> uuidpb.UUID:
+        if pipeline_file is not None:
+            with open(pipeline_file, "r") as f:
+                yaml = f.read()
+        elif pipeline is not None:
+            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
+                )
+            else:
+                yaml = pipeline
+        else:
+            raise ValueError("must specify one of 'pipeline_file' or 'pipeline'")
+
+        for model in models:
+            self.create_model(model)
+
+        stub = self._lps_stub()
+        req = lppb.CreateLogicalPipelineRequest(
+            org_id=self._org_id(),
+            name=name,
+            yaml=yaml,
+        )
+        resp: lppb.CreateLogicalPipelineResponse = stub.CreateLogicalPipeline(
+            req, metadata=self._get_request_metadata(idempotent=True)
+        )
+        return resp.id

gml/compile.py

@@ -1,19 +1,55 @@
-from torch.fx.experimental.proxy_tensor import make_fx
+# 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
 
-#from torch_mlir.torchscript import compile as torch_mlir_compile
+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 import OutputType
 from torch_mlir.dynamo import _get_decomposition_table
 
 
 def to_torch_mlir(model, example_inputs):
-    # Running the model a few times on the inputs, leads to more consistent compiled results.
-    # TODO(james): would be good to understand why this is the case.
-    for _ in range(2):
-        _ = model(*example_inputs)
+    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
 
-    model = make_fx(model, pre_dispatch=True, decomposition_table=_get_decomposition_table())(*example_inputs)
+    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,

gml/model.py

@@ -0,0 +1,150 @@
+# 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 io
+from pathlib import Path
+from typing import BinaryIO, Dict, List, Literal, Optional, TextIO, Tuple
+
+import gml.proto.src.api.corepb.v1.model_exec_pb2 as modelexecpb
+import torch
+import torch_mlir
+from gml.compile import to_torch_mlir
+from gml.preprocessing import ImagePreprocessingStep
+from gml.tensor import TensorSemantics
+
+
+class Model(abc.ABC):
+    def __init__(
+        self,
+        name: str,
+        kind: modelexecpb.ModelInfo.ModelKind,
+        storage_format: modelexecpb.ModelInfo.ModelStorageFormat,
+        input_tensor_semantics: List[TensorSemantics],
+        output_tensor_semantics: List[TensorSemantics],
+        class_labels: Optional[List[str]] = None,
+        class_labels_file: Optional[Path] = None,
+        image_preprocessing_steps: Optional[List[ImagePreprocessingStep]] = None,
+    ):
+        self.name = name
+        self.kind = kind
+        self.storage_format = storage_format
+        self.class_labels = class_labels
+        if class_labels_file:
+            self.class_labels = []
+            with open(class_labels_file, "r") as f:
+                for line in f.readlines():
+                    self.class_labels.append(line.strip())
+        self.input_tensor_semantics = input_tensor_semantics
+        self.output_tensor_semantics = output_tensor_semantics
+        self.image_preprocessing_steps = image_preprocessing_steps
+
+    def to_proto(self) -> modelexecpb.ModelInfo:
+        image_preprocessing_steps = None
+        if self.image_preprocessing_steps:
+            image_preprocessing_steps = [
+                step.to_proto() for step in self.image_preprocessing_steps
+            ]
+        return modelexecpb.ModelInfo(
+            name=self.name,
+            kind=self.kind,
+            format=self.storage_format,
+            class_labels=self.class_labels,
+            image_preprocessing_steps=image_preprocessing_steps,
+            input_tensor_semantics=[
+                semantics.to_proto() for semantics in self.input_tensor_semantics
+            ],
+            output_tensor_semantics=[
+                semantics.to_proto() for semantics in self.output_tensor_semantics
+            ],
+        )
+
+    @abc.abstractmethod
+    def collect_assets(self) -> Dict[str, TextIO | BinaryIO | Path]:
+        pass
+
+
+class TorchModel(Model):
+    def __init__(
+        self,
+        name: str,
+        torch_module: torch.nn.Module,
+        input_shapes: List[List[int]],
+        input_dtypes: List[torch.dtype],
+        **kwargs,
+    ):
+        super().__init__(
+            name,
+            modelexecpb.ModelInfo.MODEL_KIND_TORCHSCRIPT,
+            modelexecpb.ModelInfo.MODEL_STORAGE_FORMAT_MLIR_TEXT,
+            **kwargs,
+        )
+        self.torch_module = torch_module
+        self.input_shapes = input_shapes
+        self.input_dtypes = input_dtypes
+
+    def _convert_to_torch_mlir(self):
+        return to_torch_mlir(
+            self.torch_module.to("cpu"),
+            [
+                torch_mlir.TensorPlaceholder(shape, dtype)
+                for shape, dtype in zip(self.input_shapes, self.input_dtypes)
+            ],
+        )
+
+    def collect_assets(self) -> Dict[str, TextIO | BinaryIO | Path]:
+        compiled = self._convert_to_torch_mlir()
+        file = io.BytesIO(str(compiled).encode("utf-8"))
+        return {"": file}
+
+
+def _kind_str_to_kind_format_protos(
+    kind: str,
+) -> Tuple[modelexecpb.ModelInfo.ModelKind, modelexecpb.ModelInfo.ModelStorageFormat]:
+    match kind.lower():
+        case "openvino":
+            return (
+                modelexecpb.ModelInfo.MODEL_KIND_OPENVINO,
+                modelexecpb.ModelInfo.MODEL_STORAGE_FORMAT_OPENVINO,
+            )
+        case "onnx":
+            return (
+                modelexecpb.ModelInfo.MODEL_KIND_ONNX,
+                modelexecpb.ModelInfo.MODEL_STORAGE_FORMAT_PROTOBUF,
+            )
+        case "tfl":
+            return (
+                modelexecpb.ModelInfo.MODEL_KIND_TFLITE,
+                modelexecpb.ModelInfo.MODEL_STORAGE_FORMAT_FLATBUFFER,
+            )
+        case _:
+            raise ValueError("invalid model kind: {}".format(kind))
+
+
+class ModelFromFiles(Model):
+    def __init__(
+        self,
+        name: str,
+        kind: Literal["openvino", "onnx", "tfl"],
+        files: Dict[str, TextIO | BinaryIO | Path],
+        **kwargs,
+    ):
+        kind, storage_format = _kind_str_to_kind_format_protos(kind)
+        super().__init__(name=name, kind=kind, storage_format=storage_format, **kwargs)
+        self.files = files
+
+    def collect_assets(self) -> Dict[str, TextIO | BinaryIO | Path]:
+        return self.files

gml/model_utils.py

@@ -0,0 +1,33 @@
+# 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
+
+
+def prepare_ultralytics_yolo(model):
+    """Prepares an ultralytics YOLO model for export.
+
+    Ultralytics YOLO models requires setting `export=True` on some of the torch modules for exporting to work properly.
+    This function handles setting that value on the necessary modules.
+    """
+    if not hasattr(model, "model"):
+        raise ValueError(
+            "input to `prepare_ultralytics_yolo` is not a supported ultralytics yolo model"
+        )
+    if hasattr(model.model, "fuse") and callable(model.model.fuse):
+        model.model.fuse()
+
+    for _, m in model.named_modules():
+        if hasattr(m, "export"):
+            m.export = True