change worker to docker proxy

containers/main.py

@@ -109,10 +109,8 @@
     "d_prime",
 ]
 
-sparsity_params = dict(method="radius", radius_um=100)
-
 postprocessing_params = dict(
-    sparsity=sparsity_params,
+    sparsity=dict(method="radius", radius_um=100),
     waveforms_deduplicate=dict(
         ms_before=0.5,
         ms_after=1.5,
@@ -148,7 +146,11 @@
     locations=dict(method="monopolar_triangulation"),
     template_metrics=dict(upsampling_factor=10, sparsity=None),
     principal_components=dict(n_components=5, mode="by_channel_local", whiten=True),
-    quality_metrics=dict(qm_params=qm_params, metric_names=qm_metric_names, n_jobs=1),
+    quality_metrics=dict(
+        qm_params=qm_params, 
+        metric_names=qm_metric_names, 
+        n_jobs=1
+    ),
 )
 
 curation_params = dict(

docker-compose-dev.yml

@@ -1,28 +1,38 @@
 version: "3"
 
 services:
-  frontend:
-    build:
-      context: frontend
-      dockerfile: Dockerfile
-    image: si-sorting-frontend
-    container_name: si-sorting-frontend
-    command: ["npm", "run", "start"]
+  docker-proxy:
+    image: bobrik/socat
+    container_name: si-docker-proxy
+    command: "TCP4-LISTEN:2375,fork,reuseaddr UNIX-CONNECT:/var/run/docker.sock"
     ports:
-      - "5173:5173"
-    environment:
-      DEPLOY_MODE: compose
+      - "2376:2375"
     volumes:
-      - ./frontend:/app
-    depends_on:
-      - rest
+      - /var/run/docker.sock:/var/run/docker.sock
+
+  # frontend:
+  #   build:
+  #     context: frontend
+  #     dockerfile: Dockerfile
+  #   image: si-sorting-frontend
+  #   container_name: si-sorting-frontend
+  #   command: ["npm", "run", "start"]
+  #   ports:
+  #     - "5173:5173"
+  #   environment:
+  #     DEPLOY_MODE: compose
+  #   volumes:
+  #     - ./frontend:/app
+  #   depends_on:
+  #     - rest
 
   rest:
     build:
       context: rest
       dockerfile: Dockerfile
     image: si-sorting-rest
     container_name: si-sorting-rest
+    command: ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000", "--reload", "--reload-dir", "/app"]
     ports:
       - "8000:8000"
     environment:
@@ -39,33 +49,33 @@ services:
     depends_on:
       - database
 
-  worker:
-    build:
-      context: containers
-      dockerfile: Dockerfile.combined
-    image: si-sorting-worker
-    # image: ghcr.io/catalystneuro/si-sorting-worker:latest
-    container_name: si-sorting-worker
-    ports:
-      - "5000:5000"
-    environment:
-      WORKER_DEPLOY_MODE: compose
-      AWS_DEFAULT_REGION: ${AWS_DEFAULT_REGION}
-      AWS_ACCESS_KEY_ID: ${AWS_ACCESS_KEY_ID}
-      AWS_SECRET_ACCESS_KEY: ${AWS_SECRET_ACCESS_KEY}
-      DANDI_API_KEY: ${DANDI_API_KEY}
-      DANDI_API_KEY_STAGING: ${DANDI_API_KEY_STAGING}
-    volumes:
-      - ./containers:/app
-      - ./results:/results
-      - ./logs:/logs
-    deploy:
-      resources:
-        reservations:
-          devices:
-            - driver: nvidia
-              count: 1
-              capabilities: [gpu]
+  # worker:
+  #   build:
+  #     context: containers
+  #     dockerfile: Dockerfile.combined
+  #   image: si-sorting-worker
+  #   # image: ghcr.io/catalystneuro/si-sorting-worker:latest
+  #   container_name: si-sorting-worker
+  #   ports:
+  #     - "5000:5000"
+  #   environment:
+  #     WORKER_DEPLOY_MODE: compose
+  #     AWS_DEFAULT_REGION: ${AWS_DEFAULT_REGION}
+  #     AWS_ACCESS_KEY_ID: ${AWS_ACCESS_KEY_ID}
+  #     AWS_SECRET_ACCESS_KEY: ${AWS_SECRET_ACCESS_KEY}
+  #     DANDI_API_KEY: ${DANDI_API_KEY}
+  #     DANDI_API_KEY_STAGING: ${DANDI_API_KEY_STAGING}
+  #   volumes:
+  #     - ./containers:/app
+  #     - ./results:/results
+  #     - ./logs:/logs
+  #   deploy:
+  #     resources:
+  #       reservations:
+  #         devices:
+  #           - driver: nvidia
+  #             count: 1
+  #             capabilities: [gpu]
 
   database:
     image: postgres:latest

rest/Dockerfile

@@ -19,4 +19,7 @@ ENV SI_CLOUD_ENV production
 ENV PYTHONUNBUFFERED=1
 
 EXPOSE 8000
-CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000", "--reload"]
+
+WORKDIR /
+
+CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000"]

rest/README.md

@@ -0,0 +1,7 @@
+# REST API
+
+To run REST API in local environment, from the root directory of the project run:
+
+```bash
+uvicorn rest.main:app --host 0.0.0.0 --port 8000 --workers 4 --reload
+```

rest/clients/aws.py

@@ -1,7 +1,7 @@
 import boto3
 import enum
 
-from core import settings
+from ..core import settings
 
 
 class JobStatus(enum.Enum):

rest/clients/database.py

@@ -4,7 +4,7 @@
 import ast
 import json
 
-from db.models import User, DataSource, Run
+from ..db.models import User, DataSource, Run
 
 
 class DatabaseClient:

rest/clients/local_docker.py

@@ -0,0 +1,87 @@
+from pathlib import Path
+import docker
+
+from ..core.logger import logger
+from ..models.sorting import (
+    RunKwargs,
+    SourceDataKwargs,
+    RecordingKwargs,
+    PreprocessingKwargs,
+    SorterKwargs,
+    PostprocessingKwargs,
+    CurationKwargs,
+    VisualizationKwargs,
+)
+
+
+class LocalDockerClient:
+
+    def __init__(self, base_url: str = "tcp://docker-proxy:2375"):
+        self.logger = logger
+        self.client = docker.DockerClient(base_url=base_url)
+
+    def run_sorting(
+        self, 
+        run_kwargs: RunKwargs,
+        source_data_kwargs: SourceDataKwargs,
+        recording_kwargs: RecordingKwargs,
+        preprocessing_kwargs: PreprocessingKwargs,
+        sorter_kwargs: SorterKwargs,
+        postprocessing_kwargs: PostprocessingKwargs,
+        curation_kwargs: CurationKwargs,
+        visualization_kwargs: VisualizationKwargs,
+    ) -> None:
+        # Pass kwargs as environment variables to the container
+        env_vars = dict(
+            SI_RUN_KWARGS=run_kwargs.json(),
+            SI_SOURCE_DATA_KWARGS=source_data_kwargs.json(),
+            SI_RECORDING_KWARGS=recording_kwargs.json(),
+            SI_PREPROCESSING_KWARGS=preprocessing_kwargs.json(),
+            SI_SORTER_KWARGS=sorter_kwargs.json(),
+            SI_POSTPROCESSING_KWARGS=postprocessing_kwargs.json(),
+            SI_CURATION_KWARGS=curation_kwargs.json(),
+            SI_VISUALIZATION_KWARGS=visualization_kwargs.json(),
+        )
+
+        # Local volumes to mount
+        local_directory = Path(".").absolute()
+        logs_directory = local_directory / "logs"
+        results_directory = local_directory / "results"
+        volumes = {
+            logs_directory: {'bind': '/logs', 'mode': 'rw'},
+            results_directory: {'bind': '/results', 'mode': 'rw'},
+        }
+
+        container = self.client.containers.run(
+            image='python:slim',
+            command=['python', '-c', 'import os; print(os.environ.get("SI_RUN_KWARGS"))'],
+            detach=True,
+            environment=env_vars,
+            volumes=volumes,
+            device_requests=[
+                docker.types.DeviceRequest(
+                    device_ids=["0"], 
+                    capabilities=[['gpu']]
+                )
+            ]
+        )
+        # if response.status_code == 200:
+        #     self.logger.info("Success!")
+        # else:
+        #     self.logger.info(f"Error {response.status_code}: {response.content}")
+
+
+    def get_run_logs(self, run_identifier):
+        # TODO: Implement this
+        self.logger.info("Getting logs...")
+        # response = requests.get(self.url + "/logs", params={"run_identifier": run_identifier})
+        # if response.status_code == 200:
+        #     logs = response.content.decode('utf-8')
+        #     if "Error running sorter" in logs:
+        #         return "fail", logs
+        #     elif "Sorting job completed successfully!" in logs:
+        #         return "success", logs
+        #     return "running", logs
+        # else:
+        #     self.logger.info(f"Error {response.status_code}: {response.content}")
+        #     return "fail", f"Logs couldn't be retrieved. Error {response.status_code}: {response.content}"

rest/db/utils.py

@@ -2,7 +2,7 @@
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import relationship, sessionmaker
 
-from db.models import Base, User, DataSource, Run
+from .models import Base, User, DataSource, Run
 
 
 def initialize_db(db: str):

rest/main.py

@@ -5,13 +5,13 @@
 from fastapi.responses import JSONResponse
 from pathlib import Path
 
-from core.settings import settings
-from routes.user import router as router_user
-from routes.dandi import router as router_dandi
-from routes.sorting import router as router_sorting
-from routes.runs import router as router_runs
-from clients.dandi import DandiClient
-from db.utils import initialize_db
+from .core.settings import settings
+from .routes.user import router as router_user
+from .routes.dandi import router as router_dandi
+from .routes.sorting import router as router_sorting
+from .routes.runs import router as router_runs
+from .clients.dandi import DandiClient
+from .db.utils import initialize_db
 import logging