made changes to incorporate an option for adding labels to the docker image

examples/getting-started/full_workflow_build.py

@@ -3,14 +3,14 @@
 import pickle
 import numpy as np
 from typing import Mapping
-from chassisml import ChassisModel # 
-from chassis.builder import DockerBuilder # 
+from chassisml import ChassisModel #
+from chassis.builder import DockerBuilder, BuildOptions  #
 import chassis.guides as guides
 
-# load model # 
-model = pickle.load(guides.DigitsClassifier) 
+# load model #
+model = pickle.load(guides.DigitsClassifier)
 
-# define predict function # 
+# define predict function #
 def predict(input_bytes: Mapping[str, bytes]) -> dict[str, bytes]:
     inputs = np.array(json.loads(input_bytes['input']))
     inference_results = model.predict_proba(inputs)
@@ -25,9 +25,9 @@ def predict(input_bytes: Mapping[str, bytes]) -> dict[str, bytes]:
     return {'results.json': json.dumps(structured_results).encode()}
 
 # create chassis model object, add required dependencies, and define metadata
-chassis_model = ChassisModel(process_fn=predict)                # 
-chassis_model.add_requirements(["scikit-learn", "numpy"])       # 
-chassis_model.metadata.model_name = "Digits Classifier"         # 
+chassis_model = ChassisModel(process_fn=predict)                #
+chassis_model.add_requirements(["scikit-learn", "numpy"])       #
+chassis_model.metadata.model_name = "Digits Classifier"         #
 chassis_model.metadata.model_version = "0.0.1"
 chassis_model.metadata.add_input(
     key="input",
@@ -40,14 +40,15 @@ def predict(input_bytes: Mapping[str, bytes]) -> dict[str, bytes]:
     media_type="application/json",
     max_size="1M",
     description="Top digit prediction and confidence score"
-)    
+)
 
-# test model # 
+# test model #
 results = chassis_model.test(guides.DigitsSampleData)
 print(results)
 
-# build container # 
-builder = DockerBuilder(chassis_model)
+options = BuildOptions(labels={"key1": "value1"})
+# build container #
+builder = DockerBuilder(chassis_model, options)
 start_time = time.time()
 res = builder.build_image(name="my-first-chassis-model", tag="0.0.1", show_logs=True)
 end_time = time.time()

packages/chassisml/src/chassis/builder/buildable.py

@@ -264,10 +264,12 @@ def render_dockerfile(self, options: BuildOptions) -> str:
             apt_package_list = " ".join(self.apt_packages)
             run_apt_get = f"RUN apt-get update && apt-get install -y {apt_package_list} && rm -rf /var/lib/apt/lists/*"
 
+        #   TODO keys here are variables available in template
         return dockerfile_template.render(
             python_version=options.python_version,
             cuda_version=options.cuda_version,
             apt_packages=run_apt_get,
+            labels=options.labels or {},
         )
 
     def _write_additional_files(self, context: BuildContext):

packages/chassisml/src/chassis/builder/options.py

@@ -2,7 +2,7 @@
 
 import dataclasses
 import platform
-from typing import List, Optional, Union
+from typing import List, Optional, Union, Dict
 
 
 @dataclasses.dataclass
@@ -58,6 +58,7 @@ class BuildOptions:
     python_version: str = "3.9"
     cuda_version: Optional[str] = None
     server: str = "omi"
+    labels: Optional[Dict[str, str]] = None
 
 
 DefaultBuildOptions = BuildOptions()

packages/chassisml/src/chassis/builder/templates/Dockerfile

@@ -13,6 +13,11 @@ FROM python:{{ python_version }}-slim-bullseye
 
 WORKDIR /app
 
+{% for key,value in labels.items() -%}
+LABEL {{key}}="{{value}}"
+{% endfor %}
+
+
 # Copy requirements file and pip install.
 COPY requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
@@ -26,4 +31,4 @@ COPY chassis chassis
 # Copy the app and data
 COPY data data
 
-CMD ["python3", "entrypoint.py"]
+CMD ["python3", "entrypoint.py"]