Merge pull request #30 from Prismadic/docker

🧲 feat(release): Docker image

Author: mxchinegod

docker/Dockerfile

@@ -0,0 +1,13 @@
+FROM python:3.11
+RUN python3.11 -m ensurepip 
+RUN python3.11 -m pip install torch sentence-transformers pandas tqdm spacy pymilvus boto3 nats-py rich xxhash nkeys -U
+
+WORKDIR /usr/src/app
+
+# Install the required package
+RUN pip install --no-cache-dir llm-magnet -U
+
+COPY ./scripts/ ./
+
+ENTRYPOINT ["sh", "-c"]
+CMD ["python ./main.py"]

docker/README.md

@@ -0,0 +1,9 @@
+# 🧲 docker
+
+> [!NOTE]
+> Each runtime/worker has a `MAGNET_ROLE` environment variable.
+
+Current role abstractions:
+
+- "index"
+    - Stream data from a `field`, encode using an embedding model, and upload to an index.

docker/docker-compose.yaml

@@ -0,0 +1,18 @@
+version: '3.8'
+services:
+  magnet-index:
+    image: magnet
+    environment:
+      MILVUS_HOST: ""
+      MILVUS_PORT: "19530"
+      MILVUS_USER: ""
+      MILVUS_PASSWORD: ""
+      DIMENSION: ""
+      MODEL: ""
+      INDEX: ""
+      efConstruction: ""
+      M: ""
+      RESONATOR_URI: ""
+      CATEGORY: ""
+      SESSION: ""
+      STREAM: ""

docker/scripts/main.py

@@ -0,0 +1,13 @@
+import os, asyncio
+from roles import index
+from magnet.utils.globals import _f
+magnet_role = os.getenv('MAGNET_ROLE')
+
+def main():
+    if magnet_role == 'index':
+        asyncio.run(index.main())
+    else:
+        _f('fatal', 'No MAGNET_ROLE set! 😭')
+
+if __name__ == "__main__":
+    main()

docker/scripts/roles/index.py

@@ -0,0 +1,36 @@
+import os
+from magnet.ize import memory
+from magnet.ic import field
+from magnet.utils.globals import _f
+
+config = {
+    "MILVUS_URI": os.getenv('MILVUS_HOST'),
+    "MILVUS_PORT": int(os.getenv('MILVUS_PORT', 19530)),
+    "MILVUS_USER": os.getenv('MILVUS_USER'),
+    "MILVUS_PASSWORD": os.getenv('MILVUS_PASSWORD'),
+    "DIMENSION": int(os.getenv('DIMENSION', 1024)),
+    "MODEL": os.getenv('MODEL'),
+    "INDEX": os.getenv('INDEX'),
+    "INDEX_PARAMS": {
+        'metric_type': 'COSINE',
+        'index_type':'HNSW',
+        'params': {
+            "efConstruction": int(os.getenv('efConstruction', 40)),
+            "M": int(os.getenv('M', 48))
+        },
+    }
+}
+
+embedder = memory.Embedder(config)
+
+async def handle_payload(payload, msg):
+    await embedder.index(payload, msg, verbose=True)
+
+async def main():
+    try:
+        reso = field.Resonator(os.getenv('RESONATOR_URI'))
+    except Exception as e:
+        _f('warn', e)
+        main()
+    await reso.on(category=os.getenv('CATEGORY'), session=os.getenv('SESSION'), stream=os.getenv('STREAM'))
+    await reso.listen(cb=handle_payload)