From 44f9fff31add1ecc2f500208033fee0f04f57f7e Mon Sep 17 00:00:00 2001 From: luiz Date: Fri, 13 Oct 2023 15:04:51 +0200 Subject: [PATCH] wip refactor main.py --- containers/Dockerfile.ks2_5 | 2 +- containers/README.md | 3 +- containers/main.py | 68 ++++++++++++------- ...compose-dev.yml => docker-compose-dev.yaml | 0 docker-compose.yml => docker-compose.yaml | 0 rest/clients/local_docker.py | 13 +++- 6 files changed, 58 insertions(+), 28 deletions(-) rename docker-compose-dev.yml => docker-compose-dev.yaml (100%) rename docker-compose.yml => docker-compose.yaml (100%) diff --git a/containers/Dockerfile.ks2_5 b/containers/Dockerfile.ks2_5 index 602d67b..90bb79b 100644 --- a/containers/Dockerfile.ks2_5 +++ b/containers/Dockerfile.ks2_5 @@ -1,5 +1,5 @@ # Spike sorters image -FROM spikeinterface/kilosort2_5-compiled-base:0.2.0 as ks25base +FROM spikeinterface/kilosort2_5-compiled-base:0.2.0 # # NVIDIA-ready Image # FROM nvidia/cuda:11.6.2-base-ubuntu20.04 diff --git a/containers/README.md b/containers/README.md index c85cdd0..0d22f95 100644 --- a/containers/README.md +++ b/containers/README.md @@ -16,7 +16,8 @@ Basic infrastructure makes use of the following AWS services: Build docker image: ```bash -$ DOCKER_BUILDKIT=1 docker build -t -f . +$ DOCKER_BUILDKIT=1 docker build -t ghcr.io/catalystneuro/si-sorting-ks25:latest -f Dockerfile.ks2_5 . +$ DOCKER_BUILDKIT=1 docker build -t ghcr.io/catalystneuro/si-sorting-ks3:latest -f Dockerfile.ks3 . ``` Run locally: diff --git a/containers/main.py b/containers/main.py index 8fbd742..ec1103a 100644 --- a/containers/main.py +++ b/containers/main.py @@ -588,10 +588,10 @@ def main( run_at = validate_not_none(run_kwargs, "run_at") run_identifier = run_kwargs.get("run_identifier", datetime.now().strftime("%Y%m%d%H%M%S")) run_description = run_kwargs.get("run_description", "") - test_with_toy_recording = run_kwargs.get("test_with_toy_recording", "False").lower() in ("true", "1", "t") - test_with_subrecording = run_kwargs.get("test_with_subrecording", "False").lower() in ("true", "1", "t") + test_with_toy_recording = run_kwargs.get("test_with_toy_recording", "False") + test_with_subrecording = run_kwargs.get("test_with_subrecording", "False") test_subrecording_n_frames = int(run_kwargs.get("test_subrecording_n_frames", 30000)) - log_to_file = run_kwargs.get("log_to_file", "False").lower() in ("true", "1", "t") + log_to_file = run_kwargs.get("log_to_file", "False") # Get source data kwargs from ENV variables source_data_kwargs = json.loads(os.environ.get("SI_SOURCE_DATA_KWARGS", "{}")) @@ -622,27 +622,47 @@ def main( output_destination = validate_not_none(output_kwargs, "output_destination") output_path = validate_not_none(output_kwargs, "output_path") - # Run main function - main( - run_at=run_at, - run_identifier=run_identifier, - run_description=run_description, - test_with_toy_recording=test_with_toy_recording, - test_with_subrecording=test_with_subrecording, - test_subrecording_n_frames=test_subrecording_n_frames, - log_to_file=log_to_file, - source_name=source_name, - source_data_type=source_data_type, - source_data_paths=source_data_paths, - recording_kwargs=recording_kwargs, - preprocessing_kwargs=preprocessing_kwargs, - sorter_kwargs=sorter_kwargs, - postprocessing_kwargs=postprocessing_kwargs, - curation_kwargs=curation_kwargs, - visualization_kwargs=visualization_kwargs, - output_destination=output_destination, - output_path=output_path, - ) + # # Run main function + # main( + # run_at=run_at, + # run_identifier=run_identifier, + # run_description=run_description, + # test_with_toy_recording=test_with_toy_recording, + # test_with_subrecording=test_with_subrecording, + # test_subrecording_n_frames=test_subrecording_n_frames, + # log_to_file=log_to_file, + # source_name=source_name, + # source_data_type=source_data_type, + # source_data_paths=source_data_paths, + # recording_kwargs=recording_kwargs, + # preprocessing_kwargs=preprocessing_kwargs, + # sorter_kwargs=sorter_kwargs, + # postprocessing_kwargs=postprocessing_kwargs, + # curation_kwargs=curation_kwargs, + # visualization_kwargs=visualization_kwargs, + # output_destination=output_destination, + # output_path=output_path, + # ) + + print("\nRun at: ", run_at) + print("\nRun identifier: ", run_identifier) + print("\nRun description: ", run_description) + print("\nTest with toy recording: ", test_with_toy_recording) + print("\nTest with subrecording: ", test_with_subrecording) + print("\nTest subrecording n frames: ", test_subrecording_n_frames) + print("\nLog to file: ", log_to_file) + print("\nSource name: ", source_name) + print("\nSource data type: ", source_data_type) + print("\nSource data paths: ", source_data_paths) + print("\nRecording kwargs: ", recording_kwargs) + print("\nPreprocessing kwargs: ", preprocessing_kwargs) + print("\nSorter kwargs: ", sorter_kwargs) + print("\nPostprocessing kwargs: ", postprocessing_kwargs) + print("\nCuration kwargs: ", curation_kwargs) + print("\nVisualization kwargs: ", visualization_kwargs) + print("\nOutput destination: ", output_destination) + print("\nOutput path: ", output_path) + # Known issues: diff --git a/docker-compose-dev.yml b/docker-compose-dev.yaml similarity index 100% rename from docker-compose-dev.yml rename to docker-compose-dev.yaml diff --git a/docker-compose.yml b/docker-compose.yaml similarity index 100% rename from docker-compose.yml rename to docker-compose.yaml diff --git a/rest/clients/local_docker.py b/rest/clients/local_docker.py index fbd2042..c8216f7 100644 --- a/rest/clients/local_docker.py +++ b/rest/clients/local_docker.py @@ -15,6 +15,15 @@ ) +map_sorter_to_image = { + "kilosort2": "ghcr.io/catalystneuro/si-sorting-ks2:latest", + "kilosort25": "ghcr.io/catalystneuro/si-sorting-ks25:latest", + "kilosort3": "ghcr.io/catalystneuro/si-sorting-ks3:latest", + "ironclust": "ghcr.io/catalystneuro/si-sorting-ironclust:latest", + "spykingcircus": "ghcr.io/catalystneuro/si-sorting-spyking-circus:latest", +} + + class LocalDockerClient: def __init__(self, base_url: str = "tcp://docker-proxy:2375"): @@ -57,8 +66,8 @@ def run_sorting( container = self.client.containers.run( name=f'si-sorting-run-{run_kwargs.run_identifier}', - image='python:slim', - command=['python', '-c', 'import os; print(os.environ.get("SI_RUN_KWARGS"))'], + image=map_sorter_to_image[sorter_kwargs.sorter_name], + command=['python', 'main.py'], detach=True, environment=env_vars, volumes=volumes,