feat: write nginx config

apps/hiopy/loco.py

@@ -27,3 +27,31 @@ def expose_store_over_http(store, host, port):
     t = Thread(target=_run_server, args=(runner, host, port),
                daemon=True)
     t.start()
+
+
+def write_nginx_config(distributed_data_vars, port, rank):
+    import platform
+    from regex_engine import generator
+
+    with open("backend_map.conf", "w") as f:
+        f.write("map $request_uri $backend {\n")
+        for r, data_vars in enumerate(distributed_data_vars):
+            for group, vgroup, chunk_slice, data_vars in data_vars:
+                for v in data_vars:
+                    if chunk_slice == slice(None):
+                        cell_regex = r"\d+"
+                    else:
+                        cell_chunk_size = dict(zip(v.attrs["_ARRAY_DIMENSIONS"], v.chunks))["cell"]
+                        chunk_range = (chunk_slice.start//cell_chunk_size, chunk_slice.stop//cell_chunk_size)
+                        cell_regex = generator().numerical_range(*chunk_range)[1:-1]
+                    f.write(f"  \"~^/{group}/{v.name}/\\d+\\.\\d+\\.{cell_regex}\" loco_backend{r};\n")
+        f.write("}\n\n")
+        for i in range(len(distributed_data_vars)):
+            f.write(f"include \"loco_backend{i}.conf\";\n")
+
+    with open(f"loco_backend{rank}.conf", "w") as f:
+        f.write(f"""
+upstream loco_backend{rank} {{
+    server {platform.node()}:{port};
+}}
+""")

apps/hiopy/loco_nginx.conf

+error_log nginx_error.log;
+pid loco_nginx.pid;
+
+worker_processes  1;
+daemon off;
+
+events {
+  worker_connections 1024;
+}
+
+http {
+     access_log nginx_access.log;
+     include backend_map.conf; # <- this file is generated by the hiopy-workers
+     server {
+            listen 8666 default_server;
+            listen [::]:8666 default_server;
+            server_name _;
+
+
+            location / {
+                     proxy_pass http://$backend;
+                     proxy_set_header Host $host;
+                     proxy_set_header X-Real-IP $remote_addr;
+                     proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
+                     proxy_set_header X-Forwarded-Proto $scheme;
+            }
+     }
+}
\ No newline at end of file

apps/hiopy/worker.py

@@ -8,7 +8,7 @@ import logging
 from itertools import chain, groupby
 from _distribute_work import distribute_work
 from _data_handler import DataHandler
-from loco import expose_store_over_http
+from loco import expose_store_over_http, write_nginx_config
 
 parser = ArgumentParser()
 # zarr.copy_store (used with loco) does not work with zarr.open_consolidated
@@ -65,6 +65,10 @@ grouped_data_vars = {grid: list(variables)
                      }
 distributed_data_vars = distribute_work(grouped_data_vars,
                                         group_comm_size())
+
+if args.loco:
+    write_nginx_config(distributed_data_vars, args.loco_port, group_comm_rank())
+
 my_data_vars = distributed_data_vars[group_comm_rank()]
 # my_data_vars list of tuples: ((group, slice), list of variables)