Merge pull request #468 from mystic-ai/matthew/turbo-registry

Turbo Registry Integration

examples/text-to-text/twitter-roberta-base-sentiment/pipeline.yaml

@@ -15,4 +15,4 @@ pipeline_graph: new_pipeline:my_new_pipeline
 pipeline_name: twitter-roberta-base-sentiment
 description: null
 readme: null
-extras: {}
+extras: {}

pipeline/console/container/build.py

@@ -25,6 +25,9 @@ def build_container(namespace: Namespace):
 
     pipeline_config = PipelineConfig.parse_obj(pipeline_config_yaml)
 
+    if pipeline_config.extras and pipeline_config.extras.get("turbo_registry", False):
+        template = docker_templates.turbo_registry_template
+
     if not pipeline_config.runtime:
         raise ValueError("No runtime config found")
     if not pipeline_config.runtime.python:

pipeline/console/container/push.py

@@ -86,7 +86,11 @@ def push_container(namespace: Namespace):
 
     docker_client: docker.DockerClient = docker.from_env(timeout=600)
 
-    registry_info = http.get(endpoint="/v4/registry")
+    registry_params = {}
+    if pipeline_config.extras and pipeline_config.extras.get("turbo_registry", False):
+        registry_params["turbo_registry"] = "yes"
+
+    registry_info = http.get(endpoint="/v4/registry", params=registry_params)
     registry_info = registry_schemas.RegistryInformation.parse_raw(registry_info.text)
 
     upload_registry = registry_info.url

pipeline/container/docker_templates/__init__.py

@@ -13,3 +13,4 @@ def _read_template(template_name: str) -> str:
 dockerfile_template = _read_template("dockerfile_template.txt")
 pipeline_template_python = _read_template("pipeline_template.py")
 readme_template = _read_template("readme_template.md")
+turbo_registry_template = _read_template("turbo_registry_template.txt")

pipeline/container/docker_templates/turbo_registry_template.txt

@@ -0,0 +1,32 @@
+# This file is automatically generated by pipeline-ai and should not be changed by hand.
+FROM python:{python_version}-slim
+
+WORKDIR /app
+
+RUN apt update -y
+RUN pip install --no-cache-dir -U pip setuptools wheel
+
+# Install serving packages
+RUN pip install --no-cache-dir -U fastapi==0.105.0 uvicorn==0.25.0 \
+    python-multipart==0.0.6 loguru==0.7.2
+
+# Container commands
+{container_commands}
+
+# Install python dependencies
+RUN pip install --no-cache-dir {python_requirements}
+
+# Copy in files
+COPY ./ ./
+
+FROM scratch
+COPY --from=0 / /
+
+WORKDIR /app
+
+ENV PYTHONUNBUFFERED=1
+ENV PIPELINE_PATH={pipeline_path}
+ENV PIPELINE_NAME={pipeline_name}
+ENV PIPELINE_IMAGE={pipeline_image}
+
+CMD ["uvicorn", "pipeline.container.startup:create_app", "--host", "0.0.0.0", "--port", "14300"]

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pipeline-ai"
-version = "2.3.2"
+version = "2.4.0"
 description = "Pipelines for machine learning workloads."
 authors = [
   "Paul Hetherington <ph@mystic.ai>",