latch 2.41.3.dev1__tar.gz → 2.41.3.dev2__tar.gz

{latch-2.41.3.dev1/latch.egg-info → latch-2.41.3.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: latch
-Version: 2.41.3.dev1
+Version: 2.41.3.dev2
 Summary: The Latch SDK
 Author-email: kenny@latch.bio
 Classifier: Programming Language :: Python :: 3.8

{latch-2.41.3.dev1 → latch-2.41.3.dev2}/latch/resources/tasks.py

@@ -182,40 +182,6 @@ def _get_small_pod() -> Pod:
     )
 
 
-def _get_nextflow_runtime_pod() -> Pod:
-    primary_container = V1Container(name="primary")
-    resources = V1ResourceRequirements(
-        requests={"cpu": "2", "memory": "4Gi", "ephemeral-storage": "100Gi"},
-        limits={"cpu": "2", "memory": "4Gi", "ephemeral-storage": "100Gi"},
-    )
-    primary_container.resources = resources
-    volume_mounts = [V1VolumeMount(mount_path="/nf-workdir", name="nextflow-workdir")]
-    primary_container.volume_mounts = volume_mounts
-
-    return Pod(
-        annotations={
-            "io.kubernetes.cri-o.userns-mode": (
-                "private:uidmapping=0:1048576:65536;gidmapping=0:1048576:65536"
-            )
-        },
-        pod_spec=V1PodSpec(
-            runtime_class_name="sysbox-runc",
-            service_account_name="rahul-test",
-            containers=[primary_container],
-            volumes=[
-                V1Volume(
-                    name="nextflow-workdir",
-                    persistent_volume_claim=V1PersistentVolumeClaimVolumeSource(
-                        # this value will be injected by flytepropeller
-                        claim_name="nextflow-pvc-placeholder"
-                    ),
-                )
-            ],
-        ),
-        primary_container_name="primary",
-    )
-
-
 large_gpu_task = functools.partial(task, task_config=_get_large_gpu_pod())
 """This task will get scheduled on a large GPU-enabled node.
 
@@ -351,8 +317,6 @@ small_task = functools.partial(task, task_config=_get_small_pod())
      - False
 """
 
-nextflow_runtime_task = functools.partial(task, task_config=_get_nextflow_runtime_pod())
-
 
 def custom_memory_optimized_task(cpu: int, memory: int):
     """Returns a custom task configuration requesting
@@ -530,3 +494,43 @@ def custom_task(
     return functools.partial(
         task, task_config=_custom_task_config(cpu, memory, storage_gib), timeout=timeout
     )
+
+
+def nextflow_runtime_task(cpu: int, memory: int):
+    primary_container = V1Container(name="primary")
+    resources = V1ResourceRequirements(
+        requests={
+            "cpu": str(cpu),
+            "memory": f"{memory}Gi",
+            "ephemeral-storage": "20Gi",
+        },
+        limits={"cpu": str(cpu), "memory": f"{memory}Gi", "ephemeral-storage": "20Gi"},
+    )
+    primary_container.resources = resources
+    volume_mounts = [V1VolumeMount(mount_path="/nf-workdir", name="nextflow-workdir")]
+    primary_container.volume_mounts = volume_mounts
+
+    task_config = Pod(
+        annotations={
+            "io.kubernetes.cri-o.userns-mode": (
+                "private:uidmapping=0:1048576:65536;gidmapping=0:1048576:65536"
+            )
+        },
+        pod_spec=V1PodSpec(
+            runtime_class_name="sysbox-runc",
+            service_account_name="rahul-test",
+            containers=[primary_container],
+            volumes=[
+                V1Volume(
+                    name="nextflow-workdir",
+                    persistent_volume_claim=V1PersistentVolumeClaimVolumeSource(
+                        # this value will be injected by flytepropeller
+                        claim_name="nextflow-pvc-placeholder"
+                    ),
+                )
+            ],
+        ),
+        primary_container_name="primary",
+    )
+
+    return functools.partial(task, task_config=task_config)

{latch-2.41.3.dev1 → latch-2.41.3.dev2}/latch/types/directory.py

@@ -114,8 +114,10 @@ class LatchDir(FlyteDirectory):
         self._path_generated = False
 
         if is_valid_url(self.path) and remote_path is None:
+            self._raw_remote_path = path
             self._remote_directory = self.path
         else:
+            self._raw_remote_path = remote_path
             self._remote_directory = None if remote_path is None else str(remote_path)
 
         if kwargs.get("downloader") is not None:

{latch-2.41.3.dev1 → latch-2.41.3.dev2}/latch/types/metadata.py

@@ -483,6 +483,21 @@ class NextflowParameter(Generic[T], LatchParameter):
             self.type = Annotated[self.type, "samplesheet"]
 
 
+@dataclass(frozen=True)
+class NextflowRuntimeResources:
+    """Resources for Nextflow runtime tasks"""
+
+    cpus: Optional[int] = 4
+    """
+    Number of CPUs required for the task
+    """
+    memory: Optional[str] = 8
+    """
+    Memory required for the task (e.g. "1 GB")
+    """
+    storage_gib: Optional[int] = 100
+
+
 @dataclass
 class LatchMetadata:
     """Class for organizing workflow metadata
@@ -683,7 +698,23 @@ _snakemake_metadata: Optional[SnakemakeMetadata] = None
 @dataclass
 class NextflowMetadata(LatchMetadata):
     name: Optional[str] = None
+    """
+    Name of the workflow
+    """
     parameters: Dict[str, NextflowParameter] = field(default_factory=dict)
+    """
+    A dictionary mapping parameter names (strings) to `NextflowParameter` objects
+    """
+    runtime_resources: NextflowRuntimeResources = field(
+        default_factory=NextflowRuntimeResources
+    )
+    """
+    Resources (cpu/memory/storage) for Nextflow runtime task
+    """
+    output_dir: Optional[LatchDir] = None
+    """
+    Directory to dump Nextflow logs
+    """
 
     def __post_init__(self):
         if self.name is None:

{latch-2.41.3.dev1 → latch-2.41.3.dev2/latch.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: latch
-Version: 2.41.3.dev1
+Version: 2.41.3.dev2
 Summary: The Latch SDK
 Author-email: kenny@latch.bio
 Classifier: Programming Language :: Python :: 3.8

{latch-2.41.3.dev1 → latch-2.41.3.dev2}/latch.egg-info/SOURCES.txt

@@ -82,6 +82,7 @@ latch_cli/exceptions/handler.py
 latch_cli/exceptions/traceback.py
 latch_cli/nextflow/__init__.py
 latch_cli/nextflow/dependencies.py
+latch_cli/nextflow/utils.py
 latch_cli/nextflow/workflow.py
 latch_cli/services/__init__.py
 latch_cli/services/get.py
@@ -110,13 +111,7 @@ latch_cli/services/execute/utils.py
 latch_cli/services/init/__init__.py
 latch_cli/services/init/init.py
 latch_cli/services/init/__pycache__/__init__.cpython-310.pyc
-latch_cli/services/init/__pycache__/__init__.cpython-311.pyc
-latch_cli/services/init/__pycache__/__init__.cpython-38.pyc
-latch_cli/services/init/__pycache__/__init__.cpython-39.pyc
 latch_cli/services/init/__pycache__/init.cpython-310.pyc
-latch_cli/services/init/__pycache__/init.cpython-311.pyc
-latch_cli/services/init/__pycache__/init.cpython-38.pyc
-latch_cli/services/init/__pycache__/init.cpython-39.pyc
 latch_cli/services/init/assemble_and_sort/.env
 latch_cli/services/init/assemble_and_sort/LICENSE
 latch_cli/services/init/assemble_and_sort/README.md
@@ -124,22 +119,18 @@ latch_cli/services/init/assemble_and_sort/__init__.py
 latch_cli/services/init/assemble_and_sort/assemble.py
 latch_cli/services/init/assemble_and_sort/sort.py
 latch_cli/services/init/assemble_and_sort/system-requirements.txt
-latch_cli/services/init/assemble_and_sort/__pycache__/__init__.cpython-310.pyc
 latch_cli/services/init/common/.dockerignore
 latch_cli/services/init/example_conda/__init__.py
 latch_cli/services/init/example_conda/conda_task.py
 latch_cli/services/init/example_conda/environment.yaml
-latch_cli/services/init/example_conda/__pycache__/__init__.cpython-310.pyc
 latch_cli/services/init/example_docker/__init__.py
 latch_cli/services/init/example_docker/task.py
-latch_cli/services/init/example_nf_integration/latch_metadata/__pycache__/__init__.cpython-311.pyc
 latch_cli/services/init/example_nfcore/Dockerfile
 latch_cli/services/init/example_nfcore/__init__.py
 latch_cli/services/init/example_nfcore/task.py
 latch_cli/services/init/example_r/__init__.py
 latch_cli/services/init/example_r/environment.R
 latch_cli/services/init/example_r/r_task.py
-latch_cli/services/init/example_r/__pycache__/__init__.cpython-310.pyc
 latch_cli/services/init/example_snakemake/Dockerfile
 latch_cli/services/init/example_snakemake/Snakefile
 latch_cli/services/init/example_snakemake/config.yaml
@@ -152,7 +143,6 @@ latch_cli/services/init/template/LICENSE
 latch_cli/services/init/template/README.md
 latch_cli/services/init/template/__init__.py
 latch_cli/services/init/template/task.py
-latch_cli/services/init/template/__pycache__/__init__.cpython-310.pyc
 latch_cli/services/register/__init__.py
 latch_cli/services/register/constants.py
 latch_cli/services/register/register.py

{latch-2.41.3.dev1 → latch-2.41.3.dev2}/latch_cli/centromere/ctx.py

@@ -258,7 +258,7 @@ class _CentromereCtx:
                     click.secho(
                         dedent("""
                         Failed to register Nextflow workflow.
-                        Make sure the project root contains a `wf/__init__.py`
+                        Make sure the project root contains a `latch_metadata/__init__.py`
                         with a `NextflowMetadata` object defined.
                         """),
                         fg="red",

latch-2.41.3.dev2/latch_cli/nextflow/utils.py

@@ -0,0 +1,27 @@
+import os
+from typing import Optional
+
+import gql
+from latch_sdk_gql.execute import execute
+
+
+def _get_execution_name() -> Optional[str]:
+    token = os.environ.get("FLYTE_INTERNAL_EXECUTION_ID", None)
+    if token is None:
+        return None
+
+    res = execute(
+        gql.gql("""
+        query executionCreatorsByToken($token: String!) {
+            executionCreatorByToken(token: $token) {
+                flytedbId
+                info {
+                    displayName
+                }
+            }
+        }
+        """),
+        {"token": token},
+    )["executionCreatorByToken"]
+
+    return res["info"]["displayName"]

{latch-2.41.3.dev1 → latch-2.41.3.dev2}/latch_cli/nextflow/workflow.py

@@ -2,7 +2,7 @@ import re
 from dataclasses import fields, is_dataclass
 from enum import Enum
 from pathlib import Path
-from typing import Any, List, Optional, Tuple, Type, TypeVar
+from typing import Any, List, Optional, Tuple
 
 import click
 
@@ -11,27 +11,34 @@ from latch.types.directory import LatchDir
 from latch.types.file import LatchFile
 from latch_cli.snakemake.config.utils import get_preamble, type_repr
 from latch_cli.snakemake.utils import reindent
-from latch_cli.utils import identifier_from_str
+from latch_cli.utils import identifier_from_str, urljoins
 
 template = """\
+from dataclasses import dataclass
 from enum import Enum
 import os
 import subprocess
 import requests
+import shutil
 from pathlib import Path
 import typing
 import typing_extensions
 
 from latch.resources.workflow import workflow
-from latch.resources.tasks import nextflow_runtime_task, small_task
+from latch.resources.tasks import nextflow_runtime_task, custom_task
 from latch.types.file import LatchFile
 from latch.types.directory import LatchDir, LatchOutputDir
+from latch.ldata.path import LPath
 from latch_cli.nextflow.workflow import get_flag
+from latch_cli.nextflow.utils import _get_execution_name
+from latch_cli.utils import urljoins
 from latch.types import metadata
 from flytekit.core.annotation import FlyteAnnotation
 
+import latch_metadata
 
-@small_task
+
+@custom_task(cpu=0.25, memory=0.5, storage_gib=1)
 def initialize() -> str:
     token = os.environ.get("FLYTE_INTERNAL_EXECUTION_ID")
     if token is None:
@@ -42,36 +49,70 @@ def initialize() -> str:
     resp = requests.post(
         "http://nf-dispatcher-service.flyte.svc.cluster.local/provision-storage",
         headers=headers,
+        json={{
+            "storage_gib": {storage_gib},
+        }}
     )
     resp.raise_for_status()
-    print("Done.")
     return resp.json()["name"]
 
 
 {preambles}
 
 
-@nextflow_runtime_task
+@nextflow_runtime_task(cpu={cpu}, memory={memory})
 def nextflow_runtime(pvc_name: str, {param_signature}) -> None:
-    env = {{
-        **os.environ,
-        "NXF_HOME": "/root/.nextflow",
-        "K8_STORAGE_CLAIM_NAME": pvc_name,
-    }}
-    subprocess.run(
-        [
-            "/root/.latch/bin/nextflow",
-            "run",
-            "{script_dir}",
-            "-work-dir",
-            "/nf-workdir",
-            "-profile",
-            "{execution_profile}",
+    try:
+        shared_dir = Path("/nf-workdir")
+
+        shutil.copytree(
+            Path("/root"),
+            shared_dir,
+            ignore=lambda src, names: ["latch", ".latch"],
+            ignore_dangling_symlinks=True,
+            dirs_exist_ok=True,
+        )
+
+        env = {{
+            **os.environ,
+            "NXF_HOME": "/root/.nextflow",
+            "K8_STORAGE_CLAIM_NAME": pvc_name,
+        }}
+        subprocess.run(
+            [
+                "/root/.latch/bin/nextflow",
+                "run",
+                str(shared_dir / "{nf_script}"),
+                "-work-dir",
+                str(shared_dir),
+                "-profile",
+                "{execution_profile}",
 {params_to_flags}
-        ],
-        env=env,
-        check=True,
-    )
+            ],
+            env=env,
+            check=True,
+        )
+    except subprocess.CalledProcessError:
+        remote = LPath(urljoins("{remote_output_dir}", _get_execution_name(), "nextflow.log"))
+        print(f"Uploading .nextflow.log to {{remote.path}}")
+        remote.upload_from(Path("/root/.nextflow.log"))
+        raise
+    finally:
+        token = os.environ.get("FLYTE_INTERNAL_EXECUTION_ID")
+        if token is None:
+            raise RuntimeError("failed to get execution token")
+
+        headers = {{"Authorization": f"Latch-Execution-Token {{token}}"}}
+        resp = requests.post(
+            "http://nf-dispatcher-service.flyte.svc.cluster.local/finalize",
+            headers=headers,
+            json={{
+                "pvc_name": pvc_name,
+            }}
+        )
+        if resp.status_code != 200:
+            print("Failed to finalize workflow:", resp.status_code)
+
 
 
 @workflow(metadata._nextflow_metadata)
@@ -119,7 +160,9 @@ def generate_nextflow_workflow(
 ):
     assert metadata._nextflow_metadata is not None
 
+    wf_name = metadata._nextflow_metadata.name
     parameters = metadata._nextflow_metadata.parameters
+    resources = metadata._nextflow_metadata.runtime_resources
 
     flags = []
     for param_name, param in parameters.items():
@@ -142,9 +185,15 @@ def generate_nextflow_workflow(
         if len(preamble) > 0:
             preambles.append(preamble)
 
+    if metadata._nextflow_metadata.output_dir is None:
+        output_dir = "latch:///nextflow_outputs"
+    else:
+        output_dir = metadata._nextflow_metadata.output_dir._raw_remote_path
+    output_dir = urljoins(output_dir, wf_name)
+
     entrypoint = template.format(
         workflow_func_name=identifier_from_str(workflow_name),
-        script_dir=nf_script.resolve().relative_to(pkg_root.resolve()),
+        nf_script=nf_script.resolve().relative_to(pkg_root.resolve()),
         param_signature_with_defaults=", ".join(
             no_defaults + [f"{name}={val}" for name, val in defaults]
         ),
@@ -157,6 +206,10 @@ def generate_nextflow_workflow(
             execution_profile if execution_profile is not None else "standard"
         ),
         preambles="\n\n".join(preambles),
+        cpu=resources.cpus,
+        memory=resources.memory,
+        storage_gib=resources.storage_gib,
+        remote_output_dir=output_dir,
     )
 
     entrypoint_path = pkg_root / "wf" / "entrypoint.py"

{latch-2.41.3.dev1 → latch-2.41.3.dev2}/setup.py

@@ -13,7 +13,7 @@ if cur_ver < (3, 8) or cur_ver > (3, 11):
 
 setup(
     name="latch",
-    version="v2.41.3.dev1",
+    version="v2.41.3.dev2",
     author_email="kenny@latch.bio",
     description="The Latch SDK",
     packages=find_packages(),