From dc702aeb1e787278e9f38cd51337cfbf48c8e75b Mon Sep 17 00:00:00 2001 From: Umesh Date: Fri, 13 Oct 2023 11:53:11 -0500 Subject: [PATCH 1/2] Reset worker state after collection is finished. Closes #164 --- chimerapy/orchestrator/cli/__main__.py | 1 + chimerapy/orchestrator/models/pipeline_config.py | 5 +++++ 2 files changed, 6 insertions(+) diff --git a/chimerapy/orchestrator/cli/__main__.py b/chimerapy/orchestrator/cli/__main__.py index 658cb28..8129d44 100644 --- a/chimerapy/orchestrator/cli/__main__.py +++ b/chimerapy/orchestrator/cli/__main__.py @@ -115,6 +115,7 @@ async def aorchestrate(config: ChimeraPyPipelineConfig) -> None: "manager.timeout.worker-shutdown", config.timeouts.shutdown_timeout ) + await manager.async_reset(keep_workers=config.keep_remote_workers) await manager.async_shutdown() print("Shutting down local workers...") for worker in local_workers: diff --git a/chimerapy/orchestrator/models/pipeline_config.py b/chimerapy/orchestrator/models/pipeline_config.py index f7d8c4a..e10f592 100644 --- a/chimerapy/orchestrator/models/pipeline_config.py +++ b/chimerapy/orchestrator/models/pipeline_config.py @@ -148,6 +148,11 @@ class ChimeraPyPipelineConfig(BaseModel): description="The timeouts for the pipeline operation.", ) + keep_remote_workers: bool = Field( + default=True, + description="If true, keep remote workers after pipeline has finished collection.", + ) + def instantiate_manager(self) -> cpe.Manager: m = cpe.Manager( **self.manager_config.model_dump( From bbf0cd685838765c768880dec59c37c9a9b38204 Mon Sep 17 00:00:00 2001 From: Umesh Date: Sun, 15 Oct 2023 15:30:21 -0500 Subject: [PATCH 2/2] Add option to keep workers during shutdown instead --- chimerapy/orchestrator/cli/__main__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/chimerapy/orchestrator/cli/__main__.py b/chimerapy/orchestrator/cli/__main__.py index 8129d44..b7b8e86 100644 --- a/chimerapy/orchestrator/cli/__main__.py +++ b/chimerapy/orchestrator/cli/__main__.py @@ -115,8 +115,8 @@ async def aorchestrate(config: ChimeraPyPipelineConfig) -> None: "manager.timeout.worker-shutdown", config.timeouts.shutdown_timeout ) - await manager.async_reset(keep_workers=config.keep_remote_workers) - await manager.async_shutdown() + await manager.async_reset() + await manager.async_shutdown(shutdown_workers=config.keep_remote_workers) print("Shutting down local workers...") for worker in local_workers: await worker.async_shutdown()