Repository: PrefectHQ/prefect

Issue:
Expose ports in DockerRun
## Current behavior

It is currently very limiting to expose container ports when using Prefect's DockerRun. In normal Docker CLI, this would be:

`docker run --name test -p 12001:12001`

Using the Docker Python client, you can expose ports by supplying the `ports` argument and the `host_config`.

```
client.api.create_container(
    'busybox', 'ls', ports=[1111, 2222],
    host_config=client.api.create_host_config(port_bindings={
        1111: 4567,
        2222: None
    })
)
```

## Proposed behavior

In order to support this, we need DockerRun to take in a `ports` argument and then pass it to the agent. The agent currently spins up a container with the following command (ports is not exposed).

```
container = self.docker_client.create_container(
    image,
    command=get_flow_run_command(flow_run),
    environment=env_vars,
    name=container_name,
    volumes=container_mount_paths,
    host_config=self.docker_client.create_host_config(**host_config),
    networking_config=networking_config,
    labels=labels,
)
```

## Example

Supporting this would let a user do:

```
flow.run_config = DockerRun(
    image="some_image",
    ports=[12001],
    host_config={"port_bindings":{12001:12001}},
)
```



Discussion:
I'm keen to take this on. 
@s-mawjee go for it!

Resolving patch:
diff --git a/src/prefect/agent/docker/agent.py b/src/prefect/agent/docker/agent.py
--- a/src/prefect/agent/docker/agent.py
+++ b/src/prefect/agent/docker/agent.py
@@ -1,13 +1,13 @@
+import multiprocessing
 import ntpath
 import posixpath
-from packaging.version import parse as parse_version
-
-import multiprocessing
 import re
 import warnings
-from slugify import slugify
 from sys import platform
-from typing import TYPE_CHECKING, Dict, Iterable, List, Tuple, Any
+from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Tuple
+
+from packaging.version import parse as parse_version
+from slugify import slugify
 
 from prefect import config, context
 from prefect.agent import Agent
@@ -410,6 +410,8 @@ def deploy_flow(self, flow_run: GraphQLResult) -> str:
 
         # By default, auto-remove containers
         host_config: Dict[str, Any] = {"auto_remove": True}
+        # By default, no ports
+        ports = None
 
         # Set up a host gateway for local communication; check the docker version since
         # this is not supported by older versions
@@ -434,6 +436,8 @@ def deploy_flow(self, flow_run: GraphQLResult) -> str:
         if run_config is not None and run_config.host_config:
             # The host_config passed from the run_config will overwrite defaults
             host_config.update(run_config.host_config)
+        if run_config is not None and run_config.ports:
+            ports = run_config.ports
 
         networking_config = None
         # At the time of creation, you can only connect a container to a single network,
@@ -484,6 +488,7 @@ def deploy_flow(self, flow_run: GraphQLResult) -> str:
                     host_config=self.docker_client.create_host_config(**host_config),
                     networking_config=networking_config,
                     labels=labels,
+                    ports=ports,
                 )
             except docker.errors.APIError as exc:
                 if "Conflict" in str(exc) and "container name" in str(exc):
diff --git a/src/prefect/run_configs/docker.py b/src/prefect/run_configs/docker.py
--- a/src/prefect/run_configs/docker.py
+++ b/src/prefect/run_configs/docker.py
@@ -13,6 +13,8 @@ class DockerRun(RunConfig):
         - labels (Iterable[str], optional): an iterable of labels to apply to this
             run config. Labels are string identifiers used by Prefect Agents
             for selecting valid flow runs when polling for work
+        - ports (Iterable[int], optional): an iterable of ports to pass to the
+            Docker Agent which are used to expose container ports.
         - host_config (dict, optional): A dictionary or runtime arguments to pass to
             the Docker Agent. See link below for options.
             https://docker-py.readthedocs.io/en/stable/api.html#docker.api.container.ContainerApiMixin.create_host_config
@@ -38,8 +40,10 @@ def __init__(
         image: str = None,
         env: dict = None,
         labels: Iterable[str] = None,
-        host_config: dict = None
+        ports: Iterable[int] = None,
+        host_config: dict = None,
     ) -> None:
         super().__init__(env=env, labels=labels)
         self.image = image
         self.host_config = host_config
+        self.ports = ports
diff --git a/src/prefect/serialization/run_config.py b/src/prefect/serialization/run_config.py
--- a/src/prefect/serialization/run_config.py
+++ b/src/prefect/serialization/run_config.py
@@ -1,18 +1,18 @@
 from marshmallow import fields
 
-from prefect.utilities.serialization import (
-    JSONCompatible,
-    OneOfSchema,
-    ObjectSchema,
-    SortedList,
-)
 from prefect.run_configs import (
-    KubernetesRun,
-    LocalRun,
     DockerRun,
     ECSRun,
-    VertexRun,
+    KubernetesRun,
+    LocalRun,
     UniversalRun,
+    VertexRun,
+)
+from prefect.utilities.serialization import (
+    JSONCompatible,
+    ObjectSchema,
+    OneOfSchema,
+    SortedList,
 )
 
 
@@ -68,6 +68,7 @@ class DockerRunSchema(RunConfigSchemaBase):
     class Meta:
         object_class = DockerRun
 
+    ports = fields.List(fields.Int, allow_none=True)
     image = fields.String(allow_none=True)
     host_config = fields.Dict(keys=fields.String(), allow_none=True)
 


Test patch:
diff --git a/tests/agent/test_docker_agent.py b/tests/agent/test_docker_agent.py
--- a/tests/agent/test_docker_agent.py
+++ b/tests/agent/test_docker_agent.py
@@ -298,7 +298,6 @@ def test_prefect_logging_level_override_logic(
     ],
 )
 def test_docker_agent_deploy_flow(core_version, command, api):
-
     agent = DockerAgent()
     agent.deploy_flow(
         flow_run=GraphQLResult(
@@ -467,7 +466,6 @@ def test_docker_agent_deploy_flow_sets_container_name_with_slugify(
 @pytest.mark.parametrize("run_kind", ["docker", "missing", "universal"])
 @pytest.mark.parametrize("has_docker_storage", [True, False])
 def test_docker_agent_deploy_flow_run_config(api, run_kind, has_docker_storage):
-
     if has_docker_storage:
         storage = Docker(
             registry_url="testing", image_name="on-storage", image_tag="tag"
@@ -479,14 +477,16 @@ def test_docker_agent_deploy_flow_run_config(api, run_kind, has_docker_storage):
 
     if run_kind == "docker":
         env = {"TESTING": "VALUE"}
+        ports = [12001]
         host_config = {"auto_remove": False, "shm_size": "128m"}
         exp_host_config = {
             "auto_remove": False,
             "shm_size": "128m",
         }
-        run = DockerRun(image=image, env=env, host_config=host_config)
+        run = DockerRun(image=image, env=env, ports=ports, host_config=host_config)
     else:
         env = {}
+        ports = []
         host_config = {}
         exp_host_config = {
             "auto_remove": True,
@@ -519,6 +519,9 @@ def test_docker_agent_deploy_flow_run_config(api, run_kind, has_docker_storage):
     res_env = api.create_container.call_args[1]["environment"]
     for k, v in env.items():
         assert res_env[k] == v
+    res_ports = api.create_container.call_args[1]["ports"]
+    for i, v in enumerate(ports):
+        assert res_ports[i] == v
     res_host_config = api.create_host_config.call_args[1]
     for k, v in exp_host_config.items():
         assert res_host_config[k] == v
diff --git a/tests/run_configs/test_docker.py b/tests/run_configs/test_docker.py
--- a/tests/run_configs/test_docker.py
+++ b/tests/run_configs/test_docker.py
@@ -5,6 +5,7 @@ def test_no_args():
     config = DockerRun()
     assert config.env is None
     assert config.image is None
+    assert config.ports is None
     assert config.labels == set()
     assert config.host_config is None
 
@@ -15,9 +16,11 @@ def test_all_args(tmpdir):
         env={"hello": "world"},
         image="testing",
         labels=["a", "b"],
+        ports=[12001],
         host_config={"host": "config"},
     )
     assert config.env == {"hello": "world"}
     assert config.image == "testing"
     assert config.labels == {"a", "b"}
+    assert config.ports == [12001]
     assert config.host_config == {"host": "config"}
diff --git a/tests/serialization/test_run_configs.py b/tests/serialization/test_run_configs.py
--- a/tests/serialization/test_run_configs.py
+++ b/tests/serialization/test_run_configs.py
@@ -95,9 +95,7 @@ def test_serialize_local_run(config):
     [
         DockerRun(),
         DockerRun(
-            env={"test": "foo"},
-            image="testing",
-            labels=["a", "b"],
+            env={"test": "foo"}, image="testing", labels=["a", "b"], ports=[12001]
         ),
     ],
 )
@@ -105,7 +103,7 @@ def test_serialize_docker_run(config):
     msg = RunConfigSchema().dump(config)
     config2 = RunConfigSchema().load(msg)
     assert sorted(config.labels) == sorted(config2.labels)
-    fields = ["env", "image"]
+    fields = ["env", "image", "ports"]
     for field in fields:
         assert getattr(config, field) == getattr(config2, field)