snakemake redo

Justfile

@@ -1,7 +1,7 @@
 # Setup
 
 install:
-  uv sync
+  uv sync --group dev --group docs
 
 # Packaging
 
@@ -16,8 +16,8 @@ publish:
 # Testing
 
 test:
-  export TEST_TOKEN=$(cat ~/.latch/token) &&\
-    pytest -s tests
+  export TEST_TOKEN=$(cat ~/.latch/token)
+  pytest -s
 
 # Docs
 

pyproject.toml

@@ -12,7 +12,7 @@ include = ["src/**/*.py", "src/latch_cli/services/init/*"]
 
 [project]
 name = "latch"
-version = "2.55.0"
+version = "2.55.0.a6"
 description = "The Latch SDK"
 authors = [{ name = "Kenny Workman", email = "[email protected]" }]
 maintainers = [
@@ -34,13 +34,13 @@ dependencies = [
   "scp>=0.14.0",
   "boto3>=1.26.0",
   "tqdm>=4.63.0",
-  "lytekit==0.15.13",
+  "lytekit==0.15.17",
   "lytekitplugins-pods==0.6.2",
   "typing-extensions>=4.12.0",
   "apscheduler>=3.10.0",
-  "gql==3.4.0",
+  "gql>=3.5.0,<4.0.0",
   "graphql-core==3.2.3",
-  "requests-toolbelt==0.10.1",
+  "requests-toolbelt>=1.0.0,<2.0.0",
   "python-dateutil>=2.8",
   "GitPython==3.1.40",
   # for old latch develop, to be removed
@@ -73,7 +73,11 @@ classifiers = [
 
 [project.optional-dependencies]
 pandas = ["pandas>=2.0.0"]
-snakemake = ["snakemake>=7.18.0,<7.30.2", "pulp>=2.0,<2.8"]
+snakemake = [
+  "snakemake",
+  "snakemake-storage-plugin-latch",
+  "snakemake-executor-plugin-latch",
+]
 
 [project.scripts]
 latch = "latch_cli.main:main"
@@ -97,11 +101,10 @@ docs = [
 ]
 
 [tool.ruff]
+line-length = 100
 target-version = "py39"
 
 [tool.ruff.lint]
-preview = true
-
 pydocstyle = { convention = "google" }
 extend-select = [
   "F",

src/latch/resources/tasks.py

@@ -296,8 +296,7 @@ def _get_small_pod() -> Pod:
             )
         },
         pod_spec=V1PodSpec(
-            runtime_class_name="sysbox-runc",
-            containers=[primary_container],
+            runtime_class_name="sysbox-runc", containers=[primary_container]
         ),
         primary_container_name="primary",
     )
@@ -517,11 +516,7 @@ class _NGConfig:
 max_storage_gb_ish = int(max_storage_gib * Units.GiB / Units.GB)
 
 
-def _custom_task_config(
-    cpu: int,
-    memory: int,
-    storage_gib: int,
-) -> Pod:
+def _custom_task_config(cpu: int, memory: int, storage_gib: int) -> Pod:
     target_ng = None
     for ng in taint_data:
         if (
@@ -591,10 +586,7 @@ def custom_task(
     """
     if callable(cpu) or callable(memory) or callable(storage_gib):
         task_config = DynamicTaskConfig(
-            cpu=cpu,
-            memory=memory,
-            storage=storage_gib,
-            pod_config=_get_small_pod(),
+            cpu=cpu, memory=memory, storage=storage_gib, pod_config=_get_small_pod()
         )
         return functools.partial(task, task_config=task_config, timeout=timeout)
 
@@ -627,3 +619,27 @@ def nextflow_runtime_task(cpu: int, memory: int, storage_gib: int = 50):
     ]
 
     return functools.partial(task, task_config=task_config)
+
+
+def snakemake_runtime_task(*, cpu: int, memory: int, storage_gib: int = 50):
+    task_config = _custom_task_config(cpu, memory, storage_gib)
+
+    task_config.pod_spec.automount_service_account_token = True
+
+    assert len(task_config.pod_spec.containers) == 1
+    task_config.pod_spec.containers[0].volume_mounts = [
+        V1VolumeMount(mount_path="/snakemake-workdir", name="snakemake-workdir")
+    ]
+
+    task_config.pod_spec.volumes = [
+        V1Volume(
+            name="snakemake-workdir",
+            persistent_volume_claim=V1PersistentVolumeClaimVolumeSource(
+                # this value will be injected by flytepropeller
+                # ayush: this is also used by snakemake bc why not
+                claim_name="nextflow-pvc-placeholder"
+            ),
+        )
+    ]
+
+    return functools.partial(task, task_config=task_config)

src/latch/types/directory.py

@@ -9,18 +9,15 @@
 from flytekit.core.type_engine import TypeEngine, TypeTransformer
 from flytekit.exceptions.user import FlyteUserException
 from flytekit.models.literals import Literal
-from flytekit.types.directory.types import (
-    FlyteDirectory,
-    FlyteDirToMultipartBlobTransformer,
-)
-from latch_sdk_gql.execute import execute
+from flytekit.types.directory.types import FlyteDirectory, FlyteDirToMultipartBlobTransformer
 from typing_extensions import Annotated
 
 from latch.ldata.path import LPath
 from latch.types.file import LatchFile
 from latch.types.utils import format_path, is_valid_url
 from latch_cli.utils import urljoins
 from latch_cli.utils.path import normalize_path
+from latch_sdk_gql.execute import execute
 
 
 class IterdirChild(TypedDict):
@@ -93,10 +90,7 @@ def task(dir: LatchFile):
     """
 
     def __init__(
-        self,
-        path: Union[str, PathLike],
-        remote_path: Optional[PathLike] = None,
-        **kwargs,
+        self, path: Union[str, PathLike], remote_path: Optional[PathLike] = None, **kwargs
     ):
         if path is None:
             raise ValueError("Unable to instantiate LatchDir with None")
@@ -136,9 +130,7 @@ def downloader():
                     self._idempotent_set_path()
 
                     return ctx.file_access.get_data(
-                        self._remote_directory,
-                        self.path,
-                        is_multipart=True,
+                        self._remote_directory, self.path, is_multipart=True
                     )
 
             super().__init__(self.path, downloader, self._remote_directory)
@@ -252,10 +244,7 @@ def __repr__(self):
         if self.remote_path is None:
             return f"LatchDir({repr(format_path(self.local_path))})"
 
-        return (
-            f"LatchDir({repr(self.path)},"
-            f" remote_path={repr( format_path(self.remote_path))})"
-        )
+        return f"LatchDir({repr(self.path)}, remote_path={repr(format_path(self.remote_path))})"
 
     def __str__(self):
         if self.remote_path is None:
@@ -264,12 +253,7 @@ def __str__(self):
         return f"LatchDir({format_path(self.remote_path)})"
 
 
-LatchOutputDir = Annotated[
-    LatchDir,
-    FlyteAnnotation(
-        {"output": True},
-    ),
-]
+LatchOutputDir = Annotated[LatchDir, FlyteAnnotation({"output": True})]
 """A LatchDir tagged as the output of some workflow.
 
 The Latch Console uses this metadata to avoid checking for existence of the
@@ -284,24 +268,17 @@ def __init__(self):
         TypeTransformer.__init__(self, name="LatchDirPath", t=LatchDir)
 
     def to_python_value(
-        self,
-        ctx: FlyteContext,
-        lv: Literal,
-        expected_python_type: Union[Type[LatchDir], PathLike],
+        self, ctx: FlyteContext, lv: Literal, expected_python_type: Union[Type[LatchDir], PathLike]
     ) -> FlyteDirectory:
         uri = lv.scalar.blob.uri
         if expected_python_type is PathLike:
-            raise TypeError(
-                "Casting from Pathlike to LatchDir is currently not supported."
-            )
+            raise TypeError("Casting from Pathlike to LatchDir is currently not supported.")
 
         while get_origin(expected_python_type) == Annotated:
             expected_python_type = get_args(expected_python_type)[0]
 
         if not issubclass(expected_python_type, LatchDir):
-            raise TypeError(
-                f"Neither os.PathLike nor LatchDir specified {expected_python_type}"
-            )
+            raise TypeError(f"Neither os.PathLike nor LatchDir specified {expected_python_type}")
 
         # This is a local file path, like /usr/local/my_file, don't mess with it. Certainly, downloading it doesn't
         # make any sense.

src/latch/types/file.py

@@ -11,12 +11,12 @@
 from flytekit.core.type_engine import TypeEngine, TypeTransformer
 from flytekit.models.literals import Literal
 from flytekit.types.file.file import FlyteFile, FlyteFilePathTransformer
-from latch_sdk_gql.execute import execute
 from typing_extensions import Annotated
 
 from latch.ldata.path import LPath
 from latch.types.utils import format_path, is_absolute_node_path, is_valid_url
 from latch_cli.utils.path import normalize_path
+from latch_sdk_gql.execute import execute
 
 
 class LatchFile(FlyteFile):
@@ -115,9 +115,7 @@ def downloader():
                     self._idempotent_set_path(local_path_hint)
 
                     return ctx.file_access.get_data(
-                        self._remote_path,
-                        self.path,
-                        is_multipart=False,
+                        self._remote_path, self.path, is_multipart=False
                     )
 
             super().__init__(self.path, downloader, self._remote_path)
@@ -162,23 +160,15 @@ def __repr__(self):
         if self.remote_path is None:
             return f"LatchFile({repr(format_path(self.local_path))})"
 
-        return (
-            f"LatchFile({repr(self.path)},"
-            f" remote_path={repr(format_path(self.remote_path))})"
-        )
+        return f"LatchFile({repr(self.path)}, remote_path={repr(format_path(self.remote_path))})"
 
     def __str__(self):
         if self.remote_path is None:
             return "LatchFile()"
         return f"LatchFile({format_path(self.remote_path)})"
 
 
-LatchOutputFile = Annotated[
-    LatchFile,
-    FlyteAnnotation(
-        {"output": True},
-    ),
-]
+LatchOutputFile = Annotated[LatchFile, FlyteAnnotation({"output": True})]
 """A LatchFile tagged as the output of some workflow.
 
 The Latch Console uses this metadata to avoid checking for existence of the
@@ -193,21 +183,14 @@ def __init__(self):
         TypeTransformer.__init__(self, name="LatchFilePath", t=LatchFile)
 
     def to_python_value(
-        self,
-        ctx: FlyteContext,
-        lv: Literal,
-        expected_python_type: Union[Type[LatchFile], PathLike],
+        self, ctx: FlyteContext, lv: Literal, expected_python_type: Union[Type[LatchFile], PathLike]
     ) -> LatchFile:
         uri = lv.scalar.blob.uri
         if expected_python_type is PathLike:
-            raise TypeError(
-                "Casting from Pathlike to LatchFile is currently not supported."
-            )
+            raise TypeError("Casting from Pathlike to LatchFile is currently not supported.")
 
         if not issubclass(expected_python_type, LatchFile):
-            raise TypeError(
-                f"Neither os.PathLike nor LatchFile specified {expected_python_type}"
-            )
+            raise TypeError(f"Neither os.PathLike nor LatchFile specified {expected_python_type}")
 
         # This is a local file path, like /usr/local/my_file, don't mess with it. Certainly, downloading it doesn't
         # make any sense.

src/latch/types/glob.py

@@ -20,7 +20,7 @@ def file_glob(
     represented by the `remote_directory`.
 
     Args:
-        pattern: A glob pattern to match a set of files, eg. '\*.py'. Will
+        pattern: A glob pattern to match a set of files, eg. '*.py'. Will
             resolve paths with respect to the working directory of the caller.
         remote_directory: A valid latch URL pointing to a directory, eg.
             latch:///foo. This _must_ be a directory and not a file.