Skip to content

Commit d132b03

Browse files
authored
Clean tuples dict keys from workers_info in /api/v1/retire_workers. (#8996)
* Clean tuples dict keys from workers_info in /api/v1/retire_workers. * Format * Use `recursive_to_dict`.
1 parent 2d19d1c commit d132b03

File tree

2 files changed

+49
-1
lines changed

2 files changed

+49
-1
lines changed

distributed/http/scheduler/api.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import json
44

55
from distributed.http.utils import RequestHandler
6+
from distributed.utils import recursive_to_dict
67

78

89
class APIHandler(RequestHandler):
@@ -24,7 +25,7 @@ async def post(self):
2425
else:
2526
workers = params.get("workers", {})
2627
workers_info = await scheduler.retire_workers(workers=workers)
27-
self.write(json.dumps(workers_info))
28+
self.write(json.dumps(recursive_to_dict(workers_info)))
2829
except Exception as e:
2930
self.set_status(500, str(e))
3031
self.write(json.dumps({"Error": "Internal Server Error"}))

distributed/http/scheduler/tests/test_scheduler_http.py

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -718,6 +718,53 @@ async def test_retire_workers(c, s, a, b):
718718
assert len(retired_workers_info) == 2
719719

720720

721+
@gen_cluster(
722+
client=True,
723+
clean_kwargs={"threads": False},
724+
config={
725+
"distributed.scheduler.http.routes": DEFAULT_ROUTES
726+
+ ["distributed.http.scheduler.api"]
727+
},
728+
)
729+
async def test_retire_workers_with_tuple_keys(c, s, a, b):
730+
aiohttp = pytest.importorskip("aiohttp")
731+
732+
worker_address = "tcp://172.17.0.3:39571"
733+
734+
async def mock_retire_workers(*args, **kwargs):
735+
# Return problematic data structure
736+
# tuple are not json serializable
737+
return {
738+
worker_address: {
739+
"type": "Worker",
740+
"metrics": {
741+
"digests_total_since_heartbeat": {
742+
(
743+
"execute",
744+
"slowadd",
745+
"thread-cpu",
746+
"seconds",
747+
): 0.0003396660000000093,
748+
},
749+
},
750+
}
751+
}
752+
753+
# Replace the method with our mock
754+
s.retire_workers = mock_retire_workers
755+
756+
async with aiohttp.ClientSession() as session:
757+
params = {"workers": [a.address, b.address]}
758+
async with session.post(
759+
"http://localhost:%d/api/v1/retire_workers" % s.http_server.port,
760+
json=params,
761+
) as resp:
762+
assert resp.status == 200
763+
assert resp.headers["Content-Type"] == "application/json"
764+
retired_workers_info = json.loads(await resp.text())
765+
assert worker_address in retired_workers_info
766+
767+
721768
@gen_cluster(
722769
client=True,
723770
clean_kwargs={"threads": False},

0 commit comments

Comments
 (0)