make caching opt in

Signed-off-by: Ayush Kamat <[email protected]>

latch_cli/main.py

@@ -166,6 +166,17 @@ def dockerfile(pkg_root: str, snakemake: bool = False):
     default=None,
     help="Path to a Snakefile to register.",
 )
+@click.option(
+    "--cache-tasks/--no-cache-tasks",
+    "-c/-C",
+    is_flag=True,
+    default=False,
+    type=bool,
+    help=(
+        "Whether or not to cache snakemake tasks. Ignored if --snakefile is not"
+        " provided."
+    ),
+)
 @requires_login
 def register(
     pkg_root: str,
@@ -174,6 +185,7 @@ def register(
     docker_progress: str,
     yes: bool,
     snakefile: Optional[Path],
+    cache_tasks: bool,
 ):
     """Register local workflow code to Latch.
 
@@ -196,6 +208,7 @@ def register(
         progress_plain=(docker_progress == "auto" and not sys.stdout.isatty())
         or docker_progress == "plain",
         use_new_centromere=use_new_centromere,
+        cache_tasks=cache_tasks,
     )
 
 

latch_cli/services/register/register.py

@@ -196,6 +196,7 @@ def _build_and_serialize(
     dockerfile: Optional[Path] = None,
     *,
     progress_plain: bool = False,
+    cache_tasks: bool = False,
 ):
     assert ctx.pkg_root is not None
 
@@ -207,7 +208,7 @@ def _build_and_serialize(
         from ...snakemake.serialize import generate_jit_register_code
         from ...snakemake.workflow import build_jit_register_wrapper
 
-        jit_wf = build_jit_register_wrapper()
+        jit_wf = build_jit_register_wrapper(cache_tasks)
         generate_jit_register_code(
             jit_wf,
             ctx.pkg_root,
@@ -268,7 +269,8 @@ def register(
     skip_confirmation: bool = False,
     snakefile: Optional[Path] = None,
     *,
-    progress_plain=False,
+    progress_plain: bool = False,
+    cache_tasks: bool = False,
     use_new_centromere: bool = False,
 ):
     """Registers a workflow, defined as python code, with Latch.
@@ -415,6 +417,7 @@ def register(
                 td,
                 dockerfile=ctx.default_container.dockerfile,
                 progress_plain=progress_plain,
+                cache_tasks=cache_tasks,
             )
 
             if remote and snakefile is None:

latch_cli/snakemake/serialize.py

@@ -214,6 +214,7 @@ def extract_snakemake_workflow(
     jit_exec_display_name: str,
     local_to_remote_path_mapping: Optional[Dict[str, str]] = None,
     non_blob_parameters: Optional[Dict[str, Any]] = None,
+    cache_tasks: bool = False,
 ) -> SnakemakeWorkflow:
     extractor = snakemake_workflow_extractor(pkg_root, snakefile, non_blob_parameters)
     with extractor:
@@ -223,6 +224,7 @@ def extract_snakemake_workflow(
             jit_wf_version,
             jit_exec_display_name,
             local_to_remote_path_mapping,
+            cache_tasks,
         )
         wf.compile()
 

latch_cli/snakemake/workflow.py

@@ -289,9 +289,9 @@ def interface_to_parameters(
 class JITRegisterWorkflow(WorkflowBase, ClassStorageTaskResolver):
     out_parameter_name = "o0"  # must be "o0"
 
-    def __init__(
-        self,
-    ):
+    def __init__(self, cache_tasks: bool = False):
+        self.cache_tasks = cache_tasks
+
         assert metadata._snakemake_metadata is not None
 
         parameter_metadata = metadata._snakemake_metadata.parameters
@@ -495,7 +495,7 @@ def get_fn_code(
             print(f"JIT Workflow Version: {{jit_wf_version}}")
             print(f"JIT Execution Display Name: {{jit_exec_display_name}}")
 
-            wf = extract_snakemake_workflow(pkg_root, snakefile, jit_wf_version, jit_exec_display_name, local_to_remote_path_mapping, non_blob_parameters)
+            wf = extract_snakemake_workflow(pkg_root, snakefile, jit_wf_version, jit_exec_display_name, local_to_remote_path_mapping, non_blob_parameters, {self.cache_tasks})
             wf_name = wf.name
             generate_snakemake_entrypoint(wf, pkg_root, snakefile, {repr(remote_output_url)}, non_blob_parameters)
 
@@ -518,7 +518,7 @@ def get_fn_code(
 
                 protos = _recursive_list(td)
                 reg_resp = register_serialized_pkg(protos, None, version, account_id)
-                _print_reg_resp(reg_resp, image_name, silent=True)
+                _print_reg_resp(reg_resp, image_name)
 
             wf_spec_remote = f"latch:///.snakemake_latch/workflows/{wf_name}/{version}/spec"
             spec_dir = Path("spec")
@@ -603,6 +603,7 @@ def __init__(
         jit_wf_version: str,
         jit_exec_display_name: str,
         local_to_remote_path_mapping: Optional[Dict[str, str]] = None,
+        cache_tasks: bool = False,
     ):
         assert metadata._snakemake_metadata is not None
         name = metadata._snakemake_metadata.name
@@ -622,6 +623,7 @@ def __init__(
         self.return_files = return_files
         self._input_parameters = None
         self._dag = dag
+        self._cache_tasks = cache_tasks
         self.snakemake_tasks: List[SnakemakeJobTask] = []
 
         workflow_metadata = WorkflowMetadata(
@@ -742,15 +744,16 @@ def compile(self, **kwargs):
                 if getattr(task, "_metadata") is None:
                     task._metadata = TaskMetadata()
 
-                task._metadata.cache = True
-                task._metadata.cache_serialize = True
+                if self._cache_tasks:
+                    task._metadata.cache = True
+                    task._metadata.cache_serialize = True
 
-                hash = hashlib.new("sha256")
-                hash.update(job.properties().encode())
-                if job.is_script:
-                    hash.update(Path(job.rule.script).read_bytes())
+                    hash = hashlib.new("sha256")
+                    hash.update(job.properties().encode())
+                    if job.is_script:
+                        hash.update(Path(job.rule.script).read_bytes())
 
-                task._metadata.cache_version = hash.hexdigest()
+                    task._metadata.cache_version = hash.hexdigest()
 
                 self.snakemake_tasks.append(task)
 
@@ -831,8 +834,8 @@ def execute(self, **kwargs):
         return exception_scopes.user_entry_point(self._workflow_function)(**kwargs)
 
 
-def build_jit_register_wrapper() -> JITRegisterWorkflow:
-    wrapper_wf = JITRegisterWorkflow()
+def build_jit_register_wrapper(cache_tasks: bool = False) -> JITRegisterWorkflow:
+    wrapper_wf = JITRegisterWorkflow(cache_tasks)
     out_parameter_name = wrapper_wf.out_parameter_name
 
     python_interface = wrapper_wf.python_interface