From 27eab44b71168f068b7de731a7778065704530ca Mon Sep 17 00:00:00 2001
From: Anirudh Narsipur <anirudh.narsipur@gmail.com>
Date: Thu, 21 May 2026 19:07:09 -0700
Subject: [PATCH] allow specifying workflow name

---
 CHANGELOG.md                                |  6 +++++
 pyproject.toml                              |  2 +-
 src/latch/resources/workflow.py             |  6 ++++-
 src/latch_cli/centromere/ctx.py             | 14 ++++++++---
 src/latch_cli/main.py                       | 27 +++++++++++++++++++++
 src/latch_cli/nextflow/workflow.py          |  5 +++-
 src/latch_cli/services/register/register.py | 20 ++++++++++++---
 src/latch_cli/services/register/staging.py  | 12 ++++++---
 src/latch_cli/snakemake/workflow.py         | 16 +++++++++---
 src/latch_cli/utils/__init__.py             | 13 +++++++++-
 uv.lock                                     |  2 +-
 11 files changed, 105 insertions(+), 18 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4b569ff1..b3b8fcf1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -16,6 +16,12 @@ Types of changes
 
 # Latch SDK Changelog
 
+## 2.75.0 - 2026-05-21
+
+## Added
+
+* `latch register --workflow-name` option to override workflow name for registration.
+
 ## 2.74.0 - 2026-05-18
 
 ## Added
diff --git a/pyproject.toml b/pyproject.toml
index b887cf74..8fbeee54 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -12,7 +12,7 @@ include = ["src/**/*.py", "src/**/*.graphql", "src/**/py.typed", "src/latch_cli/
 
 [project]
 name = "latch"
-version = "2.74.0"
+version = "2.75.0"
 description = "The Latch SDK"
 authors = [{ name = "Kenny Workman", email = "kenny@latch.bio" }]
 maintainers = [{ name = "Kenny Workman", email = "kenny@latch.bio" }]
diff --git a/src/latch/resources/workflow.py b/src/latch/resources/workflow.py
index e0f741f1..f089ecc6 100644
--- a/src/latch/resources/workflow.py
+++ b/src/latch/resources/workflow.py
@@ -52,7 +52,11 @@ def workflow(
         if f.__doc__ is None or "__metadata__:" not in f.__doc__:
             metadata = _generate_metadata(f)
             _inject_metadata(f, metadata)
-        return _workflow(f)
+        wf_name_override = os.environ.get("LATCH_WF_NAME_OVERRIDE")
+        if wf_name_override is not None and wf_name_override.strip() == "":
+            wf_name_override = None
+
+        return _workflow(f, wf_name_override=wf_name_override)
 
     def decorator(f: Callable):
         signature = inspect.signature(f)
diff --git a/src/latch_cli/centromere/ctx.py b/src/latch_cli/centromere/ctx.py
index 2f7be946..091e050e 100644
--- a/src/latch_cli/centromere/ctx.py
+++ b/src/latch_cli/centromere/ctx.py
@@ -1,8 +1,8 @@
-from logging import getLogger
 import re
 import sys
 import traceback
 from dataclasses import dataclass
+from logging import getLogger
 from pathlib import Path
 from textwrap import dedent
 from typing import Dict, Optional, Tuple
@@ -94,6 +94,7 @@ def __init__(
         use_new_centromere: bool = False,
         overwrite: bool = False,
         dockerfile_path: Optional[Path] = None,
+        explicit_workflow_name: Optional[str] = None,
     ):
         self.use_new_centromere = use_new_centromere
         self.remote = remote
@@ -242,7 +243,11 @@ def __init__(
                     )
                     raise click.exceptions.Exit(1)
 
-                self.workflow_name = wf_name
+                self.workflow_name = (
+                    explicit_workflow_name
+                    if explicit_workflow_name is not None
+                    else wf_name
+                )
 
                 for obj in flyte_objects:
                     if obj.type != "task" or obj.dockerfile is None:
@@ -373,8 +378,6 @@ def __init__(
                     )
                     raise click.exceptions.Exit(1)
 
-                # todo(kenny): support per container task and custom workflow
-                # name for snakemake
                 self.workflow_name = f"{metadata._snakemake_metadata.name}_jit_register"
             else:
                 assert self.nf_script is not None
@@ -408,6 +411,9 @@ def __init__(
                 if name_path.exists():
                     self.workflow_name = name_path.read_text().strip()
 
+            if explicit_workflow_name is not None:
+                self.workflow_name = explicit_workflow_name
+
             assert self.workflow_name is not None
 
             if self.nucleus_check_version(self.version, self.workflow_name):
diff --git a/src/latch_cli/main.py b/src/latch_cli/main.py
index 48f188f6..e17bbe24 100644
--- a/src/latch_cli/main.py
+++ b/src/latch_cli/main.py
@@ -33,6 +33,7 @@
     get_latest_package_version,
     get_local_package_version,
     hash_directory,
+    normalize_explicit_workflow_name,
 )
 from latch_cli.workflow_config import BaseImageOptions
 from latch_sdk_gql.execute import execute as gql_execute
@@ -629,6 +630,21 @@ def image_ls():
     ls()
 
 
+def _validate_explicit_workflow_name(
+    ctx: click.Context,
+    _param: click.Parameter,
+    value: Optional[str],
+) -> Optional[str]:
+    if value is None:
+        return None
+
+    value = normalize_explicit_workflow_name(value)
+    if value is None:
+        ctx.fail("--workflow-name must not be empty.")
+
+    return value
+
+
 @latch.command("register")
 @click.argument("pkg_root", type=click.Path(exists=True, file_okay=False))
 @click.option(
@@ -689,6 +705,14 @@ def image_ls():
     type=str,
     help="Module containing Latch workflow to register. Defaults to `wf`",
 )
+@click.option(
+    "--workflow-name",
+    "explicit_workflow_name",
+    type=str,
+    default=None,
+    callback=_validate_explicit_workflow_name,
+    help="Override workflow name to use for this registration.",
+)
 @click.option(
     "--metadata-root",
     type=click.Path(exists=False, path_type=Path, file_okay=False),
@@ -762,6 +786,7 @@ def register(
     yes: bool,
     open: bool,
     workflow_module: Optional[str],
+    explicit_workflow_name: Optional[str],
     metadata_root: Optional[Path],
     snakefile: Optional[Path],
     cache_tasks: bool,
@@ -791,6 +816,7 @@ def register(
             remote=remote,
             skip_confirmation=yes,
             wf_module=workflow_module,
+            explicit_workflow_name=explicit_workflow_name,
             progress_plain=(docker_progress == "auto" and not sys.stdout.isatty())
             or docker_progress == "plain",
             dockerfile_path=dockerfile,
@@ -828,6 +854,7 @@ def register(
         skip_confirmation=yes,
         open=open,
         wf_module=workflow_module,
+        explicit_workflow_name=explicit_workflow_name,
         metadata_root=metadata_root,
         snakefile=snakefile,
         nf_script=nf_script,
diff --git a/src/latch_cli/nextflow/workflow.py b/src/latch_cli/nextflow/workflow.py
index 27add67a..2a5502a7 100644
--- a/src/latch_cli/nextflow/workflow.py
+++ b/src/latch_cli/nextflow/workflow.py
@@ -396,12 +396,15 @@ def generate_nextflow_workflow(
     dest: Path,
     *,
     execution_profile: Optional[str] = None,
+    workflow_name: Optional[str] = None,
 ):
     generate_nextflow_config(pkg_root)
 
     assert metadata._nextflow_metadata is not None
 
-    wf_name = metadata._nextflow_metadata.name
+    wf_name = workflow_name
+    if wf_name is None:
+        wf_name = metadata._nextflow_metadata.name
     assert wf_name is not None
 
     parameters = metadata._nextflow_metadata.parameters
diff --git a/src/latch_cli/services/register/register.py b/src/latch_cli/services/register/register.py
index 643862be..33127795 100644
--- a/src/latch_cli/services/register/register.py
+++ b/src/latch_cli/services/register/register.py
@@ -28,7 +28,7 @@
     serialize_pkg_in_container,
     upload_image,
 )
-from latch_cli.utils import WorkflowType
+from latch_cli.utils import WorkflowType, normalize_explicit_workflow_name
 
 log = getLogger(__name__)
 
@@ -311,6 +311,7 @@ def register(
     open: bool = False,
     skip_confirmation: bool = False,
     wf_module: Optional[str] = None,
+    explicit_workflow_name: Optional[str] = None,
     metadata_root: Optional[Path] = None,
     snakefile: Optional[Path] = None,
     nf_script: Optional[Path] = None,
@@ -374,6 +375,9 @@ def register(
             click.secho("\n`snakemake` package is not installed.", fg="red", bold=True)
             sys.exit(1)
 
+    explicit_workflow_name = normalize_explicit_workflow_name(explicit_workflow_name)
+    use_explicit_workflow_name = explicit_workflow_name is not None
+
     with _CentromereCtx(
         Path(pkg_root),
         disable_auto_version=disable_auto_version,
@@ -385,6 +389,7 @@ def register(
         use_new_centromere=use_new_centromere,
         overwrite=skip_confirmation,
         dockerfile_path=dockerfile_path,
+        explicit_workflow_name=explicit_workflow_name,
     ) as ctx:
         assert ctx.workflow_name is not None, "Unable to determine workflow name"
         assert ctx.version is not None, "Unable to determine workflow version"
@@ -397,6 +402,10 @@ def register(
             ])
         )
         click.echo(" ".join([click.style("Version:", fg="bright_blue"), ctx.version]))
+        if use_explicit_workflow_name:
+            click.echo(
+                "Using workflow name from --workflow-name; .latch/workflow_name will not be updated."
+            )
 
         if workspace_id is None:
             workspace_id = current_workspace()
@@ -467,7 +476,10 @@ def register(
             from ...snakemake.workflow import build_jit_register_wrapper
 
             sm_jit_wf = build_jit_register_wrapper(
-                cache_tasks, ctx.git_commit_hash, ctx.git_is_dirty
+                cache_tasks,
+                ctx.git_commit_hash,
+                ctx.git_is_dirty,
+                workflow_name=ctx.workflow_name,
             )
             generate_jit_register_code(
                 sm_jit_wf,
@@ -492,6 +504,7 @@ def register(
                 ctx.nf_script,
                 dest,
                 execution_profile=nf_execution_profile,
+                workflow_name=ctx.workflow_name,
             )
 
         click.secho("\nInitializing registration", bold=True)
@@ -614,7 +627,8 @@ def register(
             wf_name = ctx.workflow_name
 
             name_path = Path(pkg_root) / latch_constants.pkg_workflow_name
-            if not name_path.exists():
+            if not use_explicit_workflow_name and not name_path.exists():
+                name_path.parent.mkdir(parents=True, exist_ok=True)
                 name_path.write_text(ctx.workflow_name)
 
             while len(wf_infos) == 0:
diff --git a/src/latch_cli/services/register/staging.py b/src/latch_cli/services/register/staging.py
index 91ae54de..fb543309 100644
--- a/src/latch_cli/services/register/staging.py
+++ b/src/latch_cli/services/register/staging.py
@@ -11,7 +11,12 @@
 
 from ...centromere.ast_parsing import get_flyte_objects
 from ...constants import docker_image_name_illegal_pat, latch_constants
-from ...utils import WorkflowType, hash_directory, identifier_suffix_from_str
+from ...utils import (
+    WorkflowType,
+    hash_directory,
+    identifier_suffix_from_str,
+    normalize_explicit_workflow_name,
+)
 from ..docker.utils import dbnp, get_local_docker_client, remote_dbnp
 
 
@@ -23,6 +28,7 @@ def register_staging(
     remote: bool = False,
     skip_confirmation: bool = False,
     wf_module: Optional[str] = None,
+    explicit_workflow_name: Optional[str] = None,
     progress_plain: bool = False,
     dockerfile_path: Optional[Path] = None,
 ):
@@ -50,10 +56,10 @@ def register_staging(
         )
         raise click.exceptions.Exit(1) from e
 
-    wf_name: Optional[str] = None
+    wf_name: Optional[str] = normalize_explicit_workflow_name(explicit_workflow_name)
 
     name_path = pkg_root / latch_constants.pkg_workflow_name
-    if name_path.exists():
+    if wf_name is None and name_path.exists():
         click.echo(f"Parsing workflow name from {name_path}.")
         wf_name = name_path.read_text().strip()
 
diff --git a/src/latch_cli/snakemake/workflow.py b/src/latch_cli/snakemake/workflow.py
index 9f95f880..1cd0a30c 100644
--- a/src/latch_cli/snakemake/workflow.py
+++ b/src/latch_cli/snakemake/workflow.py
@@ -316,6 +316,7 @@ def __init__(
         cache_tasks: bool = False,
         git_commit_hash: Optional[str] = None,
         git_is_dirty: bool = False,
+        workflow_name: Optional[str] = None,
     ):
         self.cache_tasks = cache_tasks
 
@@ -367,7 +368,7 @@ def __init__(
         workflow_metadata = WorkflowMetadata(
             on_failure=WorkflowFailurePolicy.FAIL_IMMEDIATELY
         )
-        name = f"{name}_jit_register"
+        name = workflow_name if workflow_name is not None else f"{name}_jit_register"
         workflow_metadata_defaults = WorkflowMetadataDefaults(False)
         super().__init__(
             name=name,
@@ -577,8 +578,9 @@ def get_fn_code(
         remote_output_url: Optional[str],
     ):
         task_name = f"{self.name}_task"
+        fn_name = identifier_from_str(task_name)
 
-        code_block = self.get_fn_interface(fn_name=task_name)
+        code_block = self.get_fn_interface(fn_name=fn_name)
 
         code_block += reindent(
             rf"""
@@ -763,6 +765,11 @@ class _WorkflowInfoNode(TypedDict):
             1,
         )
         code_block += self.get_fn_return_stmt()
+        code_block += dedent(rf"""
+
+            {fn_name}._name = {task_name!r}
+            setattr(sys.modules[__name__], {task_name!r}, {fn_name})
+            """)
         return code_block
 
 
@@ -1032,8 +1039,11 @@ def build_jit_register_wrapper(
     cache_tasks: bool = False,
     git_commit_hash: Optional[str] = None,
     git_is_dirty: bool = False,
+    workflow_name: Optional[str] = None,
 ) -> JITRegisterWorkflow:
-    wrapper_wf = JITRegisterWorkflow(cache_tasks, git_commit_hash, git_is_dirty)
+    wrapper_wf = JITRegisterWorkflow(
+        cache_tasks, git_commit_hash, git_is_dirty, workflow_name
+    )
     out_parameter_name = wrapper_wf.out_parameter_name
 
     python_interface = wrapper_wf.python_interface
diff --git a/src/latch_cli/utils/__init__.py b/src/latch_cli/utils/__init__.py
index 70e25f41..24eab632 100644
--- a/src/latch_cli/utils/__init__.py
+++ b/src/latch_cli/utils/__init__.py
@@ -13,7 +13,7 @@
 from logging import getLogger
 from pathlib import Path
 from textwrap import dedent
-from typing import List
+from typing import List, Optional
 from urllib.parse import urljoin
 
 import click
@@ -70,6 +70,17 @@ def urljoins(*args: str, dir: bool = False) -> str:
 class AuthenticationError(RuntimeError): ...
 
 
+def normalize_explicit_workflow_name(workflow_name: Optional[str]) -> Optional[str]:
+    if workflow_name is None:
+        return None
+
+    workflow_name = workflow_name.strip()
+    if len(workflow_name) == 0:
+        return None
+
+    return workflow_name
+
+
 def get_auth_header() -> str:
     sdk_token = user_config.token
     execution_token = os.environ.get("FLYTE_INTERNAL_EXECUTION_ID")
diff --git a/uv.lock b/uv.lock
index 9e3c7fa7..205b37bd 100644
--- a/uv.lock
+++ b/uv.lock
@@ -1264,7 +1264,7 @@ wheels = [
 
 [[package]]
 name = "latch"
-version = "2.74.0"
+version = "2.75.0"
 source = { editable = "." }
 dependencies = [
     { name = "aioconsole" },