earthkit-workflows 0.4.5__py3-none-any.whl → 0.4.7__py3-none-any.whl

cascade/controller/report.py

@@ -27,6 +27,7 @@ JobId = str
 
 @dataclass
 class JobProgress:
+    started: bool
     completed: bool
     pct: (
         str | None
@@ -35,19 +36,20 @@ class JobProgress:
 
     @classmethod
     def failed(cls, failure: str) -> Self:
-        return cls(True, None, failure)
+        return cls(True, True, None, failure)
 
     @classmethod
     def progressed(cls, pct: float) -> Self:
         progress = "{:.2%}".format(pct)[:-1]
-        return cls(False, progress, None)
+        return cls(True, False, progress, None)
 
     @classmethod
     def succeeded(cls) -> Self:
-        return cls(True, None, None)
+        return cls(True, True, None, None)
 
 
-JobProgressStarted = JobProgress(False, "0.00", None)
+JobProgressStarted = JobProgress(True, False, "0.00", None)
+JobProgressEnqueued = JobProgress(False, False, None, None)
 
 
 @dataclass

cascade/executor/runner/packages.py

@@ -42,7 +42,15 @@ class PackagesEnv(AbstractContextManager):
         logger.debug(
             f"installing {len(packages)} packages: {','.join(packages[:3])}{',...' if len(packages) > 3 else ''}"
         )
-        install_command = ["uv", "pip", "install", "--prefix", self.td.name]
+        install_command = [
+            "uv",
+            "pip",
+            "install",
+            "--prefix",
+            self.td.name,
+            "--prerelease",
+            "allow",
+        ]
         if os.environ.get("VENV_OFFLINE", "") == "YES":
             install_command += ["--offline"]
         if cache_dir := os.environ.get("VENV_CACHE", ""):

cascade/gateway/__main__.py

@@ -15,14 +15,17 @@ from cascade.gateway.server import serve
 
 
 def main(
-    url: str, log_base: str | None = None, troika_config: str | None = None
+    url: str,
+    log_base: str | None = None,
+    troika_config: str | None = None,
+    max_jobs: int | 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, troika_config)
+    serve(url, log_base, troika_config, max_jobs)
 
 
 if __name__ == "__main__":

cascade/gateway/api.py

@@ -62,8 +62,9 @@ class JobProgressRequest(CascadeGatewayAPI):
 
 
 class JobProgressResponse(CascadeGatewayAPI):
-    progresses: dict[JobId, JobProgress]
+    progresses: dict[JobId, JobProgress | None]
     datasets: dict[JobId, list[DatasetId]]
+    queue_length: int
     error: str | None  # top level error
 
 

cascade/gateway/router.py

@@ -15,6 +15,7 @@ import os
 import stat
 import subprocess
 import uuid
+from collections import OrderedDict
 from dataclasses import dataclass
 from typing import Iterable
 
@@ -22,7 +23,12 @@ import orjson
 import zmq
 
 import cascade.executor.platform as platform
-from cascade.controller.report import JobId, JobProgress, JobProgressStarted
+from cascade.controller.report import (
+    JobId,
+    JobProgress,
+    JobProgressEnqueued,
+    JobProgressStarted,
+)
 from cascade.executor.comms import get_context
 from cascade.gateway.api import JobSpec, TroikaSpec
 from cascade.low.core import DatasetId
@@ -202,16 +208,29 @@ def _spawn_subprocess(
 
 class JobRouter:
     def __init__(
-        self, poller: zmq.Poller, log_base: str | None, troika_config: str | None
+        self,
+        poller: zmq.Poller,
+        log_base: str | None,
+        troika_config: str | None,
+        max_jobs: int | None,
     ):
         self.poller = poller
         self.jobs: dict[str, Job] = {}
+        self.active_jobs = 0
+        self.max_jobs = max_jobs
+        self.jobs_queue: OrderedDict[JobId, JobSpec] = OrderedDict()
         self.procs: dict[str, subprocess.Popen] = {}
         self.log_base = log_base
         self.troika_config = troika_config
 
-    def spawn_job(self, job_spec: JobSpec) -> JobId:
-        job_id = next_uuid(self.jobs.keys(), lambda: str(uuid.uuid4()))
+    def maybe_spawn(self) -> None:
+        if not self.jobs_queue:
+            return
+        if self.max_jobs is not None and self.active_jobs >= self.max_jobs:
+            logger.debug(f"already running {self.active_jobs}, no spawn")
+            return
+
+        job_id, job_spec = self.jobs_queue.popitem(False)
         base_addr = f"tcp://{platform.get_bindabble_self()}"
         socket = get_context().socket(zmq.PULL)
         port = socket.bind_to_random_port(base_addr)
@@ -222,18 +241,37 @@ class JobRouter:
         self.procs[job_id] = _spawn_subprocess(
             job_spec, full_addr, job_id, self.log_base, self.troika_config
         )
+        self.active_jobs += 1
+        return job_id
+
+    def enqueue_job(self, job_spec: JobSpec) -> JobId:
+        job_id = next_uuid(
+            set(self.jobs.keys()).union(self.jobs_queue.keys()),
+            lambda: str(uuid.uuid4()),
+        )
+        self.jobs_queue[job_id] = job_spec
+        self.maybe_spawn()
         return job_id
 
     def progress_of(
         self, job_ids: Iterable[JobId]
-    ) -> tuple[dict[JobId, JobProgress], dict[JobId, list[DatasetId]]]:
+    ) -> tuple[dict[JobId, JobProgress], dict[JobId, list[DatasetId]], int]:
         if not job_ids:
-            job_ids = self.jobs.keys()
-        progresses = {job_id: self.jobs[job_id].progress for job_id in job_ids}
+            job_ids = set(self.jobs.keys()).union(self.jobs_queue.keys())
+        progresses = {}
+        for job_id in job_ids:
+            if job_id in self.jobs:
+                progresses[job_id] = self.jobs[job_id].progress
+            elif job_id in self.jobs_queue:
+                progresses[job_id] = JobProgressEnqueued
+            else:
+                progresses[job_id] = None
         datasets = {
-            job_id: list(self.jobs[job_id].results.keys()) for job_id in job_ids
+            job_id: list(self.jobs[job_id].results.keys())
+            for job_id in job_ids
+            if job_id in self.jobs
         }
-        return progresses, datasets
+        return progresses, datasets, len(self.jobs_queue)
 
     def get_result(self, job_id: JobId, dataset_id: DatasetId) -> bytes:
         return self.jobs[job_id].results[dataset_id]
@@ -246,6 +284,8 @@ class JobRouter:
         job = self.jobs[job_id]
         if progress.completed:
             self.poller.unregister(job.socket)
+            self.active_jobs -= 1
+            self.maybe_spawn()
         if progress.failure is not None and job.progress.failure is None:
             job.progress = progress
         elif job.last_seen >= timestamp or job.progress.failure is not None:

cascade/gateway/server.py

@@ -31,16 +31,19 @@ def handle_fe(socket: zmq.Socket, jobs: JobRouter) -> bool:
     rv: api.CascadeGatewayAPI
     if isinstance(m, api.SubmitJobRequest):
         try:
-            job_id = jobs.spawn_job(m.job)
+            job_id = jobs.enqueue_job(m.job)
             rv = api.SubmitJobResponse(job_id=job_id, error=None)
         except Exception as e:
             logger.exception(f"failed to spawn a job: {m}")
             rv = api.SubmitJobResponse(job_id=None, error=repr(e))
     elif isinstance(m, api.JobProgressRequest):
         try:
-            progresses, datasets = jobs.progress_of(m.job_ids)
+            progresses, datasets, queue_length = jobs.progress_of(m.job_ids)
             rv = api.JobProgressResponse(
-                progresses=progresses, datasets=datasets, error=None
+                progresses=progresses,
+                datasets=datasets,
+                error=None,
+                queue_length=queue_length,
             )
         except Exception as e:
             logger.exception(f"failed to get progress of: {m}")
@@ -80,7 +83,10 @@ def handle_controller(socket: zmq.Socket, jobs: JobRouter) -> None:
 
 
 def serve(
-    url: str, log_base: str | None = None, troika_config: str | None = None
+    url: str,
+    log_base: str | None = None,
+    troika_config: str | None = None,
+    max_jobs: int | None = None,
 ) -> None:
     ctx = get_context()
     poller = zmq.Poller()
@@ -88,7 +94,7 @@ def serve(
     fe = ctx.socket(zmq.REP)
     fe.bind(url)
     poller.register(fe, flags=zmq.POLLIN)
-    jobs = JobRouter(poller, log_base, troika_config)
+    jobs = JobRouter(poller, log_base, troika_config, max_jobs)
 
     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.4.5"
+__version__ = "0.4.7"

{earthkit_workflows-0.4.5.dist-info → earthkit_workflows-0.4.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: earthkit-workflows
-Version: 0.4.5
+Version: 0.4.7
 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.4.5.dist-info → earthkit_workflows-0.4.7.dist-info}/RECORD

@@ -16,7 +16,7 @@ cascade/controller/act.py,sha256=WHIsk4H-Bbyl_DABX2VWhyKy_cNnp12x1nilatPCL8I,298
 cascade/controller/core.py,sha256=NqvZ5g5GNphwOpzdXbCI0_fxIzzmO97_n2xZKswK72Q,3589
 cascade/controller/impl.py,sha256=9jdTikYO8OkaNIfzatyr3Mhai5EfEhaeii9GaF9cQw4,3526
 cascade/controller/notify.py,sha256=5eSPKcxqrv9kHy7St-iIm1NttsyzcvwLhZI5dvr4cEY,5881
-cascade/controller/report.py,sha256=FD-MAWZq6pwSw2CP2m4OUBw4hzrX46vKE_FZO5NpjDU,3670
+cascade/controller/report.py,sha256=rKGYmq4nIiDqKuP_C7YSwpEAUOPdjILlDcbKkdUt30s,3772
 cascade/executor/bridge.py,sha256=WDE-GM2Bv7nUk1-nV-otMGuaRYw1-Vmd7PWploXBp6Y,8267
 cascade/executor/comms.py,sha256=-9qrKwva6WXkHRQtzSnLFy5gB3bOWuxYJP5fL6Uavw8,8736
 cascade/executor/config.py,sha256=8azy_sXdvDGO0zTNqA0pdtkXsyihM4FQ4U1W_3Dhua0,1571
@@ -28,14 +28,14 @@ cascade/executor/serde.py,sha256=z6klTOZqW_BVGrbIRNz4FN0_XTfRiKBRQuvgsQIuyAo,282
 cascade/executor/runner/__init__.py,sha256=30BM80ZyA7w3IrGiKKLSFuhRehbR2Mm99OJ8q5PJ63c,1547
 cascade/executor/runner/entrypoint.py,sha256=WyxOFGAYDQD_fXsM4H9_6xBrnAmQrCTUnljfcW6-BoM,7918
 cascade/executor/runner/memory.py,sha256=jkAV9T7-imciVcGvkV7OhRfosEpOQJU1OME7z-4ztAs,6371
-cascade/executor/runner/packages.py,sha256=OZjEOvKy8LQ2uguGZU1L7TVYz1415JOUGySRfU_D_sc,2513
+cascade/executor/runner/packages.py,sha256=zOMCKRX34j6vSIK741KV7kTFRgy3bZz_xeLrldslaL4,2633
 cascade/executor/runner/runner.py,sha256=zqpkvxdWLbwyUFaUbZmSj0KQEBNRpmF8gwVotiaamhc,4870
 cascade/gateway/__init__.py,sha256=1EzMKdLFXEucj0YWOlyVqLx4suOntitwM03T_rRubIk,829
-cascade/gateway/__main__.py,sha256=F_wft7ja5ckM0SqeXsy_u2j-Ch6OTlpbTTlYtDkvGMI,917
-cascade/gateway/api.py,sha256=-Vuo9fDqFNFIofcHZ79UB1rTWnQR3D9Pna2CjqdyHaE,3021
+cascade/gateway/__main__.py,sha256=kmfklSeA7v5ie75SBHOql-eHuY6x4eTHlItMYqCQ1Pg,969
+cascade/gateway/api.py,sha256=vPYfiuEjBeddFnCPZpr4_9ovuhGdZ3_migzKTUtvF98,3050
 cascade/gateway/client.py,sha256=1p4Tvrf-BH0LQHOES5rY1z3JNIfmXcqWG2kYl4rpcE0,4061
-cascade/gateway/router.py,sha256=RcDniyPOZnu6_HuMUrQjZ4P-PoUbUVezvYXG_ryBLUg,10399
-cascade/gateway/server.py,sha256=vb3z0TfoMvSHqczhmYgzeXGVcw2M9yGpyW0t6d57Oag,3827
+cascade/gateway/router.py,sha256=9oTkqssb3dHF24TIaAn_7oQoNfm4qkOvriufbOJxnyE,11582
+cascade/gateway/server.py,sha256=BfUKpU2nCEB_zI4BdZU_9zHYHX1WoQaLARCTxMSP0Nk,3971
 cascade/low/__init__.py,sha256=5cw2taOGITK_gFbICftzK2YLdEAnLUY5OzblFzdHss4,769
 cascade/low/builders.py,sha256=_u5X8G_EF00hFt8Anv9AXo6yPf1O8MHDmqs2kKmREl0,7073
 cascade/low/core.py,sha256=_3x4ka_pmCgZbfwFeyhq8S4M6wmh0s24VRCLhk5yQFM,6444
@@ -59,7 +59,7 @@ cascade/shm/disk.py,sha256=Fdl_pKOseaXroRp01OwqWVsdI-sSmiFizIFCdxBuMWM,2653
 cascade/shm/func.py,sha256=ZWikgnSLCmbSoW2LDRJwtjxdwTxkR00OUHAsIRQ-ChE,638
 cascade/shm/server.py,sha256=LnnNX0F6QJt5V_JLfmC3ZMHGNL5WpLY44wpB_pYDr7Y,5042
 earthkit/workflows/__init__.py,sha256=-p4anEn0YQbYWM2tbXb0Vc3wq4-m6kFhcNEgAVu5Jis,1948
-earthkit/workflows/_version.py,sha256=E5wa5mKarJfqqG7Ah0-faFraDaomimplznAnXRNZfLw,72
+earthkit/workflows/_version.py,sha256=Vm-kUKLcx3Zcn5dopRQ_3Wdf6JGosDY_tCioAs7HFG0,72
 earthkit/workflows/decorators.py,sha256=DM4QAtQ2glUUcDecwPkXcdlu4dio7MvgpcdmU5LYvD8,937
 earthkit/workflows/fluent.py,sha256=IN_sqwr7W8wbwP7wTOklgnjVe34IUCmv1ku-DWVTCJc,30179
 earthkit/workflows/mark.py,sha256=PdsXmRfhw1SyyJ74mzFPsLRqMCdlYv556fFX4bqlh9Y,1319
@@ -89,8 +89,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=nhMAC0eMLxoJamjqB5Ns0OWy0OuxEJ_YvaDFGEQITls,129
-earthkit_workflows-0.4.5.dist-info/licenses/LICENSE,sha256=73MJ7twXMKnWwmzmrMiFwUeY7c6JTvxphVggeUq9Sq4,11381
-earthkit_workflows-0.4.5.dist-info/METADATA,sha256=zR79j7-OUPmbz7_K8oXuYCfYSZJlptNepUxPS6j-CZs,1571
-earthkit_workflows-0.4.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-earthkit_workflows-0.4.5.dist-info/top_level.txt,sha256=oNrH3Km3hK5kDkTOiM-8G8OQglvZcy-gUKy7rlooWXs,17
-earthkit_workflows-0.4.5.dist-info/RECORD,,
+earthkit_workflows-0.4.7.dist-info/licenses/LICENSE,sha256=73MJ7twXMKnWwmzmrMiFwUeY7c6JTvxphVggeUq9Sq4,11381
+earthkit_workflows-0.4.7.dist-info/METADATA,sha256=8hfeDm0i94yYA1_CJRFcGrfWE1r7MCzb-HyzvIc5tPs,1571
+earthkit_workflows-0.4.7.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+earthkit_workflows-0.4.7.dist-info/top_level.txt,sha256=oNrH3Km3hK5kDkTOiM-8G8OQglvZcy-gUKy7rlooWXs,17
+earthkit_workflows-0.4.7.dist-info/RECORD,,