We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Update mypy to 0.810


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `src/py/flwr_experimental/ops/cluster.py`
Content:
```
1 # Copyright 2020 Adap GmbH. All Rights Reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 # ==============================================================================
15 """Implements compute classes for EC2."""
16 import concurrent.futures
17 from contextlib import contextmanager
18 from itertools import groupby
19 from logging import DEBUG, ERROR
20 from typing import Dict, Iterator, List, Optional, Tuple
21 
22 from paramiko.client import SSHClient
23 from paramiko.sftp_attr import SFTPAttributes
24 
25 from flwr.common.logger import log
26 
27 from .compute.adapter import Adapter
28 from .instance import Instance
29 
30 ExecInfo = Tuple[str, str]
31 
32 
33 class StartFailed(Exception):
34     """Raised when cluster could not start."""
35 
36 
37 class InstanceIdNotFound(Exception):
38     """Raised when there was no instance with given id."""
39 
40 
41 class InstanceMismatch(Exception):
42     """Raised when instances passed to create_instances do not have the same
43     values for RAM or CPU."""
44 
45 
46 class IgnoreHostKeyPolicy:
47     """Policy for accepting any unknown host key.
48 
49     This is used by `paramiko.client.SSHClient`.
50     """
51 
52     # pylint: disable=no-self-use, unused-argument
53     def missing_host_key(self, client: SSHClient, hostname: str, key: str) -> None:
54         """Simply return to ignore the host key.
55 
56         As we create and destroy machines quite regularly and don't
57         reuse them we will not store the host key in the local system to
58         avoid pollution the local known_hosts file.
59         """
60         return None
61 
62 
63 SSHCredentials = Tuple[str, str]  # username, key_filename
64 
65 
66 @contextmanager
67 def ssh_connection(
68     instance: Instance, ssh_credentials: SSHCredentials
69 ) -> Iterator[SSHClient]:
70     """Connect to server and yield SSH client."""
71     username, key_filename = ssh_credentials
72 
73     client = SSHClient()
74     client.set_missing_host_key_policy(IgnoreHostKeyPolicy)
75     client.connect(
76         hostname=instance.public_ip,
77         port=instance.ssh_port,
78         username=username,
79         key_filename=key_filename,
80     )
81 
82     yield client
83 
84     client.close()
85 
86 
87 def create_instances(adapter: Adapter, instances: List[Instance], timeout: int) -> None:
88     """Start instances and set props of each instance.
89 
90     Fails if CPU and RAM of instances are not all the same.
91     """
92     if not all(
93         [
94             ins.num_cpu == instances[0].num_cpu and ins.num_ram == instances[0].num_ram
95             for ins in instances
96         ]
97     ):
98         raise InstanceMismatch(
99             "Values of num_cpu and num_ram have to be equal for all instances."
100         )
101 
102     # As checked before that each instance has the same num_cpu and num_ram
103     # we can just take the values from the first => instances[0]
104     adapter_instances = adapter.create_instances(
105         num_cpu=instances[0].num_cpu,
106         num_ram=instances[0].num_ram,
107         num_instance=len(instances),
108         gpu=instances[0].gpu,
109         timeout=timeout,
110     )
111 
112     for i, adp_ins in enumerate(adapter_instances):
113         instance_id, private_ip, public_ip, ssh_port, state = adp_ins
114 
115         instances[i].instance_id = instance_id
116         instances[i].private_ip = private_ip
117         instances[i].public_ip = public_ip
118         instances[i].ssh_port = ssh_port
119         instances[i].state = state
120 
121 
122 def group_instances_by_specs(instances: List[Instance]) -> List[List[Instance]]:
123     """Group instances by num_cpu and num_ram."""
124     groups: List[List[Instance]] = []
125     keyfunc = lambda ins: f"{ins.num_cpu}-{ins.num_ram}"
126     instances = sorted(instances, key=keyfunc)
127     for _, group in groupby(instances, keyfunc):
128         groups.append(list(group))
129     return groups
130 
131 
132 class Cluster:
133     """Compute enviroment independend compute cluster."""
134 
135     def __init__(
136         self,
137         adapter: Adapter,
138         ssh_credentials: SSHCredentials,
139         instances: List[Instance],
140         timeout: int,
141     ):
142         """Create cluster.
143 
144         Args:
145             timeout (int): Minutes after which the machine will shutdown and terminate.
146                            This is a safety mechanism to avoid run aways cost. The user should still
147                            make sure to monitor the progress in case this mechanism fails.
148 
149         Example:
150             To start two groups of instances where the first one has one instance and the
151             second one has two instances you might define the following list of instances:
152 
153             instances = [
154                 Instance(name='server', group='server', num_cpu=2, num_ram=1.0),
155                 Instance(name='client_0', group='clients', num_cpu=4, num_ram=16.0),
156                 Instance(name='client_1', group='clients', num_cpu=4, num_ram=16.0),
157             ]
158 
159             Depending on the adapter used not every combination of vCPU and RAM might be available.
160         """
161         instance_names = {ins.name for ins in instances}
162         assert len(instance_names) == len(instances), "Instance names must be unique."
163 
164         self.adapter = adapter
165         self.ssh_credentials = ssh_credentials
166         self.instances = instances
167         self.timeout = timeout
168 
169     def get_instance(self, instance_name: str) -> Instance:
170         """Return instance by instance_name."""
171         for ins in self.instances:
172             if ins.name == instance_name:
173                 return ins
174 
175         raise InstanceIdNotFound()
176 
177     def get_instance_names(self, groups: Optional[List[str]] = None) -> List[str]:
178         """Return a list of all instance names."""
179         return [
180             ins.name for ins in self.instances if groups is None or ins.group in groups
181         ]
182 
183     def start(self) -> None:
184         """Start the instance."""
185         instance_groups = group_instances_by_specs(self.instances)
186 
187         with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
188             futures = [
189                 executor.submit(
190                     create_instances, self.adapter, instance_group, self.timeout
191                 )
192                 for instance_group in instance_groups
193             ]
194             concurrent.futures.wait(futures)
195 
196             try:
197                 for future in futures:
198                     future.result()
199             # pylint: disable=broad-except
200             except Exception as exc:
201                 log(
202                     ERROR,
203                     "Failed to start the cluster completely. Shutting down...",
204                 )
205                 log(ERROR, exc)
206 
207                 for future in futures:
208                     future.cancel()
209 
210                 self.terminate()
211                 raise StartFailed() from exc
212 
213         for ins in self.instances:
214             log(DEBUG, ins)
215 
216     def terminate(self) -> None:
217         """Terminate all instances and shutdown cluster."""
218         self.adapter.terminate_all_instances()
219 
220     def upload(
221         self, instance_name: str, local_path: str, remote_path: str
222     ) -> SFTPAttributes:
223         """Upload local file to remote instance."""
224         instance = self.get_instance(instance_name)
225 
226         with ssh_connection(instance, self.ssh_credentials) as client:
227             sftp = client.open_sftp()
228             sftp_file_attributes = sftp.put(local_path, remote_path)
229 
230         return sftp_file_attributes
231 
232     def upload_all(
233         self, local_path: str, remote_path: str
234     ) -> Dict[str, SFTPAttributes]:
235         """Upload file to all instances."""
236         results: Dict[str, SFTPAttributes] = {}
237 
238         with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
239             # Start the load operations and mark each future with its URL
240             future_to_result = {
241                 executor.submit(
242                     self.upload, instance_name, local_path, remote_path
243                 ): instance_name
244                 for instance_name in self.get_instance_names()
245             }
246 
247             for future in concurrent.futures.as_completed(future_to_result):
248                 instance_name = future_to_result[future]
249                 try:
250                     results[instance_name] = future.result()
251                 # pylint: disable=broad-except
252                 except Exception as exc:
253                     log(ERROR, (instance_name, exc))
254 
255         return results
256 
257     def exec(self, instance_name: str, command: str) -> ExecInfo:
258         """Run command on instance and return stdout."""
259         log(DEBUG, "Exec on %s: %s", instance_name, command)
260 
261         instance = self.get_instance(instance_name)
262 
263         with ssh_connection(instance, self.ssh_credentials) as client:
264             _, stdout, stderr = client.exec_command(command)
265             stdout = stdout.readlines()
266             stderr = stderr.readlines()
267 
268         print(stdout, stderr)
269 
270         return stdout, stderr
271 
272     def exec_all(
273         self, command: str, groups: Optional[List[str]] = None
274     ) -> Dict[str, ExecInfo]:
275         """Run command on all instances.
276 
277         If provided filter by group.
278         """
279         instance_names = self.get_instance_names(groups)
280 
281         results: Dict[str, ExecInfo] = {}
282 
283         with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
284             # Start the load operations and mark each future with its URL
285             future_to_result = {
286                 executor.submit(self.exec, instance_name, command): instance_name
287                 for instance_name in instance_names
288             }
289 
290             for future in concurrent.futures.as_completed(future_to_result):
291                 instance_name = future_to_result[future]
292                 try:
293                     results[instance_name] = future.result()
294                 # pylint: disable=broad-except
295                 except Exception as exc:
296                     log(ERROR, (instance_name, exc))
297 
298         return results
299
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/src/py/flwr_experimental/ops/cluster.py b/src/py/flwr_experimental/ops/cluster.py
--- a/src/py/flwr_experimental/ops/cluster.py
+++ b/src/py/flwr_experimental/ops/cluster.py
@@ -17,9 +17,9 @@
 from contextlib import contextmanager
 from itertools import groupby
 from logging import DEBUG, ERROR
-from typing import Dict, Iterator, List, Optional, Tuple
+from typing import Dict, Iterator, List, Optional, Tuple, Type, Union, cast
 
-from paramiko.client import SSHClient
+from paramiko.client import MissingHostKeyPolicy, SSHClient
 from paramiko.sftp_attr import SFTPAttributes
 
 from flwr.common.logger import log
@@ -27,7 +27,7 @@
 from .compute.adapter import Adapter
 from .instance import Instance
 
-ExecInfo = Tuple[str, str]
+ExecInfo = Tuple[List[str], List[str]]
 
 
 class StartFailed(Exception):
@@ -70,11 +70,18 @@
     """Connect to server and yield SSH client."""
     username, key_filename = ssh_credentials
 
+    instance_ssh_port: int = cast(int, instance.ssh_port)
+    ignore_host_key_policy: Union[
+        Type[MissingHostKeyPolicy], MissingHostKeyPolicy
+    ] = cast(
+        Union[Type[MissingHostKeyPolicy], MissingHostKeyPolicy], IgnoreHostKeyPolicy
+    )
+
     client = SSHClient()
-    client.set_missing_host_key_policy(IgnoreHostKeyPolicy)
+    client.set_missing_host_key_policy(ignore_host_key_policy)
     client.connect(
-        hostname=instance.public_ip,
-        port=instance.ssh_port,
+        hostname=str(instance.public_ip),
+        port=instance_ssh_port,
         username=username,
         key_filename=key_filename,
     )
@@ -225,7 +232,9 @@
 
         with ssh_connection(instance, self.ssh_credentials) as client:
             sftp = client.open_sftp()
-            sftp_file_attributes = sftp.put(local_path, remote_path)
+
+            if sftp is not None:
+                sftp_file_attributes = sftp.put(local_path, remote_path)
 
         return sftp_file_attributes
 
@@ -262,12 +271,12 @@
 
         with ssh_connection(instance, self.ssh_credentials) as client:
             _, stdout, stderr = client.exec_command(command)
-            stdout = stdout.readlines()
-            stderr = stderr.readlines()
+            lines_stdout = stdout.readlines()
+            lines_stderr = stderr.readlines()
 
-        print(stdout, stderr)
+        print(lines_stdout, lines_stderr)
 
-        return stdout, stderr
+        return lines_stdout, lines_stderr
 
     def exec_all(
         self, command: str, groups: Optional[List[str]] = None

{"golden_diff": "diff --git a/src/py/flwr_experimental/ops/cluster.py b/src/py/flwr_experimental/ops/cluster.py\n--- a/src/py/flwr_experimental/ops/cluster.py\n+++ b/src/py/flwr_experimental/ops/cluster.py\n@@ -17,9 +17,9 @@\n from contextlib import contextmanager\n from itertools import groupby\n from logging import DEBUG, ERROR\n-from typing import Dict, Iterator, List, Optional, Tuple\n+from typing import Dict, Iterator, List, Optional, Tuple, Type, Union, cast\n \n-from paramiko.client import SSHClient\n+from paramiko.client import MissingHostKeyPolicy, SSHClient\n from paramiko.sftp_attr import SFTPAttributes\n \n from flwr.common.logger import log\n@@ -27,7 +27,7 @@\n from .compute.adapter import Adapter\n from .instance import Instance\n \n-ExecInfo = Tuple[str, str]\n+ExecInfo = Tuple[List[str], List[str]]\n \n \n class StartFailed(Exception):\n@@ -70,11 +70,18 @@\n     \"\"\"Connect to server and yield SSH client.\"\"\"\n     username, key_filename = ssh_credentials\n \n+    instance_ssh_port: int = cast(int, instance.ssh_port)\n+    ignore_host_key_policy: Union[\n+        Type[MissingHostKeyPolicy], MissingHostKeyPolicy\n+    ] = cast(\n+        Union[Type[MissingHostKeyPolicy], MissingHostKeyPolicy], IgnoreHostKeyPolicy\n+    )\n+\n     client = SSHClient()\n-    client.set_missing_host_key_policy(IgnoreHostKeyPolicy)\n+    client.set_missing_host_key_policy(ignore_host_key_policy)\n     client.connect(\n-        hostname=instance.public_ip,\n-        port=instance.ssh_port,\n+        hostname=str(instance.public_ip),\n+        port=instance_ssh_port,\n         username=username,\n         key_filename=key_filename,\n     )\n@@ -225,7 +232,9 @@\n \n         with ssh_connection(instance, self.ssh_credentials) as client:\n             sftp = client.open_sftp()\n-            sftp_file_attributes = sftp.put(local_path, remote_path)\n+\n+            if sftp is not None:\n+                sftp_file_attributes = sftp.put(local_path, remote_path)\n \n         return sftp_file_attributes\n \n@@ -262,12 +271,12 @@\n \n         with ssh_connection(instance, self.ssh_credentials) as client:\n             _, stdout, stderr = client.exec_command(command)\n-            stdout = stdout.readlines()\n-            stderr = stderr.readlines()\n+            lines_stdout = stdout.readlines()\n+            lines_stderr = stderr.readlines()\n \n-        print(stdout, stderr)\n+        print(lines_stdout, lines_stderr)\n \n-        return stdout, stderr\n+        return lines_stdout, lines_stderr\n \n     def exec_all(\n         self, command: str, groups: Optional[List[str]] = None\n", "issue": "Update mypy to 0.810\n\n", "before_files": [{"content": "# Copyright 2020 Adap GmbH. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n\"\"\"Implements compute classes for EC2.\"\"\"\nimport concurrent.futures\nfrom contextlib import contextmanager\nfrom itertools import groupby\nfrom logging import DEBUG, ERROR\nfrom typing import Dict, Iterator, List, Optional, Tuple\n\nfrom paramiko.client import SSHClient\nfrom paramiko.sftp_attr import SFTPAttributes\n\nfrom flwr.common.logger import log\n\nfrom .compute.adapter import Adapter\nfrom .instance import Instance\n\nExecInfo = Tuple[str, str]\n\n\nclass StartFailed(Exception):\n    \"\"\"Raised when cluster could not start.\"\"\"\n\n\nclass InstanceIdNotFound(Exception):\n    \"\"\"Raised when there was no instance with given id.\"\"\"\n\n\nclass InstanceMismatch(Exception):\n    \"\"\"Raised when instances passed to create_instances do not have the same\n    values for RAM or CPU.\"\"\"\n\n\nclass IgnoreHostKeyPolicy:\n    \"\"\"Policy for accepting any unknown host key.\n\n    This is used by `paramiko.client.SSHClient`.\n    \"\"\"\n\n    # pylint: disable=no-self-use, unused-argument\n    def missing_host_key(self, client: SSHClient, hostname: str, key: str) -> None:\n        \"\"\"Simply return to ignore the host key.\n\n        As we create and destroy machines quite regularly and don't\n        reuse them we will not store the host key in the local system to\n        avoid pollution the local known_hosts file.\n        \"\"\"\n        return None\n\n\nSSHCredentials = Tuple[str, str]  # username, key_filename\n\n\n@contextmanager\ndef ssh_connection(\n    instance: Instance, ssh_credentials: SSHCredentials\n) -> Iterator[SSHClient]:\n    \"\"\"Connect to server and yield SSH client.\"\"\"\n    username, key_filename = ssh_credentials\n\n    client = SSHClient()\n    client.set_missing_host_key_policy(IgnoreHostKeyPolicy)\n    client.connect(\n        hostname=instance.public_ip,\n        port=instance.ssh_port,\n        username=username,\n        key_filename=key_filename,\n    )\n\n    yield client\n\n    client.close()\n\n\ndef create_instances(adapter: Adapter, instances: List[Instance], timeout: int) -> None:\n    \"\"\"Start instances and set props of each instance.\n\n    Fails if CPU and RAM of instances are not all the same.\n    \"\"\"\n    if not all(\n        [\n            ins.num_cpu == instances[0].num_cpu and ins.num_ram == instances[0].num_ram\n            for ins in instances\n        ]\n    ):\n        raise InstanceMismatch(\n            \"Values of num_cpu and num_ram have to be equal for all instances.\"\n        )\n\n    # As checked before that each instance has the same num_cpu and num_ram\n    # we can just take the values from the first => instances[0]\n    adapter_instances = adapter.create_instances(\n        num_cpu=instances[0].num_cpu,\n        num_ram=instances[0].num_ram,\n        num_instance=len(instances),\n        gpu=instances[0].gpu,\n        timeout=timeout,\n    )\n\n    for i, adp_ins in enumerate(adapter_instances):\n        instance_id, private_ip, public_ip, ssh_port, state = adp_ins\n\n        instances[i].instance_id = instance_id\n        instances[i].private_ip = private_ip\n        instances[i].public_ip = public_ip\n        instances[i].ssh_port = ssh_port\n        instances[i].state = state\n\n\ndef group_instances_by_specs(instances: List[Instance]) -> List[List[Instance]]:\n    \"\"\"Group instances by num_cpu and num_ram.\"\"\"\n    groups: List[List[Instance]] = []\n    keyfunc = lambda ins: f\"{ins.num_cpu}-{ins.num_ram}\"\n    instances = sorted(instances, key=keyfunc)\n    for _, group in groupby(instances, keyfunc):\n        groups.append(list(group))\n    return groups\n\n\nclass Cluster:\n    \"\"\"Compute enviroment independend compute cluster.\"\"\"\n\n    def __init__(\n        self,\n        adapter: Adapter,\n        ssh_credentials: SSHCredentials,\n        instances: List[Instance],\n        timeout: int,\n    ):\n        \"\"\"Create cluster.\n\n        Args:\n            timeout (int): Minutes after which the machine will shutdown and terminate.\n                           This is a safety mechanism to avoid run aways cost. The user should still\n                           make sure to monitor the progress in case this mechanism fails.\n\n        Example:\n            To start two groups of instances where the first one has one instance and the\n            second one has two instances you might define the following list of instances:\n\n            instances = [\n                Instance(name='server', group='server', num_cpu=2, num_ram=1.0),\n                Instance(name='client_0', group='clients', num_cpu=4, num_ram=16.0),\n                Instance(name='client_1', group='clients', num_cpu=4, num_ram=16.0),\n            ]\n\n            Depending on the adapter used not every combination of vCPU and RAM might be available.\n        \"\"\"\n        instance_names = {ins.name for ins in instances}\n        assert len(instance_names) == len(instances), \"Instance names must be unique.\"\n\n        self.adapter = adapter\n        self.ssh_credentials = ssh_credentials\n        self.instances = instances\n        self.timeout = timeout\n\n    def get_instance(self, instance_name: str) -> Instance:\n        \"\"\"Return instance by instance_name.\"\"\"\n        for ins in self.instances:\n            if ins.name == instance_name:\n                return ins\n\n        raise InstanceIdNotFound()\n\n    def get_instance_names(self, groups: Optional[List[str]] = None) -> List[str]:\n        \"\"\"Return a list of all instance names.\"\"\"\n        return [\n            ins.name for ins in self.instances if groups is None or ins.group in groups\n        ]\n\n    def start(self) -> None:\n        \"\"\"Start the instance.\"\"\"\n        instance_groups = group_instances_by_specs(self.instances)\n\n        with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:\n            futures = [\n                executor.submit(\n                    create_instances, self.adapter, instance_group, self.timeout\n                )\n                for instance_group in instance_groups\n            ]\n            concurrent.futures.wait(futures)\n\n            try:\n                for future in futures:\n                    future.result()\n            # pylint: disable=broad-except\n            except Exception as exc:\n                log(\n                    ERROR,\n                    \"Failed to start the cluster completely. Shutting down...\",\n                )\n                log(ERROR, exc)\n\n                for future in futures:\n                    future.cancel()\n\n                self.terminate()\n                raise StartFailed() from exc\n\n        for ins in self.instances:\n            log(DEBUG, ins)\n\n    def terminate(self) -> None:\n        \"\"\"Terminate all instances and shutdown cluster.\"\"\"\n        self.adapter.terminate_all_instances()\n\n    def upload(\n        self, instance_name: str, local_path: str, remote_path: str\n    ) -> SFTPAttributes:\n        \"\"\"Upload local file to remote instance.\"\"\"\n        instance = self.get_instance(instance_name)\n\n        with ssh_connection(instance, self.ssh_credentials) as client:\n            sftp = client.open_sftp()\n            sftp_file_attributes = sftp.put(local_path, remote_path)\n\n        return sftp_file_attributes\n\n    def upload_all(\n        self, local_path: str, remote_path: str\n    ) -> Dict[str, SFTPAttributes]:\n        \"\"\"Upload file to all instances.\"\"\"\n        results: Dict[str, SFTPAttributes] = {}\n\n        with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:\n            # Start the load operations and mark each future with its URL\n            future_to_result = {\n                executor.submit(\n                    self.upload, instance_name, local_path, remote_path\n                ): instance_name\n                for instance_name in self.get_instance_names()\n            }\n\n            for future in concurrent.futures.as_completed(future_to_result):\n                instance_name = future_to_result[future]\n                try:\n                    results[instance_name] = future.result()\n                # pylint: disable=broad-except\n                except Exception as exc:\n                    log(ERROR, (instance_name, exc))\n\n        return results\n\n    def exec(self, instance_name: str, command: str) -> ExecInfo:\n        \"\"\"Run command on instance and return stdout.\"\"\"\n        log(DEBUG, \"Exec on %s: %s\", instance_name, command)\n\n        instance = self.get_instance(instance_name)\n\n        with ssh_connection(instance, self.ssh_credentials) as client:\n            _, stdout, stderr = client.exec_command(command)\n            stdout = stdout.readlines()\n            stderr = stderr.readlines()\n\n        print(stdout, stderr)\n\n        return stdout, stderr\n\n    def exec_all(\n        self, command: str, groups: Optional[List[str]] = None\n    ) -> Dict[str, ExecInfo]:\n        \"\"\"Run command on all instances.\n\n        If provided filter by group.\n        \"\"\"\n        instance_names = self.get_instance_names(groups)\n\n        results: Dict[str, ExecInfo] = {}\n\n        with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:\n            # Start the load operations and mark each future with its URL\n            future_to_result = {\n                executor.submit(self.exec, instance_name, command): instance_name\n                for instance_name in instance_names\n            }\n\n            for future in concurrent.futures.as_completed(future_to_result):\n                instance_name = future_to_result[future]\n                try:\n                    results[instance_name] = future.result()\n                # pylint: disable=broad-except\n                except Exception as exc:\n                    log(ERROR, (instance_name, exc))\n\n        return results\n", "path": "src/py/flwr_experimental/ops/cluster.py"}], "after_files": [{"content": "# Copyright 2020 Adap GmbH. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n\"\"\"Implements compute classes for EC2.\"\"\"\nimport concurrent.futures\nfrom contextlib import contextmanager\nfrom itertools import groupby\nfrom logging import DEBUG, ERROR\nfrom typing import Dict, Iterator, List, Optional, Tuple, Type, Union, cast\n\nfrom paramiko.client import MissingHostKeyPolicy, SSHClient\nfrom paramiko.sftp_attr import SFTPAttributes\n\nfrom flwr.common.logger import log\n\nfrom .compute.adapter import Adapter\nfrom .instance import Instance\n\nExecInfo = Tuple[List[str], List[str]]\n\n\nclass StartFailed(Exception):\n    \"\"\"Raised when cluster could not start.\"\"\"\n\n\nclass InstanceIdNotFound(Exception):\n    \"\"\"Raised when there was no instance with given id.\"\"\"\n\n\nclass InstanceMismatch(Exception):\n    \"\"\"Raised when instances passed to create_instances do not have the same\n    values for RAM or CPU.\"\"\"\n\n\nclass IgnoreHostKeyPolicy:\n    \"\"\"Policy for accepting any unknown host key.\n\n    This is used by `paramiko.client.SSHClient`.\n    \"\"\"\n\n    # pylint: disable=no-self-use, unused-argument\n    def missing_host_key(self, client: SSHClient, hostname: str, key: str) -> None:\n        \"\"\"Simply return to ignore the host key.\n\n        As we create and destroy machines quite regularly and don't\n        reuse them we will not store the host key in the local system to\n        avoid pollution the local known_hosts file.\n        \"\"\"\n        return None\n\n\nSSHCredentials = Tuple[str, str]  # username, key_filename\n\n\n@contextmanager\ndef ssh_connection(\n    instance: Instance, ssh_credentials: SSHCredentials\n) -> Iterator[SSHClient]:\n    \"\"\"Connect to server and yield SSH client.\"\"\"\n    username, key_filename = ssh_credentials\n\n    instance_ssh_port: int = cast(int, instance.ssh_port)\n    ignore_host_key_policy: Union[\n        Type[MissingHostKeyPolicy], MissingHostKeyPolicy\n    ] = cast(\n        Union[Type[MissingHostKeyPolicy], MissingHostKeyPolicy], IgnoreHostKeyPolicy\n    )\n\n    client = SSHClient()\n    client.set_missing_host_key_policy(ignore_host_key_policy)\n    client.connect(\n        hostname=str(instance.public_ip),\n        port=instance_ssh_port,\n        username=username,\n        key_filename=key_filename,\n    )\n\n    yield client\n\n    client.close()\n\n\ndef create_instances(adapter: Adapter, instances: List[Instance], timeout: int) -> None:\n    \"\"\"Start instances and set props of each instance.\n\n    Fails if CPU and RAM of instances are not all the same.\n    \"\"\"\n    if not all(\n        [\n            ins.num_cpu == instances[0].num_cpu and ins.num_ram == instances[0].num_ram\n            for ins in instances\n        ]\n    ):\n        raise InstanceMismatch(\n            \"Values of num_cpu and num_ram have to be equal for all instances.\"\n        )\n\n    # As checked before that each instance has the same num_cpu and num_ram\n    # we can just take the values from the first => instances[0]\n    adapter_instances = adapter.create_instances(\n        num_cpu=instances[0].num_cpu,\n        num_ram=instances[0].num_ram,\n        num_instance=len(instances),\n        gpu=instances[0].gpu,\n        timeout=timeout,\n    )\n\n    for i, adp_ins in enumerate(adapter_instances):\n        instance_id, private_ip, public_ip, ssh_port, state = adp_ins\n\n        instances[i].instance_id = instance_id\n        instances[i].private_ip = private_ip\n        instances[i].public_ip = public_ip\n        instances[i].ssh_port = ssh_port\n        instances[i].state = state\n\n\ndef group_instances_by_specs(instances: List[Instance]) -> List[List[Instance]]:\n    \"\"\"Group instances by num_cpu and num_ram.\"\"\"\n    groups: List[List[Instance]] = []\n    keyfunc = lambda ins: f\"{ins.num_cpu}-{ins.num_ram}\"\n    instances = sorted(instances, key=keyfunc)\n    for _, group in groupby(instances, keyfunc):\n        groups.append(list(group))\n    return groups\n\n\nclass Cluster:\n    \"\"\"Compute enviroment independend compute cluster.\"\"\"\n\n    def __init__(\n        self,\n        adapter: Adapter,\n        ssh_credentials: SSHCredentials,\n        instances: List[Instance],\n        timeout: int,\n    ):\n        \"\"\"Create cluster.\n\n        Args:\n            timeout (int): Minutes after which the machine will shutdown and terminate.\n                           This is a safety mechanism to avoid run aways cost. The user should still\n                           make sure to monitor the progress in case this mechanism fails.\n\n        Example:\n            To start two groups of instances where the first one has one instance and the\n            second one has two instances you might define the following list of instances:\n\n            instances = [\n                Instance(name='server', group='server', num_cpu=2, num_ram=1.0),\n                Instance(name='client_0', group='clients', num_cpu=4, num_ram=16.0),\n                Instance(name='client_1', group='clients', num_cpu=4, num_ram=16.0),\n            ]\n\n            Depending on the adapter used not every combination of vCPU and RAM might be available.\n        \"\"\"\n        instance_names = {ins.name for ins in instances}\n        assert len(instance_names) == len(instances), \"Instance names must be unique.\"\n\n        self.adapter = adapter\n        self.ssh_credentials = ssh_credentials\n        self.instances = instances\n        self.timeout = timeout\n\n    def get_instance(self, instance_name: str) -> Instance:\n        \"\"\"Return instance by instance_name.\"\"\"\n        for ins in self.instances:\n            if ins.name == instance_name:\n                return ins\n\n        raise InstanceIdNotFound()\n\n    def get_instance_names(self, groups: Optional[List[str]] = None) -> List[str]:\n        \"\"\"Return a list of all instance names.\"\"\"\n        return [\n            ins.name for ins in self.instances if groups is None or ins.group in groups\n        ]\n\n    def start(self) -> None:\n        \"\"\"Start the instance.\"\"\"\n        instance_groups = group_instances_by_specs(self.instances)\n\n        with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:\n            futures = [\n                executor.submit(\n                    create_instances, self.adapter, instance_group, self.timeout\n                )\n                for instance_group in instance_groups\n            ]\n            concurrent.futures.wait(futures)\n\n            try:\n                for future in futures:\n                    future.result()\n            # pylint: disable=broad-except\n            except Exception as exc:\n                log(\n                    ERROR,\n                    \"Failed to start the cluster completely. Shutting down...\",\n                )\n                log(ERROR, exc)\n\n                for future in futures:\n                    future.cancel()\n\n                self.terminate()\n                raise StartFailed() from exc\n\n        for ins in self.instances:\n            log(DEBUG, ins)\n\n    def terminate(self) -> None:\n        \"\"\"Terminate all instances and shutdown cluster.\"\"\"\n        self.adapter.terminate_all_instances()\n\n    def upload(\n        self, instance_name: str, local_path: str, remote_path: str\n    ) -> SFTPAttributes:\n        \"\"\"Upload local file to remote instance.\"\"\"\n        instance = self.get_instance(instance_name)\n\n        with ssh_connection(instance, self.ssh_credentials) as client:\n            sftp = client.open_sftp()\n\n            if sftp is not None:\n                sftp_file_attributes = sftp.put(local_path, remote_path)\n\n        return sftp_file_attributes\n\n    def upload_all(\n        self, local_path: str, remote_path: str\n    ) -> Dict[str, SFTPAttributes]:\n        \"\"\"Upload file to all instances.\"\"\"\n        results: Dict[str, SFTPAttributes] = {}\n\n        with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:\n            # Start the load operations and mark each future with its URL\n            future_to_result = {\n                executor.submit(\n                    self.upload, instance_name, local_path, remote_path\n                ): instance_name\n                for instance_name in self.get_instance_names()\n            }\n\n            for future in concurrent.futures.as_completed(future_to_result):\n                instance_name = future_to_result[future]\n                try:\n                    results[instance_name] = future.result()\n                # pylint: disable=broad-except\n                except Exception as exc:\n                    log(ERROR, (instance_name, exc))\n\n        return results\n\n    def exec(self, instance_name: str, command: str) -> ExecInfo:\n        \"\"\"Run command on instance and return stdout.\"\"\"\n        log(DEBUG, \"Exec on %s: %s\", instance_name, command)\n\n        instance = self.get_instance(instance_name)\n\n        with ssh_connection(instance, self.ssh_credentials) as client:\n            _, stdout, stderr = client.exec_command(command)\n            lines_stdout = stdout.readlines()\n            lines_stderr = stderr.readlines()\n\n        print(lines_stdout, lines_stderr)\n\n        return lines_stdout, lines_stderr\n\n    def exec_all(\n        self, command: str, groups: Optional[List[str]] = None\n    ) -> Dict[str, ExecInfo]:\n        \"\"\"Run command on all instances.\n\n        If provided filter by group.\n        \"\"\"\n        instance_names = self.get_instance_names(groups)\n\n        results: Dict[str, ExecInfo] = {}\n\n        with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:\n            # Start the load operations and mark each future with its URL\n            future_to_result = {\n                executor.submit(self.exec, instance_name, command): instance_name\n                for instance_name in instance_names\n            }\n\n            for future in concurrent.futures.as_completed(future_to_result):\n                instance_name = future_to_result[future]\n                try:\n                    results[instance_name] = future.result()\n                # pylint: disable=broad-except\n                except Exception as exc:\n                    log(ERROR, (instance_name, exc))\n\n        return results\n", "path": "src/py/flwr_experimental/ops/cluster.py"}]}