latch 2.41.3.dev0__tar.gz → 2.41.3.dev2__tar.gz

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

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

{latch-2.41.3.dev0 → 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.dev0 → 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.dev0 → 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.dev0 → latch-2.41.3.dev2/latch.egg-info}/PKG-INFO

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

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

@@ -80,6 +80,10 @@ latch_cli/exceptions/cache.py
 latch_cli/exceptions/errors.py
 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
 latch_cli/services/get_executions.py
@@ -107,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
@@ -121,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
@@ -149,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.dev0 → 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/dependencies.py

@@ -0,0 +1,91 @@
+import shutil
+import subprocess
+from concurrent.futures import ProcessPoolExecutor
+from ctypes import c_int
+from multiprocessing.managers import SyncManager
+from pathlib import Path
+from urllib.parse import urljoin
+
+import boto3
+import click
+
+from latch_cli import tinyrequests
+from latch_cli.utils import dedent
+
+
+def _do_download(
+    url: str,
+    output_path: Path,
+    total_count: int,
+    counter,
+    lock,
+):  # todo(ayush): figure out the right type annotation for counter/lock
+    res = tinyrequests.get(url)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    output_path.write_bytes(res.content)
+
+    with lock:
+        counter.value += 1
+        progress_str = f"{counter.value}/{total_count}"
+
+        click.echo("\x1b[0K", nl=False)
+        click.secho(progress_str, dim=True, italic=True, nl=False)
+        click.echo(f"\x1b[{len(progress_str)}D", nl=False)
+
+
+def download_nf_jars(pkg_root: Path):
+    s3_resource = boto3.resource("s3")
+    bucket = s3_resource.Bucket("latch-public")
+
+    objects = list(bucket.objects.filter(Prefix=".nextflow/"))
+
+    click.secho("  Downloading Nextflow binaries: \x1b[?25l", italic=True, nl=False)
+
+    with SyncManager() as man:
+        counter = man.Value(c_int, 0)
+        lock = man.Lock()
+        with ProcessPoolExecutor() as exec:
+            for obj in objects:
+                url = urljoin(
+                    "https://latch-public.s3.us-west-2.amazonaws.com/", obj.key
+                )
+                obj_path = pkg_root / ".latch" / obj.key
+
+                exec.submit(_do_download, url, obj_path, len(objects), counter, lock)
+
+    click.echo("\x1b[0K", nl=False)
+    click.secho("Done. \x1b[?25h", italic=True)
+
+
+def ensure_nf_dependencies(pkg_root: Path, *, force_redownload: bool = False):
+    try:
+        subprocess.run(["java", "--version"], check=True, capture_output=True)
+    except (subprocess.CalledProcessError, FileNotFoundError) as e:
+        click.secho(
+            dedent("""\
+            Java is not installed - this is a requirement to run Nextflow.
+            Please install Java and try again.
+            """),
+            fg="red",
+        )
+        raise click.exceptions.Exit(1) from e
+
+    nf_executable = pkg_root / ".latch" / "bin" / "nextflow"
+    nf_jars = pkg_root / ".latch" / ".nextflow"
+
+    if force_redownload:
+        nf_executable.unlink(missing_ok=True)
+        if nf_jars.exists():
+            shutil.rmtree(nf_jars)
+
+    if not nf_executable.exists():
+        res = tinyrequests.get(
+            "https://latch-public.s3.us-west-2.amazonaws.com/nextflow"
+        )
+        nf_executable.parent.mkdir(parents=True, exist_ok=True)
+
+        nf_executable.write_bytes(res.content)
+        nf_executable.chmod(0o700)
+
+    if not nf_jars.exists():
+        download_nf_jars(pkg_root)

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.dev2/latch_cli/nextflow/workflow.py

@@ -0,0 +1,222 @@
+import re
+from dataclasses import fields, is_dataclass
+from enum import Enum
+from pathlib import Path
+from typing import Any, List, Optional, Tuple
+
+import click
+
+import latch.types.metadata as metadata
+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, 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, 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
+
+
+@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:
+        raise RuntimeError("failed to get execution token")
+
+    print("Provisioning shared storage volume...")
+    headers = {{"Authorization": f"Latch-Execution-Token {{token}}"}}
+    resp = requests.post(
+        "http://nf-dispatcher-service.flyte.svc.cluster.local/provision-storage",
+        headers=headers,
+        json={{
+            "storage_gib": {storage_gib},
+        }}
+    )
+    resp.raise_for_status()
+    return resp.json()["name"]
+
+
+{preambles}
+
+
+@nextflow_runtime_task(cpu={cpu}, memory={memory})
+def nextflow_runtime(pvc_name: str, {param_signature}) -> None:
+    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,
+        )
+    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)
+def {workflow_func_name}({param_signature_with_defaults}) -> None:
+    pvc_name: str = initialize()
+    nextflow_runtime(pvc_name=pvc_name, {param_args})
+
+"""
+
+
+def _get_flags_for_dataclass(name: str, val: Any) -> List[str]:
+    assert is_dataclass(val)
+
+    flags = []
+    for f in fields(val):
+        flags.extend(get_flag(f"{name}.{f.name}", getattr(val, f.name)))
+
+    return flags
+
+
+def get_flag(name: str, val: Any) -> List[str]:
+    flag = f"--{name}"
+
+    if isinstance(val, bool):
+        return [flag] if val else []
+    elif isinstance(val, LatchFile) or isinstance(val, LatchDir):
+        if val.remote_path is not None:
+            return [flag, val.remote_path]
+
+        return [flag, str(val.path)]
+    elif is_dataclass(val):
+        return _get_flags_for_dataclass(name, val)
+    elif isinstance(val, Enum):
+        return [flag, getattr(val, "value")]
+    else:
+        return [flag, str(val)]
+
+
+def generate_nextflow_workflow(
+    pkg_root: Path,
+    workflow_name: str,
+    nf_script: Path,
+    *,
+    execution_profile: Optional[str] = None,
+):
+    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():
+        flags.append(reindent(f"*get_flag({repr(param_name)}, {param_name})", 3))
+
+    defaults: List[Tuple[str, str]] = []
+    no_defaults: List[str] = []
+    preambles: List[str] = []
+    for param_name, param in parameters.items():
+        sig = f"{param_name}: {type_repr(param.type)}"
+        if param.default is not None:
+            if isinstance(param.default, Enum):
+                defaults.append((sig, param.default))
+            else:
+                defaults.append((sig, repr(param.default)))
+        else:
+            no_defaults.append(sig)
+
+        preamble = get_preamble(param.type)
+        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),
+        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]
+        ),
+        param_signature=", ".join(no_defaults + [name for name, _ in defaults]),
+        param_args=", ".join(
+            f"{param_name}={param_name}" for param_name in parameters.keys()
+        ),
+        params_to_flags=",\n".join(flags),
+        execution_profile=(
+            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"
+    entrypoint_path.parent.mkdir(exist_ok=True)
+    entrypoint_path.write_text(entrypoint)
+
+    click.secho(
+        f"Nextflow workflow written to {pkg_root / 'wf' / 'entrypoint.py'}",
+        fg="green",
+    )

{latch-2.41.3.dev0 → latch-2.41.3.dev2}/latch_cli/snakemake/config/utils.py

@@ -192,7 +192,7 @@ def parse_value(t: Type, v: JSONValue):
 def is_primitive_type(
     typ: Type,
 ) -> TypeGuard[Union[Type[None], Type[str], Type[bool], Type[int], Type[float]]]:
-    return typ in {Type[None], str, bool, int, float}
+    return typ in {type(None), str, bool, int, float}
 
 
 def is_primitive_value(val: object) -> TypeGuard[Union[None, str, bool, int, float]]:
@@ -231,11 +231,12 @@ def type_repr(t: Type, *, add_namespace: bool = False) -> str:
     if get_origin(t) is Annotated:
         args = get_args(t)
         assert len(args) > 1
-        assert isinstance(args[1], FlyteAnnotation)
-        return (
-            f"typing_extensions.Annotated[{type_repr(args[0], add_namespace=add_namespace)},"
-            f" FlyteAnnotation({repr(args[1].data)})]"
-        )
+        if isinstance(args[1], FlyteAnnotation):
+            return (
+                f"typing_extensions.Annotated[{type_repr(args[0], add_namespace=add_namespace)},"
+                f" FlyteAnnotation({repr(args[1].data)})]"
+            )
+        return type_repr(args[0], add_namespace=add_namespace)
 
     return t.__name__
 
@@ -250,6 +251,16 @@ def dataclass_repr(typ: Type) -> str:
     return "\n".join(lines) + "\n\n\n"
 
 
+def enum_repr(typ: Type) -> str:
+    assert issubclass(typ, Enum), typ
+
+    lines = [f"class {typ.__name__}(Enum):"]
+    for name, val in typ._member_map_.items():
+        lines.append(f"    {name} = {repr(val.value)}")
+
+    return "\n".join(lines) + "\n\n\n"
+
+
 def get_preamble(typ: Type) -> str:
     if get_origin(typ) is Annotated:
         args = get_args(typ)
@@ -262,6 +273,9 @@ def get_preamble(typ: Type) -> str:
     if get_origin(typ) in {Union, list}:
         return "".join([get_preamble(t) for t in get_args(typ)])
 
+    if issubclass(typ, Enum):
+        return enum_repr(typ)
+
     assert is_dataclass(typ), typ
 
     preamble = "".join([get_preamble(f.type) for f in fields(typ)])

{latch-2.41.3.dev0 → 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.dev0",
+    version="v2.41.3.dev2",
     author_email="kenny@latch.bio",
     description="The Latch SDK",
     packages=find_packages(),