use latch_metadata for docker credentials

latch/types/__init__.py

@@ -2,6 +2,7 @@
 from latch.types.file import LatchFile, LatchOutputFile
 from latch.types.glob import file_glob
 from latch.types.metadata import (
+    DockerMetadata,
     Fork,
     ForkBranch,
     LatchAppearanceType,

latch/types/metadata.py

@@ -521,10 +521,17 @@ def _parameter_str(t: Tuple[str, LatchParameter]):
         ).strip("\n ")
 
 
+class DockerMetadata:
+    def __init__(self, username: str, secret_name: str):
+        self.username = username
+        self.secret_name = secret_name
+
+
 @dataclass
 class SnakemakeMetadata(LatchMetadata):
     output_dir: Optional[LatchDir] = None
     name: Optional[str] = None
+    docker_metadata: Optional[DockerMetadata] = None
     parameters: Dict[str, SnakemakeParameter] = field(default_factory=dict)
 
     def __post_init__(self):

latch_cli/main.py

@@ -185,16 +185,6 @@ def dockerfile(pkg_root: str, snakemake: bool = False):
         " provided."
     ),
 )
-@click.option(
-    "--requires-docker-login",
-    is_flag=True,
-    default=False,
-    type=bool,
-    help=(
-        "Whether or not to use Latch Secrets to login to private docker registry."
-        " Ignored if --snakefile is not provided."
-    ),
-)
 @requires_login
 def register(
     pkg_root: str,
@@ -205,7 +195,6 @@ def register(
     open: bool,
     snakefile: Optional[Path],
     cache_tasks: bool,
-    requires_docker_login: bool,
 ):
     """Register local workflow code to Latch.
 
@@ -230,7 +219,6 @@ def register(
         or docker_progress == "plain",
         use_new_centromere=use_new_centromere,
         cache_tasks=cache_tasks,
-        docker_login=requires_docker_login,
     )
 
 

latch_cli/services/register/register.py

@@ -198,7 +198,6 @@ def _build_and_serialize(
     *,
     progress_plain: bool = False,
     cache_tasks: bool = False,
-    docker_login: bool = False,
 ):
     assert ctx.pkg_root is not None
 
@@ -210,7 +209,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(cache_tasks, docker_login)
+        jit_wf = build_jit_register_wrapper(cache_tasks)
         generate_jit_register_code(
             jit_wf,
             ctx.pkg_root,
@@ -276,7 +275,6 @@ def register(
     snakefile: Optional[Path] = None,
     progress_plain: bool = False,
     cache_tasks: bool = False,
-    docker_login: bool = False,
     use_new_centromere: bool = False,
 ):
     """Registers a workflow, defined as python code, with Latch.
@@ -424,7 +422,6 @@ def register(
                 dockerfile=ctx.default_container.dockerfile,
                 progress_plain=progress_plain,
                 cache_tasks=cache_tasks,
-                docker_login=docker_login,
             )
 
             if remote and snakefile is None:

latch_cli/snakemake/serialize.py

@@ -25,6 +25,8 @@
 from snakemake.workflow import Workflow, WorkflowError
 from typing_extensions import Self
 
+from latch.types.metadata import DockerMetadata
+
 from ..services.register.utils import import_module_by_path
 from .serialize_utils import (
     EntityCache,
@@ -210,7 +212,6 @@ def extract_snakemake_workflow(
     local_to_remote_path_mapping: Optional[Dict[str, str]] = None,
     non_blob_parameters: Optional[Dict[str, Any]] = None,
     cache_tasks: bool = False,
-    docker_login: bool = False,
 ) -> SnakemakeWorkflow:
     extractor = snakemake_workflow_extractor(pkg_root, snakefile, non_blob_parameters)
     with extractor:
@@ -221,7 +222,6 @@ def extract_snakemake_workflow(
             jit_exec_display_name,
             local_to_remote_path_mapping,
             cache_tasks,
-            docker_login,
         )
         wf.compile()
 

latch_cli/snakemake/workflow.py

@@ -74,6 +74,7 @@
 from latch.resources.tasks import custom_task
 from latch.types.directory import LatchDir
 from latch.types.file import LatchFile
+from latch.types.metadata import DockerMetadata
 from latch_cli.snakemake.config.utils import type_repr
 
 from ..utils import identifier_suffix_from_str
@@ -314,9 +315,8 @@ def interface_to_parameters(
 class JITRegisterWorkflow(WorkflowBase, ClassStorageTaskResolver):
     out_parameter_name = "o0"  # must be "o0"
 
-    def __init__(self, cache_tasks: bool = False, docker_login: bool = False):
+    def __init__(self, cache_tasks: bool = False):
         self.cache_tasks = cache_tasks
-        self.docker_login = docker_login
 
         assert metadata._snakemake_metadata is not None
 
@@ -527,7 +527,15 @@ 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, {self.cache_tasks}, {self.docker_login})
+            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)
 
@@ -640,7 +648,6 @@ def __init__(
         jit_exec_display_name: str,
         local_to_remote_path_mapping: Optional[Dict[str, str]] = None,
         cache_tasks: bool = False,
-        docker_login: bool = False,
     ):
         assert metadata._snakemake_metadata is not None
         name = metadata._snakemake_metadata.name
@@ -661,7 +668,7 @@ def __init__(
         self._input_parameters = None
         self._dag = dag
         self._cache_tasks = cache_tasks
-        self._docker_login = docker_login
+        self._docker_metadata = metadata._snakemake_metadata.docker_metadata
         self.snakemake_tasks: List[SnakemakeJobTask] = []
 
         workflow_metadata = WorkflowMetadata(
@@ -875,10 +882,8 @@ def execute(self, **kwargs):
         return exception_scopes.user_entry_point(self._workflow_function)(**kwargs)
 
 
-def build_jit_register_wrapper(
-    cache_tasks: bool = False, docker_login: bool = False
-) -> JITRegisterWorkflow:
-    wrapper_wf = JITRegisterWorkflow(cache_tasks, docker_login)
+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
@@ -1068,9 +1073,6 @@ def __init__(
         self._python_outputs = outputs
         self._target_file_for_input_param = target_file_for_input_param
         self._target_file_for_output_param = target_file_for_output_param
-        self._docker_login = (
-            self.wf._docker_login and self.job.container_img_url is not None
-        )
 
         self._task_function = task_fn_placeholder
 
@@ -1397,15 +1399,18 @@ def get_fn_code(
             1,
         )
 
-        if self._docker_login:
+        if (
+            self.wf._docker_metadata is not None
+            and self.job.container_img_url is not None
+        ):
             code_block += reindent(
                 rf"""
 
                 print("\n\n\nLogging into Docker\n")
                 from latch.functions.secrets import get_secret
+                docker_usr = {self.wf._docker_metadata.username}
                 try:
-                    docker_usr = get_secret("DOCKER_USERNAME")
-                    docker_pwd = get_secret("DOCKER_PASSWORD")
+                    docker_pwd = get_secret({self.wf._docker_metadata.secret_name})
                 except ValueError as e:
                     print("Failed to get Docker credentials:", e)
                     sys.exit(1)