diff --git a/airflow/providers/docker/operators/docker_swarm.py b/airflow/providers/docker/operators/docker_swarm.py index a1f3f0b2bcf2d..9a5067940f46a 100644 --- a/airflow/providers/docker/operators/docker_swarm.py +++ b/airflow/providers/docker/operators/docker_swarm.py @@ -105,6 +105,9 @@ class DockerSwarmOperator(DockerOperator): :type mode: docker.types.ServiceMode :param networks: List of network names or IDs or NetworkAttachmentConfig to attach the service to. :type networks: List[Union[str, NetworkAttachmentConfig]] + :param placement: Placement instructions for the scheduler. If a list is passed instead, + it is assumed to be a list of constraints as part of a Placement object. + :type placement: Union[types.Placement, List[types.Placement]] """ def __init__( @@ -116,6 +119,7 @@ def __init__( secrets: Optional[List[types.SecretReference]] = None, mode: Optional[types.ServiceMode] = None, networks: Optional[List[Union[str, types.NetworkAttachmentConfig]]] = None, + placement: Optional[Union[types.Placement, List[types.Placement]]] = None, **kwargs, ) -> None: super().__init__(image=image, **kwargs) @@ -126,6 +130,7 @@ def __init__( self.secrets = secrets self.mode = mode self.networks = networks + self.placement = placement def execute(self, context) -> None: self.cli = self._get_cli() @@ -153,6 +158,7 @@ def _run_service(self) -> None: restart_policy=types.RestartPolicy(condition='none'), resources=types.Resources(mem_limit=self.mem_limit), networks=self.networks, + placement=self.placement, ), name=f'airflow-{get_random_string()}', labels={'name': f'airflow__{self.dag_id}__{self.task_id}'}, diff --git a/tests/providers/docker/operators/test_docker_swarm.py b/tests/providers/docker/operators/test_docker_swarm.py index 09207b425dbb8..d2c0c9c5dd953 100644 --- a/tests/providers/docker/operators/test_docker_swarm.py +++ b/tests/providers/docker/operators/test_docker_swarm.py @@ -71,6 +71,7 @@ def _client_service_logs_effect(): secrets=[types.SecretReference(secret_id="dummy_secret_id", secret_name="dummy_secret_name")], mode=types.ServiceMode(mode="replicated", replicas=3), networks=["dummy_network"], + placement=types.Placement(constraints=["node.labels.region==east"]), ) operator.execute(None) @@ -79,6 +80,7 @@ def _client_service_logs_effect(): restart_policy=mock_obj, resources=mock_obj, networks=["dummy_network"], + placement=types.Placement(constraints=["node.labels.region==east"]), ) types_mock.ContainerSpec.assert_called_once_with( image='ubuntu:latest',