add container port support to Task by Hera (#411)

src/hera/task.py

@@ -35,6 +35,7 @@
 from hera.metric import Metric, Metrics
 from hera.operator import Operator
 from hera.parameter import Parameter
+from hera.port import ContainerPort
 from hera.resource_template import ResourceTemplate
 from hera.resources import Resources
 from hera.retry_strategy import RetryStrategy
@@ -170,6 +171,8 @@ class Task(IO):
         Any built-in/custom Prometheus metrics to track.
     sidecars: Optional[List[Sidecar]] = None
         List of sidecars to create for the main pods of the container that runs the task.
+    ports: Optional[List[ContainerPort]] = None
+        List of ports to create for the main pods of the container that runs the task.
 
     Notes
     -----
@@ -213,6 +216,7 @@ def __init__(
         timeout: Optional[str] = None,
         metrics: Optional[Union[Metric, List[Metric], Metrics]] = None,
         sidecars: Optional[List[Sidecar]] = None,
+        ports: Optional[List[ContainerPort]] = None,
     ):
         if dag and source:
             raise ValueError("Cannot use both `dag` and `source`")
@@ -271,6 +275,7 @@ def __init__(
         self.is_exit_task: bool = False
         self.depends: Optional[str] = None
         self.when: Optional[str] = None
+        self.ports = ports
 
         self.validate()
 
@@ -923,6 +928,7 @@ def _build_container_kwargs(self) -> Dict:
             command=self.get_command(),
             resources=self.resources.build() if self.resources else None,
             args=self.get_args(),
+            ports=None if self.ports is None else [p.build() for p in self.ports],
             env=env,
             env_from=env_from,
             working_dir=self.working_dir,

tests/test_task.py

@@ -21,6 +21,7 @@
     ConfigMapEnv,
     ConfigMapEnvFrom,
     ConfigMapVolume,
+    ContainerPort,
     EmptyDirVolume,
     Env,
     ExistingVolume,
@@ -214,6 +215,16 @@ def test_volume_mounts_returns_expected_volumes(self, no_op):
         assert vs[3].name
         assert vs[3].mount_path == "/v3"
 
+    def test_container_port_returns_expected_ports(self, no_op):
+        t = Task(
+            "t",
+            no_op,
+            ports=[ContainerPort(8080, name="test-port")],
+        )
+        cp = t.ports
+        assert cp[0].name == "test-port"
+        assert cp[0].container_port == 8080
+
     def test_gpu_toleration_returns_expected_toleration(self):
         tn = GPUToleration
         assert tn.key == "nvidia.com/gpu"