Merge pull request #551 from Sahil-u07/improve-docker-compose

feat: improve docker-compose generation with restart, depends_on, network, and ZMQ mode

Author: pradeeban

concore_cli/cli.py

@@ -75,8 +75,16 @@ def init(name, template, interactive):
     is_flag=True,
     help="Generate docker-compose.yml in output directory (docker type only)",
 )
-def build(workflow_file, source, output, type, auto_build, compose):
+@click.option(
+    "--zmq",
+    is_flag=True,
+    help="Configure compose for ZMQ networking mode (requires --compose)",
+)
+def build(workflow_file, source, output, type, auto_build, compose, zmq):
     """Compile a concore workflow into executable scripts"""
+    if zmq and not compose:
+        console.print("[red]Error:[/red] --zmq requires --compose")
+        sys.exit(1)
     try:
         build_workflow(
             workflow_file,
@@ -86,6 +94,7 @@ def build(workflow_file, source, output, type, auto_build, compose):
             auto_build,
             console,
             compose=compose,
+            zmq_mode=zmq,
         )
     except Exception as e:
         console.print(f"[red]Error:[/red] {str(e)}")

concore_cli/commands/build.py

@@ -75,9 +75,15 @@ def _parse_docker_run_line(line):
     }
 
 
-def _write_docker_compose(output_path):
+def _write_docker_compose(output_path, console, zmq_mode=False):
     run_script = output_path / "run"
     if not run_script.exists():
+        console.print(
+            f"[yellow]Warning:[/yellow] No docker run script found in {output_path}."
+        )
+        console.print(
+            "[dim]Tip: run concore build --type docker first, then use --compose[/dim]"
+        )
         return None
 
     services = []
@@ -89,15 +95,10 @@ def _write_docker_compose(output_path):
     if not services:
         return None
 
-    compose_lines = [
-        "networks:",
-        "  concore-net:",
-        "    driver: bridge",
-        "",
-        "services:",
-    ]
+    compose_lines = ["services:"]
 
     named_volumes = set()
+    previous_service_name = None
     for index, service in enumerate(services, start=1):
         service_name = re.sub(r"[^A-Za-z0-9_.-]", "-", service["container_name"]).strip(
             "-."
@@ -107,14 +108,11 @@ def _write_docker_compose(output_path):
         elif not service_name[0].isalpha():
             service_name = f"service-{service_name}"
 
-        compose_lines.append(f"  {_yaml_quote(service_name)}:")
+        compose_lines.append(f"  {service_name}:")
         compose_lines.append(f"    image: {_yaml_quote(service['image'])}")
         compose_lines.append(
             f"    container_name: {_yaml_quote(service['container_name'])}"
         )
-        compose_lines.append("    restart: on-failure")
-        compose_lines.append("    networks:")
-        compose_lines.append("      - concore-net")
 
         if service["volumes"]:
             compose_lines.append("    volumes:")
@@ -124,12 +122,28 @@ def _write_docker_compose(output_path):
                 if re.match(r"^[a-zA-Z0-9_-]+$", part1):
                     named_volumes.add(part1)
 
+        compose_lines.append("    restart: on-failure")
+        if zmq_mode:
+            compose_lines.append("    environment:")
+            compose_lines.append("      - CONCORE_TRANSPORT=zmq")
+        if index > 1 and previous_service_name:
+            compose_lines.append("    depends_on:")
+            compose_lines.append(f"      - {previous_service_name}")
+        compose_lines.append("    networks:")
+        compose_lines.append("      - concore_net")
+        previous_service_name = service_name
+
     if named_volumes:
         compose_lines.append("")
         compose_lines.append("volumes:")
         for v in sorted(named_volumes):
             compose_lines.append(f"  {v}:")
 
+    compose_lines.append("")
+    compose_lines.append("networks:")
+    compose_lines.append("  concore_net:")
+    compose_lines.append("    driver: bridge")
+
     compose_lines.append("")
     compose_path = output_path / "docker-compose.yml"
     compose_path.write_text("\n".join(compose_lines), encoding="utf-8")
@@ -144,6 +158,7 @@ def build_workflow(
     auto_build,
     console,
     compose=False,
+    zmq_mode=False,
 ):
     workflow_path = Path(workflow_file).resolve()
     source_path = Path(source).resolve()
@@ -238,7 +253,9 @@ def build_workflow(
             )
 
             if compose:
-                compose_path = _write_docker_compose(output_path)
+                compose_path = _write_docker_compose(
+                    output_path, console, zmq_mode=zmq_mode
+                )
                 if compose_path is not None:
                     console.print(
                         f"[green]✓[/green] Compose file written to [cyan]{compose_path}[/cyan]"

tests/test_cli.py

@@ -384,8 +384,8 @@ def test_build_command_docker_compose_single_node(self):
             self.assertIn("container_name: 'N1'", compose_content)
             self.assertIn("image: 'docker-script'", compose_content)
             self.assertIn("networks:", compose_content)
-            self.assertIn("concore-net:", compose_content)
-            self.assertIn("- concore-net", compose_content)
+            self.assertIn("concore_net:", compose_content)
+            self.assertIn("- concore_net", compose_content)
             self.assertIn("restart: on-failure", compose_content)
 
             metadata = json.loads(Path("out/STUDY.json").read_text())

tests/test_compose_generation.py

@@ -0,0 +1,85 @@
+from concore_cli.commands.build import _write_docker_compose
+from rich.console import Console
+from pathlib import Path
+
+
+def _fake_run_script(output_dir, services):
+    lines = [
+        f"docker run --name {s['name']} -v /study:/study {s['image']} &"
+        for s in services
+    ]
+    (Path(output_dir) / "run").write_text("\n".join(lines))
+
+
+def test_compose_has_restart_policy(tmp_path):
+    _fake_run_script(tmp_path, [{"name": "node1", "image": "concore/py"}])
+    path = _write_docker_compose(tmp_path, Console(quiet=True))
+    assert path is not None
+    content = path.read_text()
+    assert "restart: on-failure" in content
+
+
+def test_compose_has_network_section(tmp_path):
+    _fake_run_script(tmp_path, [{"name": "node1", "image": "concore/py"}])
+    path = _write_docker_compose(tmp_path, Console(quiet=True))
+    content = path.read_text()
+    assert "concore_net" in content
+    assert "networks:" in content
+
+
+def test_compose_depends_on_second_service(tmp_path):
+    _fake_run_script(
+        tmp_path,
+        [
+            {"name": "controller", "image": "concore/py"},
+            {"name": "plant", "image": "concore/cpp"},
+        ],
+    )
+    path = _write_docker_compose(tmp_path, Console(quiet=True))
+    content = path.read_text()
+    assert "depends_on" in content
+    assert "controller" in content
+
+
+def test_compose_first_service_has_no_depends_on(tmp_path):
+    _fake_run_script(
+        tmp_path,
+        [
+            {"name": "controller", "image": "concore/py"},
+            {"name": "plant", "image": "concore/cpp"},
+        ],
+    )
+    path = _write_docker_compose(tmp_path, Console(quiet=True))
+    lines = path.read_text().splitlines()
+    controller_idx = next(i for i, line in enumerate(lines) if "controller:" in line)
+    plant_idx = next(i for i, line in enumerate(lines) if "plant:" in line)
+    section = lines[controller_idx:plant_idx]
+    assert not any("depends_on" in line for line in section)
+
+
+def test_zmq_mode_adds_env(tmp_path):
+    _fake_run_script(tmp_path, [{"name": "node1", "image": "concore/py"}])
+    path = _write_docker_compose(tmp_path, Console(quiet=True), zmq_mode=True)
+    content = path.read_text()
+    assert "CONCORE_TRANSPORT=zmq" in content
+
+
+def test_no_zmq_env_in_default_mode(tmp_path):
+    _fake_run_script(tmp_path, [{"name": "node1", "image": "concore/py"}])
+    path = _write_docker_compose(tmp_path, Console(quiet=True), zmq_mode=False)
+    content = path.read_text()
+    assert "CONCORE_TRANSPORT" not in content
+
+
+def test_missing_run_script_returns_none(tmp_path):
+    result = _write_docker_compose(tmp_path, Console(quiet=True))
+    assert result is None
+
+
+def test_zmq_without_compose_errors():
+    from click.testing import CliRunner
+    from concore_cli.cli import cli
+
+    runner = CliRunner()
+    result = runner.invoke(cli, ["build", "wf.graphml", "--zmq"])
+    assert result.exit_code != 0