earthkit-workflows 0.3.4__py3-none-any.whl → 0.3.6__py3-none-any.whl

cascade/benchmarks/__main__.py

@@ -23,10 +23,10 @@ Make sure you correctly configure:
 
 import logging
 import logging.config
+import multiprocessing
 import os
 import subprocess
 from concurrent.futures import ThreadPoolExecutor
-from multiprocessing import Process
 from time import perf_counter_ns
 
 import fire
@@ -36,7 +36,7 @@ import cascade.low.into
 from cascade.controller.impl import run
 from cascade.executor.bridge import Bridge
 from cascade.executor.comms import callback
-from cascade.executor.config import logging_config
+from cascade.executor.config import logging_config, logging_config_filehandler
 from cascade.executor.executor import Executor
 from cascade.executor.msg import BackboneAddress, ExecutorShutdown
 from cascade.low.core import JobInstance
@@ -105,8 +105,14 @@ def launch_executor(
     i: int,
     shm_vol_gb: int | None,
     gpu_count: int,
+    log_base: str | None,
 ):
-    logging.config.dictConfig(logging_config)
+    if log_base is not None:
+        log_base = f"{log_base}.host{i}"
+        log_path = f"{log_base}.txt"
+        logging.config.dictConfig(logging_config_filehandler(log_path))
+    else:
+        logging.config.dictConfig(logging_config)
     logger.info(f"will set {gpu_count} gpus on host {i}")
     os.environ["CASCADE_GPU_COUNT"] = str(gpu_count)
     executor = Executor(
@@ -116,6 +122,7 @@ def launch_executor(
         f"h{i}",
         portBase,
         shm_vol_gb,
+        log_base,
     )
     executor.register()
     executor.recv_loop()
@@ -126,9 +133,14 @@ def run_locally(
     hosts: int,
     workers: int,
     portBase: int = 12345,
+    log_base: str | None = None,
     report_address: str | None = None,
 ):
-    logging.config.dictConfig(logging_config)
+    if log_base is not None:
+        log_path = f"{log_base}.controller.txt"
+        logging.config.dictConfig(logging_config_filehandler(log_path))
+    else:
+        logging.config.dictConfig(logging_config)
     launch = perf_counter_ns()
     preschedule = precompute(job)
     c = f"tcp://localhost:{portBase}"
@@ -139,9 +151,10 @@ def run_locally(
             gpu_count = get_gpu_count()
         else:
             gpu_count = 0
-        p = Process(
+        # NOTE forkserver/spawn seem to forget venv, we need fork
+        p = multiprocessing.get_context("fork").Process(
             target=launch_executor,
-            args=(job, c, workers, portBase + 1 + i * 10, i, None, gpu_count),
+            args=(job, c, workers, portBase + 1 + i * 10, i, None, gpu_count, log_base),
         )
         p.start()
         ps.append(p)
@@ -171,6 +184,7 @@ def main_local(
     job: str | None = None,
     instance: str | None = None,
     port_base: int = 12345,
+    log_base: str | None = None,
 ) -> None:
     jobInstance = get_job(job, instance)
     run_locally(
@@ -179,6 +193,7 @@ def main_local(
         workers_per_host,
         report_address=report_address,
         portBase=port_base,
+        log_base=log_base,
     )
 
 

cascade/executor/executor.py

@@ -28,7 +28,7 @@ import cascade.shm.client as shm_client
 from cascade.executor.comms import GraceWatcher, Listener, ReliableSender, callback
 from cascade.executor.comms import default_message_resend_ms as resend_grace_ms
 from cascade.executor.comms import default_timeout_ms as comms_default_timeout_ms
-from cascade.executor.config import logging_config
+from cascade.executor.config import logging_config, logging_config_filehandler
 from cascade.executor.data_server import start_data_server
 from cascade.executor.msg import (
     Ack,
@@ -70,6 +70,7 @@ class Executor:
         host: HostId,
         portBase: int,
         shm_vol_gb: int | None = None,
+        log_base: str | None = None,
     ) -> None:
         self.job_instance = job_instance
         self.param_source = param_source(job_instance.edges)
@@ -78,6 +79,7 @@ class Executor:
         self.workers: dict[WorkerId, BaseProcess | None] = {
             WorkerId(host, f"w{i}"): None for i in range(workers)
         }
+        self.log_base = log_base
 
         self.datasets: set[DatasetId] = set()
         self.heartbeat_watcher = GraceWatcher(grace_ms=heartbeat_grace_ms)
@@ -92,17 +94,25 @@ class Executor:
         shm_port = portBase + 2
         shm_api.publish_client_port(shm_port)
         ctx = get_context("fork")
+        if log_base:
+            shm_logging = logging_config_filehandler(f"{log_base}.shm.txt")
+        else:
+            shm_logging = logging_config
         self.shm_process = ctx.Process(
             target=shm_server,
             args=(
                 shm_port,
                 shm_vol_gb * (1024**3) if shm_vol_gb else None,
-                logging_config,
+                shm_logging,
                 f"sCasc{host}",
             ),
         )
         self.shm_process.start()
         self.daddress = address_of(portBase + 1)
+        if log_base:
+            dsr_logging = logging_config_filehandler(f"{log_base}.dsr.txt")
+        else:
+            dsr_logging = logging_config
         self.data_server = ctx.Process(
             target=start_data_server,
             args=(
@@ -110,7 +120,7 @@ class Executor:
                 self.daddress,
                 self.host,
                 shm_port,
-                logging_config,
+                dsr_logging,
             ),
         )
         self.data_server.start()
@@ -181,6 +191,7 @@ class Executor:
                 job=self.job_instance,
                 param_source=self.param_source,
                 callback=self.mlistener.address,
+                log_base=self.log_base,
             )
             p = ctx.Process(target=entrypoint, kwargs={"runnerContext": runnerContext})
             p.start()
@@ -222,17 +233,17 @@ class Executor:
         procFail = lambda ex: ex is not None and ex != 0
         for k, e in self.workers.items():
             if e is None:
-                ValueError(f"process on {k} is not alive")
+                raise ValueError(f"process on {k} is not alive")
             elif procFail(e.exitcode):
-                ValueError(
+                raise ValueError(
                     f"process on {k} failed to terminate correctly: {e.pid} -> {e.exitcode}"
                 )
         if procFail(self.shm_process.exitcode):
-            ValueError(
+            raise ValueError(
                 f"shm server {self.shm_process.pid} failed with {self.shm_process.exitcode}"
             )
         if procFail(self.data_server.exitcode):
-            ValueError(
+            raise ValueError(
                 f"data server {self.data_server.pid} failed with {self.data_server.exitcode}"
             )
         if self.heartbeat_watcher.is_breach() > 0:

cascade/executor/runner/entrypoint.py

@@ -17,7 +17,7 @@ import zmq
 
 import cascade.executor.serde as serde
 from cascade.executor.comms import callback
-from cascade.executor.config import logging_config
+from cascade.executor.config import logging_config, logging_config_filehandler
 from cascade.executor.msg import (
     BackboneAddress,
     DatasetPublished,
@@ -44,6 +44,7 @@ class RunnerContext:
     job: JobInstance
     callback: BackboneAddress
     param_source: dict[TaskId, dict[int | str, DatasetId]]
+    log_base: str | None
 
     def project(self, taskSequence: TaskSequence) -> ExecutionContext:
         schema_lookup: dict[DatasetId, str] = {}
@@ -92,7 +93,11 @@ def execute_sequence(
 
 
 def entrypoint(runnerContext: RunnerContext):
-    logging.config.dictConfig(logging_config)
+    if runnerContext.log_base:
+        log_path = f"{runnerContext.log_base}.{runnerContext.workerId.worker}"
+        logging.config.dictConfig(logging_config_filehandler(log_path))
+    else:
+        logging.config.dictConfig(logging_config)
     ctx = zmq.Context()
     socket = ctx.socket(zmq.PULL)
     socket.bind(worker_address(runnerContext.workerId))

cascade/gateway/__main__.py

@@ -10,13 +10,17 @@ import logging.config
 
 import fire
 
-from cascade.executor.config import logging_config
+from cascade.executor.config import logging_config, logging_config_filehandler
 from cascade.gateway.server import serve
 
 
-def main(url: str) -> None:
-    logging.config.dictConfig(logging_config)
-    serve(url)
+def main(url: str, log_base: str | None = None) -> None:
+    if log_base:
+        log_path = f"{log_base}/gateway.txt"
+        logging.config.dictConfig(logging_config_filehandler(log_path))
+    else:
+        logging.config.dictConfig(logging_config)
+    serve(url, log_base)
 
 
 if __name__ == "__main__":

cascade/gateway/router.py

@@ -43,7 +43,9 @@ class Job:
 local_job_port = 12345
 
 
-def _spawn_local(job_spec: JobSpec, addr: str, job_id: JobId) -> subprocess.Popen:
+def _spawn_local(
+    job_spec: JobSpec, addr: str, job_id: JobId, log_base: str | None
+) -> subprocess.Popen:
     base = [
         "python",
         "-m",
@@ -68,11 +70,15 @@ def _spawn_local(job_spec: JobSpec, addr: str, job_id: JobId) -> subprocess.Pope
         f"{job_spec.hosts}",
     ]
     report = ["--report_address", f"{addr},{job_id}"]
+    if log_base:
+        logs = ["--log_base", f"{log_base}/job.{job_id}"]
+    else:
+        logs = []
     global local_job_port
     portBase = ["--port_base", str(local_job_port)]
     local_job_port += 1 + job_spec.hosts * job_spec.workers_per_host * 10
     return subprocess.Popen(
-        base + infra + report + portBase, env={**os.environ, **job_spec.envvars}
+        base + infra + report + portBase + logs, env={**os.environ, **job_spec.envvars}
     )
 
 
@@ -105,29 +111,40 @@ def _spawn_slurm(job_spec: JobSpec, addr: str, job_id: JobId) -> subprocess.Pope
     )
 
 
-def _spawn_subprocess(job_spec: JobSpec, addr: str, job_id: JobId) -> subprocess.Popen:
+def _spawn_subprocess(
+    job_spec: JobSpec, addr: str, job_id: JobId, log_base: str | None
+) -> subprocess.Popen:
     if job_spec.use_slurm:
+        if log_base is not None:
+            raise ValueError(f"unexpected {log_base=}")
         return _spawn_slurm(job_spec, addr, job_id)
     else:
-        return _spawn_local(job_spec, addr, job_id)
+        return _spawn_local(job_spec, addr, job_id, log_base)
 
 
 class JobRouter:
-    def __init__(self, poller: zmq.Poller):
+    def __init__(self, poller: zmq.Poller, log_base: str | None):
         self.poller = poller
         self.jobs: dict[str, Job] = {}
         self.procs: dict[str, subprocess.Popen] = {}
+        self.log_base = log_base
 
     def spawn_job(self, job_spec: JobSpec) -> JobId:
         job_id = next_uuid(self.jobs.keys(), lambda: str(uuid.uuid4()))
-        base_addr = f"tcp://{getfqdn()}"
+        if job_spec.use_slurm:
+            base_addr = f"tcp://{getfqdn()}"
+        else:
+            # NOTE on macos, it seems getfqdn does not give zmq-bindable addr
+            base_addr = "tcp://localhost"
         socket = get_context().socket(zmq.PULL)
         port = socket.bind_to_random_port(base_addr)
         full_addr = f"{base_addr}:{port}"
         logger.debug(f"will spawn job {job_id} and listen on {full_addr}")
         self.poller.register(socket, flags=zmq.POLLIN)
         self.jobs[job_id] = Job(socket, JobProgressStarted, -1, {})
-        self.procs[job_id] = _spawn_subprocess(job_spec, full_addr, job_id)
+        self.procs[job_id] = _spawn_subprocess(
+            job_spec, full_addr, job_id, self.log_base
+        )
         return job_id
 
     def progress_of(

cascade/gateway/server.py

@@ -79,14 +79,14 @@ def handle_controller(socket: zmq.Socket, jobs: JobRouter) -> None:
         jobs.put_result(report.job_id, dataset_id, result)
 
 
-def serve(url: str) -> None:
+def serve(url: str, log_base: str | None = None) -> None:
     ctx = get_context()
     poller = zmq.Poller()
 
     fe = ctx.socket(zmq.REP)
     fe.bind(url)
     poller.register(fe, flags=zmq.POLLIN)
-    jobs = JobRouter(poller)
+    jobs = JobRouter(poller, log_base)
 
     logger.debug("entering recv loop")
     is_break = False

earthkit/workflows/_version.py

@@ -1,2 +1,2 @@
 # Do not change! Do not track in version control!
-__version__ = "0.3.4"
+__version__ = "0.3.6"

earthkit/workflows/fluent.py

@@ -424,12 +424,8 @@ class Action:
             )
 
         # Applies operation to every node, keeping node array structure
-        new_nodes = xr.DataArray(
-            np.empty(self.nodes.shape, dtype=object),
-            coords=self.nodes.coords,
-            dims=self.nodes.dims,
-            attrs=self.nodes.attrs,
-        )
+        new_nodes = np.empty(self.nodes.shape, dtype=object)
+
         it = np.nditer(self.nodes, flags=["multi_index", "refs_ok"])
         node_payload = payload
         for node in it:
@@ -441,7 +437,14 @@ class Action:
                 num_outputs=len(yields[1]) if yields else 1,
             )
 
-        return type(self)(new_nodes, yields)
+        new_nodes_xr = xr.DataArray(
+            new_nodes,
+            coords=self.nodes.coords,
+            dims=self.nodes.dims,
+            attrs=self.nodes.attrs,
+        )
+
+        return type(self)(new_nodes_xr, yields)
 
     def reduce(
         self,

{earthkit_workflows-0.3.4.dist-info → earthkit_workflows-0.3.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: earthkit-workflows
-Version: 0.3.4
+Version: 0.3.6
 Summary: Earthkit Workflows is a Python library for declaring earthkit task DAGs, as well as scheduling and executing them on heterogeneous computing systems.
 Author-email: "European Centre for Medium-Range Weather Forecasts (ECMWF)" <software.support@ecmwf.int>
 License-Expression: Apache-2.0

{earthkit_workflows-0.3.4.dist-info → earthkit_workflows-0.3.6.dist-info}/RECORD

@@ -1,7 +1,7 @@
 cascade/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cascade/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 cascade/benchmarks/__init__.py,sha256=Gu8kEApmJ2zsIhT2zpm1-6n84-OwWnz-0vO8UHYtBzo,528
-cascade/benchmarks/__main__.py,sha256=M9MSOYdsOQZdXjB_0sfl5ltv0cuAatip_-e-Wi78FFQ,6897
+cascade/benchmarks/__main__.py,sha256=LyYwIAGLgZD4Fq7Kxb3vFXu3oDqA6MzzICH9h-bpLL8,7517
 cascade/benchmarks/anemoi.py,sha256=qtAI03HdtAmcksCgjIEZyNyUNzMp370KF4lAh5g4cOk,1077
 cascade/benchmarks/generators.py,sha256=NK4fFisWsZdMkA2Auzrn-P7G5D9AKpo2JVnqXE44YT8,2169
 cascade/benchmarks/job1.py,sha256=NY1k9PvkUZODCIDO_zSNwC9sFiMYpBwOaWB7FMSkt2o,4594
@@ -17,20 +17,20 @@ cascade/executor/bridge.py,sha256=vrs-5_Qt2mgkAD7Mzi43Xt_q7tpXX6i1UOPfqZSxHfs,81
 cascade/executor/comms.py,sha256=-9qrKwva6WXkHRQtzSnLFy5gB3bOWuxYJP5fL6Uavw8,8736
 cascade/executor/config.py,sha256=rA4WeCNbdJJ3FdOKJ6WN3_VUorYW3cqdMfKUYPSyj0Y,1471
 cascade/executor/data_server.py,sha256=xLIbLkWn8PnJl4lMP8ADHa2S0EgPwr0-bH7_Sib_Y70,13701
-cascade/executor/executor.py,sha256=KAwCJbv-kO1itPSggEKfVszZYMHEpr7bwgubmW3Qlds,12822
+cascade/executor/executor.py,sha256=SqMVM8BvCNM2r2Zbg9kxSxwFADAaoBU7nCMtfzktsgI,13282
 cascade/executor/msg.py,sha256=QW7Me-8Sin-x-f4M4bzvO7_av8MRkjnabQN6Ch3x22c,4230
 cascade/executor/serde.py,sha256=z6klTOZqW_BVGrbIRNz4FN0_XTfRiKBRQuvgsQIuyAo,2827
 cascade/executor/runner/__init__.py,sha256=30BM80ZyA7w3IrGiKKLSFuhRehbR2Mm99OJ8q5PJ63c,1547
-cascade/executor/runner/entrypoint.py,sha256=kpZ0stgWn5JforZhJVjVCt5RZOgrAL0DBo4A33hauz4,5886
+cascade/executor/runner/entrypoint.py,sha256=paTrkURhI0Vvxb36BUO8QMohIXJdTDVI03o2vNr8VN8,6135
 cascade/executor/runner/memory.py,sha256=EhFhZIFiDo1wDiNuw2gpeUi15yAVDW0hxD7cvX0m0Ho,5299
 cascade/executor/runner/packages.py,sha256=OZjEOvKy8LQ2uguGZU1L7TVYz1415JOUGySRfU_D_sc,2513
 cascade/executor/runner/runner.py,sha256=zqpkvxdWLbwyUFaUbZmSj0KQEBNRpmF8gwVotiaamhc,4870
 cascade/gateway/__init__.py,sha256=1EzMKdLFXEucj0YWOlyVqLx4suOntitwM03T_rRubIk,829
-cascade/gateway/__main__.py,sha256=OcT5Amo0tE1-3StHuTOeQCaABmMBO3XLDK4t4b8oeeQ,647
+cascade/gateway/__main__.py,sha256=x6-DQin6ICvalHT9YcghGyVMoykEATOdN5ON9IeHPYA,862
 cascade/gateway/api.py,sha256=-7HTUhK9idszVCwiVwyHMcNx7n6qRcyPWsLx2e19n3A,2511
 cascade/gateway/client.py,sha256=1p4Tvrf-BH0LQHOES5rY1z3JNIfmXcqWG2kYl4rpcE0,4061
-cascade/gateway/router.py,sha256=p1QT7cMAzLgmfvNYqEtENsnHPPzDbIXcnNuWoI5aejM,7154
-cascade/gateway/server.py,sha256=srqmtOuzzon5GV5jR1wUVh744Ct336BQs5Gzd3mhsPA,3733
+cascade/gateway/router.py,sha256=iN-dc3L46aEy0EV57NNKYwaqIu0Au9kImu1pg-UbxwE,7680
+cascade/gateway/server.py,sha256=tsOyKtVFs5EZmWrjKdi9JwWxK0DG207oSa9OQ-4zN3M,3772
 cascade/low/__init__.py,sha256=5cw2taOGITK_gFbICftzK2YLdEAnLUY5OzblFzdHss4,769
 cascade/low/builders.py,sha256=_u5X8G_EF00hFt8Anv9AXo6yPf1O8MHDmqs2kKmREl0,7073
 cascade/low/core.py,sha256=txya9rgks2b1ze9yLvFvrZCs8sCCtDUlfNwz4sHgybM,5994
@@ -53,9 +53,9 @@ cascade/shm/disk.py,sha256=Fdl_pKOseaXroRp01OwqWVsdI-sSmiFizIFCdxBuMWM,2653
 cascade/shm/func.py,sha256=ZWikgnSLCmbSoW2LDRJwtjxdwTxkR00OUHAsIRQ-ChE,638
 cascade/shm/server.py,sha256=5Ub9bnBmDto9BwfjX3h3sJeiLzZN4lawgtLfvK-vcMU,5036
 earthkit/workflows/__init__.py,sha256=f17AdiV9g4eRN8m4dUnSU58RoLRqk1e6iMRrQiBUSKk,1880
-earthkit/workflows/_version.py,sha256=GJ26tuVcIHne6TgyEAWjeri3FpD3x0dNxpHKKDV8nY0,72
+earthkit/workflows/_version.py,sha256=RCglqs61OUYaUdhn0AKIjoOZXQBCSTS0C7aXjmtZiuA,72
 earthkit/workflows/decorators.py,sha256=DM4QAtQ2glUUcDecwPkXcdlu4dio7MvgpcdmU5LYvD8,937
-earthkit/workflows/fluent.py,sha256=OpN2ncU0UeIyvDcdMe4woH7ZGtUhxFMB3ZQTU0nCshU,30141
+earthkit/workflows/fluent.py,sha256=IN_sqwr7W8wbwP7wTOklgnjVe34IUCmv1ku-DWVTCJc,30179
 earthkit/workflows/mark.py,sha256=PdsXmRfhw1SyyJ74mzFPsLRqMCdlYv556fFX4bqlh9Y,1319
 earthkit/workflows/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 earthkit/workflows/taskgraph.py,sha256=RsT1Qlng1uPZSaSBNqE8vFsoI5J8DDcQl468YPX-kCY,4460
@@ -83,8 +83,8 @@ earthkit/workflows/graph/split.py,sha256=t-Sji5eZb01QO1szqmDNTodDDALqdo-0R0x1ESs
 earthkit/workflows/graph/transform.py,sha256=BZ8n7ePUnuGgoHkMqZC3SLzifu4oq6q6t6vka0khFtg,3842
 earthkit/workflows/graph/visit.py,sha256=MP-aFSqOl7aqJY2i7QTgY4epqb6yM7_lK3ofvOqfahw,1755
 earthkit/workflows/plugins/__init__.py,sha256=WcX4qbEhgTXabIbogydtzNmZ2tB_SuW6NzNkOYQfS-Y,61
-earthkit_workflows-0.3.4.dist-info/licenses/LICENSE,sha256=73MJ7twXMKnWwmzmrMiFwUeY7c6JTvxphVggeUq9Sq4,11381
-earthkit_workflows-0.3.4.dist-info/METADATA,sha256=Okk7IF7eNbGRHPBKk03eMgB1X1IcrsuyPdtfdCL7S6A,1571
-earthkit_workflows-0.3.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-earthkit_workflows-0.3.4.dist-info/top_level.txt,sha256=oNrH3Km3hK5kDkTOiM-8G8OQglvZcy-gUKy7rlooWXs,17
-earthkit_workflows-0.3.4.dist-info/RECORD,,
+earthkit_workflows-0.3.6.dist-info/licenses/LICENSE,sha256=73MJ7twXMKnWwmzmrMiFwUeY7c6JTvxphVggeUq9Sq4,11381
+earthkit_workflows-0.3.6.dist-info/METADATA,sha256=yu-WoCd76So0jElLV9J1woPnwCYcJHoRla8bga7qsgM,1571
+earthkit_workflows-0.3.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+earthkit_workflows-0.3.6.dist-info/top_level.txt,sha256=oNrH3Km3hK5kDkTOiM-8G8OQglvZcy-gUKy7rlooWXs,17
+earthkit_workflows-0.3.6.dist-info/RECORD,,