latch 2.40.4.dev6__tar.gz → 2.40.4.dev8__tar.gz

{latch-2.40.4.dev6/latch.egg-info → latch-2.40.4.dev8}/PKG-INFO

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

{latch-2.40.4.dev6 → latch-2.40.4.dev8}/latch/types/metadata.py

@@ -4,7 +4,6 @@ from enum import Enum
 from pathlib import Path
 from textwrap import indent
 from typing import (
-    Annotated,
     Any,
     Callable,
     ClassVar,
@@ -22,7 +21,7 @@ from typing import (
 
 import click
 import yaml
-from typing_extensions import TypeAlias
+from typing_extensions import Annotated, TypeAlias
 
 from latch_cli.utils import identifier_suffix_from_str
 

{latch-2.40.4.dev6 → latch-2.40.4.dev8/latch.egg-info}/PKG-INFO

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

{latch-2.40.4.dev6 → latch-2.40.4.dev8}/latch_cli/extras/common/utils.py

@@ -50,6 +50,21 @@ def is_downloadable_blob_type(typ: Type):
     return True
 
 
+def is_output_dir(typ: Type) -> bool:
+    if typ is LatchOutputDir:
+        return True
+
+    origin = get_origin(typ)
+    if origin is Union:
+        return all([
+            is_output_dir(sub_typ)
+            for sub_typ in get_args(typ)
+            if sub_typ is not type(None)
+        ])
+
+    return False
+
+
 def type_repr(t: Type, *, add_namespace: bool = False) -> str:
     if get_origin(t) == Annotated:
         return type_repr(get_args(t)[0])

{latch-2.40.4.dev6 → latch-2.40.4.dev8}/latch_cli/extras/nextflow/build.py

@@ -183,7 +183,7 @@ def build_from_nextflow_dag(
                 execution_profile=execution_profile,
                 wf=wf,
                 cpu=vertex.cpu,
-                memory=vertex.memoryBytes,
+                memory_bytes=vertex.memoryBytes,
                 storage_gib=ephemeral_storage_gib,
             )
 

{latch-2.40.4.dev6 → latch-2.40.4.dev8}/latch_cli/extras/nextflow/file_persistence.py

@@ -2,7 +2,7 @@ import json
 import os
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Dict, List, Optional, Set, Union, cast
+from typing import Dict, List, Optional, Set, Union
 from urllib.parse import urlparse
 
 import click
@@ -54,7 +54,9 @@ class PathData:
     remote: Optional[str] = None
 
 
-def _extract_paths(parameter: JSONValue, res: List[PathData]):
+def _extract_paths(
+    parameter: JSONValue, res: List[PathData], is_download: bool = False
+):
     if not isinstance(parameter, Dict):
         raise ValueError(f"malformed parameter: {parameter}")
 
@@ -77,7 +79,7 @@ def _extract_paths(parameter: JSONValue, res: List[PathData]):
             return
 
         res.append(PathData(parameter=parameter, local=local, remote=remote))
-    elif "string" in parameter:
+    elif "string" in parameter and is_download:
         v = parameter["string"]
         assert isinstance(v, str)
 
@@ -87,13 +89,13 @@ def _extract_paths(parameter: JSONValue, res: List[PathData]):
         parameter["path"] = v
         del parameter["string"]
 
-        _extract_paths(parameter, res)
+        _extract_paths(parameter, res, is_download)
     elif "list" in parameter:
         v = parameter["list"]
         assert isinstance(v, List)
 
         for x in v:
-            _extract_paths(x, res)
+            _extract_paths(x, res, is_download)
     elif "map" in parameter:
         v = parameter["map"]
         assert isinstance(v, List)
@@ -103,8 +105,8 @@ def _extract_paths(parameter: JSONValue, res: List[PathData]):
             if "key" not in x or "value" not in x:
                 raise ValueError(f"malformed map entry: {x}")
 
-            _extract_paths(x["key"], res)
-            _extract_paths(x["value"], res)
+            _extract_paths(x["key"], res, is_download)
+            _extract_paths(x["value"], res, is_download)
 
 
 def _get_execution_name() -> Optional[str]:
@@ -135,10 +137,10 @@ def download_files(
     path_data: List[PathData] = []
     for channel in channels:
         if type(channel) == dict and "value" in channel:
-            _extract_paths(channel["value"], path_data)
+            _extract_paths(channel["value"], path_data, is_download=True)
         elif type(channel) == list:
             for param in channel:
-                _extract_paths(param, path_data)
+                _extract_paths(param, path_data, is_download=True)
 
     remote = _get_remote(outdir)
 
@@ -219,7 +221,11 @@ def _upload(local: Path, remote: str):
 
 
 # todo(ayush): use crc or something to avoid reuploading unchanged files
-def upload_files(channels: Dict[str, List[JSONValue]], outdir: LatchDir):
+def upload_files(
+    channels: Dict[str, List[JSONValue]],
+    outdir: LatchDir,
+    publish_dir: Optional[LatchDir] = None,
+):
     path_data: List[PathData] = []
     for channel in channels.values():
         if type(channel) == dict and "value" in channel:
@@ -252,7 +258,8 @@ def upload_files(channels: Dict[str, List[JSONValue]], outdir: LatchDir):
 
     for file in published_files:
         relative_path = Path(file).relative_to(Path.home())
-        local_to_remote[file] = urljoins(remote_parent, str(relative_path))
+        dir = publish_dir.remote_path if publish_dir is not None else remote_parent
+        local_to_remote[file] = urljoins(dir, str(relative_path))
 
     for local, remote in local_to_remote.items():
         _upload(local, remote)

{latch-2.40.4.dev6 → latch-2.40.4.dev8}/latch_cli/extras/nextflow/tasks/base.py

@@ -48,7 +48,7 @@ class NextflowBaseTask(PythonAutoContainerTask[Pod]):
         nf_task_type: NFTaskType,
         # ayush: 0.5/1 is too small and stuff starts taking a while to run
         cpu: float = 1,
-        memory: float = 2,
+        memory_gib: float = 2,
         storage_gib: int = 500,
     ):
         self.id = id
@@ -73,13 +73,17 @@ class NextflowBaseTask(PythonAutoContainerTask[Pod]):
 
         self.branches = branches
 
+        self.cpu = cpu
+        self.memory_gib = memory_gib
+        self.storage_gib = storage_gib
+
         super().__init__(
             task_type=SdkTaskType.SIDECAR_TASK,
             task_type_version=2,
             name=f"{name}_{id}",
             interface=interface,
             task_config=custom_task(
-                cpu=cpu, memory=memory, storage_gib=storage_gib
+                cpu=cpu, memory=memory_gib, storage_gib=storage_gib
             ).keywords["task_config"],
             task_resolver=NextflowBaseTaskResolver(),
         )

{latch-2.40.4.dev6 → latch-2.40.4.dev8}/latch_cli/extras/nextflow/tasks/operator.py

@@ -1,6 +1,6 @@
 import json
 from pathlib import Path
-from typing import Annotated, Dict, List, Mapping, Type, get_args, get_origin
+from typing import Dict, List, Mapping, Type
 
 from latch.types.directory import LatchDir
 from latch.types.file import LatchFile
@@ -39,7 +39,7 @@ class NextflowOperatorTask(NextflowBaseTask):
             wf,
             NFTaskType.Operator,
             cpu=4,
-            memory=8,
+            memory_gib=8,
         )
 
     def get_fn_interface(self):

{latch-2.40.4.dev6 → latch-2.40.4.dev8}/latch_cli/extras/nextflow/tasks/process.py

@@ -27,9 +27,17 @@ class NextflowProcessTask(NextflowBaseTask):
         execution_profile: Optional[str],
         wf: NextflowWorkflow,
         cpu: Optional[int] = None,
-        memory: Optional[float] = None,
+        memory_bytes: Optional[int] = None,
         storage_gib: int = 500,
     ):
+        if cpu is None:
+            cpu = 4
+
+        if memory_bytes is None:
+            memory_gib = 16
+        else:
+            memory_gib = memory_bytes // (1024**3)
+
         super().__init__(
             inputs,
             outputs,
@@ -38,6 +46,9 @@ class NextflowProcessTask(NextflowBaseTask):
             {},
             wf,
             NFTaskType.Process,
+            cpu=cpu,
+            memory_gib=memory_gib,
+            storage_gib=storage_gib,
         )
 
         self.wf_inputs = {}
@@ -108,19 +119,6 @@ class NextflowProcessTask(NextflowBaseTask):
                         else:
                             wf_paths["{k}"] = Path("/root") / "{k}"
 
-                    """,
-                    0,
-                )
-            elif is_blob_type(typ):
-                code_block += reindent(
-                    f"""
-                    if {k} is not None:
-                        if not {pre_execute}:
-                            {k}_p = Path("/root/").resolve() # superhack
-                            wf_paths["{k}"] = {k}_p
-                        else:
-                            wf_paths["{k}"] = Path("/root") / "{k}"
-
                     """,
                     0,
                 )
@@ -225,15 +223,15 @@ class NextflowProcessTask(NextflowBaseTask):
 
             def allocate_cpu({input_name}: {type_repr(input_t)}) -> int:
                 res = _read_resources()
-                return max(1, res['cpu_cores']) if res.get('cpu_cores') is not None else 16
+                return max(1, res['cpu_cores']) if res.get('cpu_cores') is not None else {self.cpu}
 
             def allocate_memory({input_name}: {type_repr(input_t)}) -> int:
                 res = _read_resources()
-                return max(1, res['memory_bytes'] // 1024**3) if res.get('memory_bytes') is not None else 48
+                return max(1, res['memory_bytes'] // 1024**3) if res.get('memory_bytes') is not None else {self.memory_gib}
 
             def allocate_disk({input_name}: {type_repr(input_t)}) -> int:
                 res = _read_resources()
-                return max(1, res['disk_bytes'] // 1024**3) if res.get('disk_bytes') is not None else 500
+                return max(1, res['disk_bytes'] // 1024**3) if res.get('disk_bytes') is not None else {self.storage_gib}
 
             def get_resources({input_name}: {type_repr(input_t)}):
         """)
@@ -283,37 +281,6 @@ class NextflowProcessTask(NextflowBaseTask):
             0,
         ).replace("||return|str||", return_str)
 
-    def get_fn_return_stmt(self):
-        results: List[str] = []
-
-        res_type = list(self._python_outputs.values())[0]
-
-        if res_type is None:
-            return "    return None"
-
-        assert is_dataclass(res_type)
-
-        for field in fields(res_type):
-            results.append(
-                reindent(
-                    rf"""
-                    {field.name}=out_channels.get(f"{field.name}")
-                    """,
-                    2,
-                ).rstrip()
-            )
-
-        return_str = ",\n".join(results)
-
-        return reindent(
-            rf"""
-                    return {res_type.__name__}(
-                ||return|str||
-                    )
-            """,
-            0,
-        ).replace("||return|str||", return_str)
-
     def get_fn_code(self, nf_script_path_in_container: Path):
         code_block = dedent(self.get_fn_interface(nf_script_path_in_container))
 
@@ -333,7 +300,8 @@ class NextflowProcessTask(NextflowBaseTask):
 
             print(out_channels)
 
-            upload_files(out_channels, LatchDir({repr(self.wf.output_directory.remote_path)}))
+            publish_dir = {self.wf.publish_dir_param}
+            upload_files(out_channels, LatchDir({repr(self.wf.output_directory.remote_path)}), publish_dir)
 
             out_channels = {{k: json.dumps(v) for k, v in out_channels.items()}}
 

{latch-2.40.4.dev6 → latch-2.40.4.dev8}/latch_cli/extras/nextflow/workflow.py

@@ -15,7 +15,11 @@ from flytekit.core.workflow import (
 from flytekit.exceptions import scopes as exception_scopes
 
 from latch.types import metadata
-from latch_cli.extras.common.utils import is_blob_type, is_downloadable_blob_type
+from latch_cli.extras.common.utils import (
+    is_blob_type,
+    is_downloadable_blob_type,
+    is_output_dir,
+)
 
 from .dag import DAG
 
@@ -33,6 +37,9 @@ from .dag import DAG
 
 
 def _get_flags_to_params(key: str, t: Type, flags: Dict[str, str]):
+    if is_output_dir(t):
+        return
+
     if is_blob_type(t):
         flags[f"--{key}"] = f"wf_paths['wf_{key}']"
     elif is_dataclass(t):
@@ -87,6 +94,16 @@ class NextflowWorkflow(WorkflowBase, ClassStorageTaskResolver):
             if is_downloadable_blob_type(v.type)
         }
 
+        self.publish_dir_param = None
+        for k, v in metadata._nextflow_metadata.parameters.items():
+            if is_output_dir(v.type):
+                if self.publish_dir_param is not None:
+                    click.secho(
+                        "Only one LatchOutputDir parameter is allowed.", fg="red"
+                    )
+                    raise click.exceptions.Exit(1)
+                self.publish_dir_param = f"wf_{k}"
+
         name = metadata._nextflow_metadata.name
         assert name is not None
 

{latch-2.40.4.dev6 → latch-2.40.4.dev8}/setup.py

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