femtocrux 0.3.0__py3-none-any.whl → 0.4.1__py3-none-any.whl

femtocrux/VERSION

@@ -1 +1 @@
-0.3.0
+0.4.1

femtocrux/__init__.py

@@ -1,6 +1,12 @@
-from .client.client import CompilerClient, TFLiteModel, FQIRModel
+from .client.client import CompilerClient, TFLiteModel, FQIRModel, ManagedCompilerClient
 from .version import __version__
 
 # PEP 8 definiton of public API
 # https://peps.python.org/pep-0008/#public-and-internal-interfaces
-__all__ = ["CompilerClient", "TFLiteModel", "FQIRModel", "__version__"]
+__all__ = [
+    "CompilerClient",
+    "TFLiteModel",
+    "FQIRModel",
+    "__version__",
+    "ManagedCompilerClient",
+]

femtocrux/client/client.py

@@ -12,6 +12,7 @@ import queue
 import sys
 import time
 from typing import Any, List, Tuple, Union
+from contextlib import contextmanager
 
 from fmot.fqir import GraphProto
 
@@ -25,7 +26,7 @@ import femtocrux.grpc.compiler_service_pb2_grpc as cs_pb2_grpc
 __docker_registry__ = "ghcr.io"
 
 
-def __get_docker_image_name__() -> str:
+def _get_docker_image_name() -> str:
     """
     Returns the docker image name. For testing, override with the
     FEMTOCRUX_SERVER_IMAGE_NAME environment variable.
@@ -46,11 +47,11 @@ def __get_docker_image_name__() -> str:
         return remote_image_name
 
 
-__docker_image_name__ = __get_docker_image_name__()
+__docker_image_name__ = _get_docker_image_name()
 
 
 # Set up logging
-def __init_logger__():
+def _init_logger():
     """Init a basic logger to stderr."""
     logger = logging.getLogger(__name__)
     logger.setLevel(logging.INFO)
@@ -61,10 +62,10 @@ def __init_logger__():
     return logger
 
 
-logger = __init_logger__()
+logger = _init_logger()
 
 
-def __env_var_to_bool__(varname: str, default: bool = False) -> bool:
+def _env_var_to_bool(varname: str, default: bool = False) -> bool:
     """Parse an environment varaible as a boolean."""
     try:
         value = os.environ[varname]
@@ -90,24 +91,22 @@ class Model:
     :class:`~femtocrux.client.client.TFLiteModel`.
     """
 
-    def __get_message__(self, options: dict = {}) -> cs_pb2.model:
+    def _get_message(self, options: dict = {}) -> cs_pb2.model:
         # Format the options
         options_struct = google.protobuf.struct_pb2.Struct()
         options_struct.update(options)
 
         # Construct the model with IR
-        return cs_pb2.model(
-            **{self.__ir_name__: self.__get_ir__()}, options=options_struct
-        )
+        return cs_pb2.model(**{self._ir_name: self._get_ir()}, options=options_struct)
 
     @property
-    def __ir_name__(self) -> str:
+    def _ir_name(self) -> str:
         """
         Subclass overrides this to tell which 'ir' field is being returned.
         """
         return NotImplementedError("Subclass must override with IR type.")
 
-    def __get_ir__(self) -> Tuple[str, Any]:
+    def _get_ir(self) -> Tuple[str, Any]:
         """
         Subclass overrides this to implement the 'ir' field of the model's
         grpc message.
@@ -135,10 +134,10 @@ class FQIRModel(Model):
     sequence_dim: int = None
 
     @property
-    def __ir_name__(self) -> str:
+    def _ir_name(self) -> str:
         return "fqir"
 
-    def __get_ir__(self) -> Any:
+    def _get_ir(self) -> Any:
         # Serialize FQIR via pickle
         model = pickle.dumps(self.graph_proto)
 
@@ -167,10 +166,10 @@ class TFLiteModel(Model):
     signature_name: str = None
 
     @property
-    def __ir_name__(self) -> str:
+    def _ir_name(self) -> str:
         return "tflite"
 
-    def __get_ir__(self) -> Any:
+    def _get_ir(self) -> Any:
         return cs_pb2.tflite(model=self.flatbuffer, signature_name=self.signature_name)
 
 
@@ -185,31 +184,31 @@ class Simulator:
 
         # Create an event stream fed by a queue
         self.request_queue = queue.SimpleQueue()
-        request_iterator = iter(self.request_queue.get, self.__request_sentinel__)
-        self.response_iterator = client.__simulate__(request_iterator)
+        request_iterator = iter(self.request_queue.get, self._request_sentinel)
+        self.response_iterator = client._simulate(request_iterator)
 
         # Compile the model with the first message
-        model_msg = model.__get_message__(options)
+        model_msg = model._get_message(options)
         simulation_start_msg = cs_pb2.simulation_input(model=model_msg)
-        self.__send_request__(simulation_start_msg)
+        self._send_request(simulation_start_msg)
 
         # Check compilation status
-        self.__get_response__()
+        self._get_response()
 
     def __del__(self):
         """Close any open streams."""
-        self.__send_request__(self.__request_sentinel__)
+        self._send_request(self._request_sentinel)
 
-    def __send_request__(self, msg):
+    def _send_request(self, msg):
         self.request_queue.put(msg)
 
-    def __get_response__(self):
+    def _get_response(self):
         response = next(self.response_iterator)
-        self.client.__check_status__(response.status)
+        self.client._check_status(response.status)
         return response
 
     @property
-    def __request_sentinel__(self) -> Any:
+    def _request_sentinel(self) -> Any:
         """Sentinel value to close the request queue."""
         return None
 
@@ -218,7 +217,7 @@ class Simulator:
         inputs: Union[np.array, Iterable[np.array]],
         quantize_inputs: bool = False,
         dequantize_outputs: bool = False,
-        sim_duration: float = None,
+        input_period: float = None,
     ) -> List[np.array]:
         """
         Simulates the model on the given inputs.
@@ -243,7 +242,8 @@ class Simulator:
         static power consumption. For example, in a streaming model, this should be the
         time elapsed between model invocations. By default, this is just the estimated
         latency of the model.
-        :type sim_duration: float, optional
+        :type input_period (float, optional): Duration between each input in a sequence,
+            in seconds.
 
         :rtype: list
         :return: The output tensors.
@@ -260,11 +260,11 @@ class Simulator:
                 data=[numpy_to_ndarray(x) for x in inputs],
                 quantize_inputs=quantize_inputs,
                 dequantize_outputs=dequantize_outputs,
-                sim_duration=sim_duration,
+                input_period=input_period,
             )
         )
-        self.__send_request__(simulation_request)
-        response = self.__get_response__()
+        self._send_request(simulation_request)
+        response = self._get_response()
 
         return [ndarray_to_numpy(x) for x in response.data], response.report
 
@@ -279,21 +279,21 @@ class CompilerClientImpl:
     def __init__(self, channel, stub):
         self.channel = channel
         self.stub = stub
-        self.__check_version__()
+        self._check_version()
 
-    def __check_status__(self, status):
+    def _check_status(self, status):
         """Check a status response, raising an exception if unsuccessful."""
         if not status.success:
             raise RuntimeError(
                 "Client received error from compiler server:\n%s" % status.msg
             )
 
-    def __check_version__(self):
+    def _check_version(self):
         """Verify the server's version matches the client."""
 
         from femtocrux.version import __version__ as client_version
 
-        server_version = self.__server_version__()
+        server_version = self._server_version()
         assert (
             client_version == server_version
         ), """
@@ -319,17 +319,17 @@ class CompilerClientImpl:
         :return: A zip archive of compiler artifacts.
         """
 
-        response = self.stub.compile(model.__get_message__(options))
-        self.__check_status__(response.status)
+        response = self.stub.compile(model._get_message(options))
+        self._check_status(response.status)
         return response.bitfile
 
-    def __ping__(self, message: bytes) -> None:
+    def _ping(self, message: bytes) -> None:
         """Pings the server with a message."""
         response = self.stub.ping(cs_pb2.data(data=message))
         if response.data != message:
             raise RuntimeError("Server response does not match request data!")
 
-    def __simulate__(self, in_stream: Iterable) -> Iterable:
+    def _simulate(self, in_stream: Iterable) -> Iterable:
         """Calls the 'simulator' bidirectional streaming RPC."""
         return self.stub.simulate(in_stream)
 
@@ -348,7 +348,7 @@ class CompilerClientImpl:
         """
         return Simulator(client=self, model=model, options=options)
 
-    def __server_version__(self) -> str:
+    def _server_version(self) -> str:
         """Queries the femtocrux version running on the server."""
         response = self.stub.version(google.protobuf.empty_pb2.Empty())
         return response.version
@@ -373,24 +373,49 @@ class CompilerClient(CompilerClientImpl):
         self.container = None  # For __del__
 
         # Start a new docker server
-        self.container = self.__create_docker_server__(docker_kwargs)
-        self.__wait_for_server_ready__()
-        self.__init_network_info__(self.container)
+        self.container = self._create_docker_server(docker_kwargs)
+        self._wait_for_server_ready()
+        self._init_network_info(self.container)
 
         # Establish a connection to the server
-        self.channel = self.__connect__()
+        self.channel = self._connect()
 
         # Initialize the client on this channel
         self.stub = cs_pb2_grpc.CompileStub(self.channel)
         super().__init__(self.channel, self.stub)
 
+    @property
+    def status(self):
+        if self.container is not None:
+            return self.container.status
+        else:
+            return "exited"
+
+    @property
+    def name(self):
+        if self.container is not None:
+            return self.container.name
+        else:
+            return None
+
+    def close(self):
+        if self.container is not None:
+            try:
+                self.container.stop()
+            except Exception as e:
+                logger.info(f"Image already closed... skipping close\n{e}")
+
     def __del__(self):
         """Reclaim system resources."""
         if self.container is not None:
-            self.container.kill()
-            self.container = None
+            cli = docker.DockerClient()
+            try:
+                container = cli.containers.get(self.name)
+                container.stop()
+            except docker.errors.NotFound:
+                pass
 
-    def __get_docker_api_client__(self):
+    def _get_docker_api_client(self):
         """Get a client to the Docker daemon."""
         try:
             return docker.from_env()
@@ -400,7 +425,7 @@ class CompilerClient(CompilerClientImpl):
                     Please ensure it is installed and running."""
             ) from exc
 
-    def __init_network_info__(self, container):
+    def _init_network_info(self, container):
         """
         For local connections only.
 
@@ -424,7 +449,7 @@ class CompilerClient(CompilerClientImpl):
         socket = bound_sockets[0]  # In case of multiple, take the first one
         self.__channel_port__ = socket["HostPort"]
 
-    def __connect__(self) -> Any:
+    def _connect(self) -> Any:
         """Establishes a gRPC connection to the server."""
 
         # Open a gRPC channel to the server
@@ -467,20 +492,20 @@ class CompilerClient(CompilerClientImpl):
         return self.__channel_port__
 
     @property
-    def __container_port__(self) -> int:
+    def _container_port(self) -> int:
         """Port used inside the container."""
         return 50051
 
     @property
-    def __container_label__(self) -> str:
+    def _container_label(self) -> str:
         """Label attached to identify containers started by this client."""
         return "femtocrux_server"
 
-    def __get_unused_container_name__(self) -> str:
+    def _get_unused_container_name(self) -> str:
         """Get an unused container name."""
 
         # Search for an unused name
-        client = self.__get_docker_api_client__()
+        client = self._get_docker_api_client()
         container_idx = 0
         while True:
             name = "femtocrux_server_%d" % container_idx
@@ -492,7 +517,7 @@ class CompilerClient(CompilerClientImpl):
 
             container_idx += 1
 
-    def __pull_docker_image__(self):
+    def _pull_docker_image(self):
         """Pull the Docker image from remote."""
 
         logger.info(
@@ -506,7 +531,7 @@ class CompilerClient(CompilerClientImpl):
         )
 
         # Log in to Github
-        client = self.__get_docker_api_client__()
+        client = self._get_docker_api_client()
         while True:
             # Get the password
             manual_pass = True
@@ -561,7 +586,7 @@ class CompilerClient(CompilerClientImpl):
 
         logger.info("Download completed.")
 
-    def __create_docker_server__(
+    def _create_docker_server(
         self, docker_kwargs: dict[str, Any] = None
     ) -> docker.models.containers.Container:
         """
@@ -571,7 +596,7 @@ class CompilerClient(CompilerClientImpl):
             docker_kwargs = {}
 
         # Get a client for the Docker daemon
-        client = self.__get_docker_api_client__()
+        client = self._get_docker_api_client()
 
         # Pull the image, if not available
         existing_image_names = [
@@ -583,7 +608,7 @@ class CompilerClient(CompilerClientImpl):
             image_not_found_msg = (
                 "Failed to find the docker image %s locally." % __docker_image_name__
             )
-            if not __env_var_to_bool__("FS_ALLOW_DOCKER_PULL", default=True):
+            if not _env_var_to_bool("FS_ALLOW_DOCKER_PULL", default=True):
                 raise RuntimeError(
                     """
                     %s
@@ -594,29 +619,29 @@ class CompilerClient(CompilerClientImpl):
 
             # Pull the image from remote
             logger.info(image_not_found_msg)
-            self.__pull_docker_image__()
+            self._pull_docker_image()
 
         # Bind a random port on the host to the container's gRPC port
-        port_interface = {self.__container_port__: None}
-        command = "--port %s" % self.__container_port__
+        port_interface = {self._container_port: None}
+        command = "--port %s" % self._container_port
 
         # Start a container running the server
         container = client.containers.run(
             __docker_image_name__,
             command=command,  # Appends entrypoint with args
             detach=True,
-            labels=[self.__container_label__],
+            labels=[self._container_label],
             stderr=True,
             stdout=True,
             ports=port_interface,
-            name=self.__get_unused_container_name__(),
+            name=self._get_unused_container_name(),
             auto_remove=True,
             **docker_kwargs,
         )
 
         return container
 
-    def __wait_for_server_ready__(self):
+    def _wait_for_server_ready(self):
         """
         Block until the Docker container is ready to handle requests.
         """
@@ -670,6 +695,15 @@ class CompilerClient(CompilerClientImpl):
             )
 
 
+@contextmanager
+def ManagedCompilerClient(docker_kwargs: dict[str, Any] = None) -> CompilerClient:
+    client = CompilerClient(docker_kwargs=docker_kwargs)
+    try:
+        yield client
+    finally:
+        client.close()
+
+
 if __name__ == "__main__":
     logging.basicConfig(stream=sys.stdout, level=logging.INFO)
     client = CompilerClient()

femtocrux/grpc/compiler_service_pb2.py

@@ -15,7 +15,7 @@ from google.protobuf import empty_pb2 as google_dot_protobuf_dot_empty__pb2
 from google.protobuf import struct_pb2 as google_dot_protobuf_dot_struct__pb2
 
 
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x16\x63ompiler_service.proto\x12\nfscompiler\x1a\x1bgoogle/protobuf/empty.proto\x1a\x1cgoogle/protobuf/struct.proto\"g\n\x04\x66qir\x12\r\n\x05model\x18\x01 \x01(\x0c\x12\x16\n\tbatch_dim\x18\x02 \x01(\x03H\x00\x88\x01\x01\x12\x19\n\x0csequence_dim\x18\x03 \x01(\x03H\x01\x88\x01\x01\x42\x0c\n\n_batch_dimB\x0f\n\r_sequence_dim\"G\n\x06tflite\x12\r\n\x05model\x18\x01 \x01(\x0c\x12\x1b\n\x0esignature_name\x18\x02 \x01(\tH\x00\x88\x01\x01\x42\x11\n\x0f_signature_name\"\x90\x01\n\x05model\x12 \n\x04\x66qir\x18\x01 \x01(\x0b\x32\x10.fscompiler.fqirH\x00\x12$\n\x06tflite\x18\x02 \x01(\x0b\x32\x12.fscompiler.tfliteH\x00\x12-\n\x07options\x18\x03 \x01(\x0b\x32\x17.google.protobuf.StructH\x01\x88\x01\x01\x42\x04\n\x02irB\n\n\x08_options\"&\n\x06status\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0b\n\x03msg\x18\x02 \x01(\t\"I\n\x12\x63ompiled_artifacts\x12\x0f\n\x07\x62itfile\x18\x01 \x01(\x0c\x12\"\n\x06status\x18\x02 \x01(\x0b\x32\x12.fscompiler.status\"&\n\x07ndarray\x12\x0c\n\x04\x64\x61ta\x18\x01 \x03(\x02\x12\r\n\x05shape\x18\x02 \x03(\x03\"\x14\n\x04\x64\x61ta\x12\x0c\n\x04\x64\x61ta\x18\x01 \x01(\x0c\"\x95\x01\n\x0fsimulation_data\x12!\n\x04\x64\x61ta\x18\x01 \x03(\x0b\x32\x13.fscompiler.ndarray\x12\x17\n\x0fquantize_inputs\x18\x02 \x01(\x08\x12\x1a\n\x12\x64\x65quantize_outputs\x18\x03 \x01(\x08\x12\x19\n\x0csim_duration\x18\x04 \x01(\x02H\x00\x88\x01\x01\x42\x0f\n\r_sim_duration\"t\n\x10simulation_input\x12\"\n\x05model\x18\x01 \x01(\x0b\x32\x11.fscompiler.modelH\x00\x12+\n\x04\x64\x61ta\x18\x02 \x01(\x0b\x32\x1b.fscompiler.simulation_dataH\x00\x42\x0f\n\rmodel_or_data\"j\n\x11simulation_output\x12!\n\x04\x64\x61ta\x18\x01 \x03(\x0b\x32\x13.fscompiler.ndarray\x12\x0e\n\x06report\x18\x02 \x01(\t\x12\"\n\x06status\x18\x03 \x01(\x0b\x32\x12.fscompiler.status\"\x1f\n\x0cversion_info\x12\x0f\n\x07version\x18\x01 \x01(\t2\x85\x02\n\x07\x43ompile\x12>\n\x07\x63ompile\x12\x11.fscompiler.model\x1a\x1e.fscompiler.compiled_artifacts\"\x00\x12,\n\x04ping\x12\x10.fscompiler.data\x1a\x10.fscompiler.data\"\x00\x12M\n\x08simulate\x12\x1c.fscompiler.simulation_input\x1a\x1d.fscompiler.simulation_output\"\x00(\x01\x30\x01\x12=\n\x07version\x12\x16.google.protobuf.Empty\x1a\x18.fscompiler.version_info\"\x00\x62\x06proto3')
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x16\x63ompiler_service.proto\x12\nfscompiler\x1a\x1bgoogle/protobuf/empty.proto\x1a\x1cgoogle/protobuf/struct.proto\"g\n\x04\x66qir\x12\r\n\x05model\x18\x01 \x01(\x0c\x12\x16\n\tbatch_dim\x18\x02 \x01(\x03H\x00\x88\x01\x01\x12\x19\n\x0csequence_dim\x18\x03 \x01(\x03H\x01\x88\x01\x01\x42\x0c\n\n_batch_dimB\x0f\n\r_sequence_dim\"G\n\x06tflite\x12\r\n\x05model\x18\x01 \x01(\x0c\x12\x1b\n\x0esignature_name\x18\x02 \x01(\tH\x00\x88\x01\x01\x42\x11\n\x0f_signature_name\"\x90\x01\n\x05model\x12 \n\x04\x66qir\x18\x01 \x01(\x0b\x32\x10.fscompiler.fqirH\x00\x12$\n\x06tflite\x18\x02 \x01(\x0b\x32\x12.fscompiler.tfliteH\x00\x12-\n\x07options\x18\x03 \x01(\x0b\x32\x17.google.protobuf.StructH\x01\x88\x01\x01\x42\x04\n\x02irB\n\n\x08_options\"&\n\x06status\x12\x0f\n\x07success\x18\x01 \x01(\x08\x12\x0b\n\x03msg\x18\x02 \x01(\t\"I\n\x12\x63ompiled_artifacts\x12\x0f\n\x07\x62itfile\x18\x01 \x01(\x0c\x12\"\n\x06status\x18\x02 \x01(\x0b\x32\x12.fscompiler.status\"&\n\x07ndarray\x12\x0c\n\x04\x64\x61ta\x18\x01 \x03(\x02\x12\r\n\x05shape\x18\x02 \x03(\x03\"\x14\n\x04\x64\x61ta\x12\x0c\n\x04\x64\x61ta\x18\x01 \x01(\x0c\"\xc1\x01\n\x0fsimulation_data\x12!\n\x04\x64\x61ta\x18\x01 \x03(\x0b\x32\x13.fscompiler.ndarray\x12\x17\n\x0fquantize_inputs\x18\x02 \x01(\x08\x12\x1a\n\x12\x64\x65quantize_outputs\x18\x03 \x01(\x08\x12\x19\n\x0csim_duration\x18\x04 \x01(\x02H\x00\x88\x01\x01\x12\x19\n\x0cinput_period\x18\x05 \x01(\x02H\x01\x88\x01\x01\x42\x0f\n\r_sim_durationB\x0f\n\r_input_period\"t\n\x10simulation_input\x12\"\n\x05model\x18\x01 \x01(\x0b\x32\x11.fscompiler.modelH\x00\x12+\n\x04\x64\x61ta\x18\x02 \x01(\x0b\x32\x1b.fscompiler.simulation_dataH\x00\x42\x0f\n\rmodel_or_data\"j\n\x11simulation_output\x12!\n\x04\x64\x61ta\x18\x01 \x03(\x0b\x32\x13.fscompiler.ndarray\x12\x0e\n\x06report\x18\x02 \x01(\t\x12\"\n\x06status\x18\x03 \x01(\x0b\x32\x12.fscompiler.status\"\x1f\n\x0cversion_info\x12\x0f\n\x07version\x18\x01 \x01(\t2\x85\x02\n\x07\x43ompile\x12>\n\x07\x63ompile\x12\x11.fscompiler.model\x1a\x1e.fscompiler.compiled_artifacts\"\x00\x12,\n\x04ping\x12\x10.fscompiler.data\x1a\x10.fscompiler.data\"\x00\x12M\n\x08simulate\x12\x1c.fscompiler.simulation_input\x1a\x1d.fscompiler.simulation_output\"\x00(\x01\x30\x01\x12=\n\x07version\x12\x16.google.protobuf.Empty\x1a\x18.fscompiler.version_info\"\x00\x62\x06proto3')
 
 _globals = globals()
 _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
@@ -37,13 +37,13 @@ if _descriptor._USE_C_DESCRIPTORS == False:
   _globals['_DATA']._serialized_start=577
   _globals['_DATA']._serialized_end=597
   _globals['_SIMULATION_DATA']._serialized_start=600
-  _globals['_SIMULATION_DATA']._serialized_end=749
-  _globals['_SIMULATION_INPUT']._serialized_start=751
-  _globals['_SIMULATION_INPUT']._serialized_end=867
-  _globals['_SIMULATION_OUTPUT']._serialized_start=869
-  _globals['_SIMULATION_OUTPUT']._serialized_end=975
-  _globals['_VERSION_INFO']._serialized_start=977
-  _globals['_VERSION_INFO']._serialized_end=1008
-  _globals['_COMPILE']._serialized_start=1011
-  _globals['_COMPILE']._serialized_end=1272
+  _globals['_SIMULATION_DATA']._serialized_end=793
+  _globals['_SIMULATION_INPUT']._serialized_start=795
+  _globals['_SIMULATION_INPUT']._serialized_end=911
+  _globals['_SIMULATION_OUTPUT']._serialized_start=913
+  _globals['_SIMULATION_OUTPUT']._serialized_end=1019
+  _globals['_VERSION_INFO']._serialized_start=1021
+  _globals['_VERSION_INFO']._serialized_end=1052
+  _globals['_COMPILE']._serialized_start=1055
+  _globals['_COMPILE']._serialized_end=1316
 # @@protoc_insertion_point(module_scope)

femtocrux/server/server.py

@@ -1,7 +1,7 @@
 import argparse
 from collections.abc import Iterable
 import concurrent
-import google.protobuf
+import google.protobuf.json_format
 import grpc
 import logging
 import pickle
@@ -36,7 +36,7 @@ class CompileServicer(cs_pb2_grpc.CompileServicer):
         self.logger.setLevel(logging.DEBUG)
         self.logger.info("Starting compile server.")
 
-    def __get_fqir_compiler__(self, model: cs_pb2.model) -> CompilerFrontend:
+    def _get_fqir_compiler(self, model: cs_pb2.model) -> CompilerFrontend:
         """Get a Torch compiler from an FQIR a model message"""
         # Deserialize FQIR
         fqir = model.fqir
@@ -49,19 +49,19 @@ class CompileServicer(cs_pb2_grpc.CompileServicer):
             seq_dim=field_or_none(fqir, "sequence_dim"),
         )
 
-    def __get_tflite_compiler__(self, model: cs_pb2.model) -> CompilerFrontend:
+    def _get_tflite_compiler(self, model: cs_pb2.model) -> CompilerFrontend:
         """Get a TFLite compiler from a model message"""
         tflite = model.tflite
         return TFLiteCompiler(
             tflite.model, signature=field_or_none(tflite, "signature_name")
         )
 
-    def __compile_model__(self, model: cs_pb2.model, context) -> CompilerFrontend:
+    def _compile_model(self, model: cs_pb2.model, context) -> CompilerFrontend:
         """Compile a model, for simulation or bitfile generation."""
         # Get a compiler for the model
         model_type_map = {
-            "fqir": self.__get_fqir_compiler__,
-            "tflite": self.__get_tflite_compiler__,
+            "fqir": self._get_fqir_compiler,
+            "tflite": self._get_tflite_compiler,
         }
         model_type = model.WhichOneof("ir")
         compiler = model_type_map[model_type](model)
@@ -85,7 +85,7 @@ class CompileServicer(cs_pb2_grpc.CompileServicer):
 
         # Compile the model
         try:
-            compiler = self.__compile_model__(model, context)
+            compiler = self._compile_model(model, context)
             bitfile = compiler.dump_bitfile(encrypt=True)
         except Exception as exc:
             msg = "Compiler raised exception:\n%s" % (format_exception_from_exc(exc))
@@ -120,7 +120,7 @@ class CompileServicer(cs_pb2_grpc.CompileServicer):
 
             # Attempt compilation
             try:
-                compiler = self.__compile_model__(model_request.model, context)
+                compiler = self._compile_model(model_request.model, context)
             except Exception as exc:
                 msg = "Compiler raised exception:\n%s" % (
                     format_exception_from_exc(exc)
@@ -158,7 +158,7 @@ class CompileServicer(cs_pb2_grpc.CompileServicer):
                     data,
                     quantize_inputs=sim_data.quantize_inputs,
                     dequantize_outputs=sim_data.dequantize_outputs,
-                    input_period=field_or_none(sim_data, "sim_duration"),
+                    input_period=field_or_none(sim_data, "input_period"),
                 )
             except Exception as exc:
                 msg = "Simulator raised exception:\n%s" % (

{femtocrux-0.3.0.dist-info → femtocrux-0.4.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: femtocrux
-Version: 0.3.0
+Version: 0.4.1
 Summary: Femtosense Compiler
 Home-page: https://github.com/femtosense/femtocrux
 Author: Femtosense
@@ -37,7 +37,6 @@ The package itself is a thin Python wrapper communicating to the actual compiler
 
 Supported ML model representations:
  - Femtosense Quantized IR (FQIR)
- - Tensorflow Lite (TFLite)
 
 ## Installation
 
@@ -52,16 +51,51 @@ pip install femtocrux
 
 ## Basic Usage
 
+### `ManagedCompilerClient`
+`ManagedCompilerClient` is the recommended way to use femtocrux within a context manager.
+
+```python
+from femtocrux import ManagedCompilerClient, FQIRModel
+
+# fqir_graph = ... # Assuming we have an FQIR graph using fmot
+# inputs = ...     # Assuming we have numpy array inputs to the model
+
+fqir_model = FQIRModel(fqir_graph, batch_dim=0, sequence_dim=1)
+with ManagedCompilerClient() as client:
+    
+    # simulate execution, view power, energy, and latency metrics
+    simulator = client.simulate(fqir_model)
+    outputs, metrics = simulator.simulate(inputs)
+
+    # compile the model to a bitfile
+    bitstream = client.compile(fqir_model)
+    with open('my_bitfile.zip', 'wb') as f: 
+        f.write(bitstream)
+```
+
+The bitfile can be used to generate program files to run on the SPU, using [femtodriverpub](https://github.com/femtosense/femtodriverpub.git)
+
+### `CompilerClient`
+`CompilerClient` is another interface to the compiler. This will be deprecated in future releases, so we recommend using `ManagedCompilerClient`.
+
 ```python
-# fqir_model = ... # Assuming we have an FQIRModel
-# inputs = ... # Assuming we have inputs for the model
+from femtocrux import CompilerClient, FQIRModel
 
-from femtocrux import CompilerClient
+# fqir_graph = ... # Assuming we have an FQIR graph using fmot
+# inputs = ...     # Assuming we have numpy array inputs to the model
+
+fqir_model = FQIRModel(fqir_graph, batch_dim=0, sequence_dim=1)
 client = CompilerClient()
+    
+# simulate execution, view power, energy, and latency metrics
 simulator = client.simulate(fqir_model)
-outputs = simulator.simulate(inputs)
-```
+outputs, metrics = simulator.simulate(inputs)
 
-## APIs
+# compile the model to a bitfile
+bitstream = client.compile(fqir_model)
+with open('my_bitfile.zip', 'wb') as f: 
+    f.write(bitstream)
 
-Coming soon...
+# close the client (important to avoid background docker containers that continue running)
+client.close()
+```

{femtocrux-0.3.0.dist-info → femtocrux-0.4.1.dist-info}/RECORD

@@ -1,21 +1,21 @@
 femtocrux/ENV_REQUIREMENTS.sh,sha256=t_O1B4hJAMgxvH9gwp1qls6eVFmhSYBJe64KmuK_H-4,1389
 femtocrux/PY_REQUIREMENTS,sha256=HZOIodZiEjQF6kbKT1PxBU96fU8W7fypimUfIeJaoAQ,303
-femtocrux/VERSION,sha256=2RXMldbKj0euKXcT7UbU5cXZnd0p_Dxh4mO98wXytbA,6
-femtocrux/__init__.py,sha256=fQdrj_d4a2T4xIcKoTdTLokaZsDP5uAbW7c39YPzkX8,271
+femtocrux/VERSION,sha256=9iGEzuh4fy9pQcggQaVyXU7cmqKT6-Xb9mRAboLsH-E,6
+femtocrux/__init__.py,sha256=yIWd9I2PEXCn_PKIILAN3mkWeTf0tgtVualeTIHNxfQ,342
 femtocrux/version.py,sha256=uNg2kHxQo6oUN1ah7s9_85rCZVRoTHGPD1GAQPZW4lw,164
 femtocrux/client/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-femtocrux/client/client.py,sha256=VP6KsESnzUAI49tAsEFGzmWEK0_cnv0e4Mm2izOkOOU,21910
+femtocrux/client/client.py,sha256=PNaXL_XTe_7_r14WWKmp1RnvAW-kvYzDXcimseA8sEo,22726
 femtocrux/grpc/__init__.py,sha256=uiMHQt5I2eAKJqI3Zh0h1Gm7cmPR4PbaGS71nCJQCGw,169
-femtocrux/grpc/compiler_service_pb2.py,sha256=XaawNGpFBdURP5JJScnzfA8JY-_J18vpVJ-Blzl5WF0,4339
+femtocrux/grpc/compiler_service_pb2.py,sha256=RUBcT-S2sKr5kveHQQ7EMyAe5rGU_8JZVbhKAG693Fg,4424
 femtocrux/grpc/compiler_service_pb2_grpc.py,sha256=L9EQFYMFVACngUrAEmSJDgUTSPSi3rDoG7TERg-pFq4,7002
 femtocrux/server/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 femtocrux/server/exceptions.py,sha256=lI6n471n5QKf5G3aL_1kuBVEItD-jBgithVVpPDwNYc,609
 femtocrux/server/healthcheck.py,sha256=ehqAwnv0D0zpy-AUZAPwv8rp874DZCwUmP8nzdXzZvI,1565
-femtocrux/server/server.py,sha256=KAH6paHEQrdw9PIv7GxSc_O2BBAKQHys3s8Ez5NkY1I,7928
+femtocrux/server/server.py,sha256=tmXVleZQB59oFzdmut3na4NnDvr0gmxphXF3N3MQx6I,7919
 femtocrux/util/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 femtocrux/util/utils.py,sha256=FZ8cssDom4B3FDbVU_ew4Cf3wOWjo2w1jwcbnLzoYnM,1003
-femtocrux-0.3.0.dist-info/LICENSE,sha256=eN9ZI1xHjUmFvN3TEeop5kBGXRUBfbsl55KBNBYYFqI,36
-femtocrux-0.3.0.dist-info/METADATA,sha256=XiRTQKnAplRmDKUGlWjh9XlhbqiyyKFTuXAoZI2Auzg,2109
-femtocrux-0.3.0.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
-femtocrux-0.3.0.dist-info/top_level.txt,sha256=BkTttlioC3je__8577wxRieZqY3Abu7FOOdMnmYbcNI,10
-femtocrux-0.3.0.dist-info/RECORD,,
+femtocrux-0.4.1.dist-info/LICENSE,sha256=eN9ZI1xHjUmFvN3TEeop5kBGXRUBfbsl55KBNBYYFqI,36
+femtocrux-0.4.1.dist-info/METADATA,sha256=JjAykujLMZmW7rGL2Rx9HfmKkrk8UM8D6VH7S92ueZQ,3536
+femtocrux-0.4.1.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+femtocrux-0.4.1.dist-info/top_level.txt,sha256=BkTttlioC3je__8577wxRieZqY3Abu7FOOdMnmYbcNI,10
+femtocrux-0.4.1.dist-info/RECORD,,

{femtocrux-0.3.0.dist-info → femtocrux-0.4.1.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.41.2)
+Generator: bdist_wheel (0.42.0)
 Root-Is-Purelib: true
 Tag: py3-none-any