inject git commit information into workflow

latch/resources/workflow.py

@@ -1,10 +1,10 @@
 import inspect
+import os
 from dataclasses import is_dataclass
 from textwrap import dedent
 from typing import Callable, Dict, Union, get_args, get_origin
 
 import click
-import os
 from flytekit import workflow as _workflow
 from flytekit.core.workflow import PythonFunctionWorkflow
 
@@ -98,6 +98,15 @@ def decorator(f: Callable):
                 )
                 raise click.exceptions.Exit(1)
 
+        git_hash = os.environ.get("GIT_COMMIT_HASH")
+        is_dirty = os.environ.get("GIT_IS_DIRTY")
+
+        if git_hash is not None:
+            metadata._non_standard["git_commit_hash"] = git_hash
+            metadata._non_standard["git_is_dirty"] = (
+                False if is_dirty is None else is_dirty == "True"
+            )
+
         _inject_metadata(f, metadata)
 
         # note(aidan): used for only serialize_in_container

latch_cli/centromere/ctx.py

@@ -31,6 +31,7 @@
     generate_temporary_ssh_credentials,
     hash_directory,
 )
+from latch_cli.workflow_config import AutoVersionMethod
 
 
 @dataclass
@@ -52,7 +53,7 @@ class _CentromereCtx:
     dkr_client: Optional[docker.APIClient] = None
     ssh_client: Optional[paramiko.SSHClient] = None
     pkg_root: Optional[Path] = None  # root
-    disable_auto_version: bool = False
+    version_method: AutoVersionMethod = False
     image_full = None
     version = None
     serialize_dir = None
@@ -75,11 +76,14 @@ class _CentromereCtx:
     internal_ip: Optional[str] = None
     username: Optional[str] = None
 
+    git_commit_hash: Optional[str] = None
+    git_is_dirty: bool = False
+
     def __init__(
         self,
         pkg_root: Path,
         *,
-        disable_auto_version: bool = False,
+        version_method: AutoVersionMethod = False,
         remote: bool = False,
         metadata_root: Optional[Path] = None,
         snakefile: Optional[Path] = None,
@@ -88,7 +92,7 @@ def __init__(
     ):
         self.use_new_centromere = use_new_centromere
         self.remote = remote
-        self.disable_auto_version = disable_auto_version
+        self.version_method = version_method
 
         try:
             self.token = retrieve_or_login()
@@ -298,10 +302,48 @@ def __init__(
                 )
             self.version = self.version.strip()
 
-            if not self.disable_auto_version:
-                hash = hash_directory(self.pkg_root)
-                self.version = f"{self.version}-{hash[:6]}"
-                click.echo(f"  {self.version}\n")
+            from git import GitError, Repo
+
+            try:
+                repo = Repo(pkg_root)
+                self.git_commit_hash = repo.head.commit.hexsha
+                self.git_is_dirty = repo.is_dirty()
+            except GitError:
+                pass
+            except Exception as e:
+                click.secho(
+                    "WARN: Exception occured while getting git hash from"
+                    f" {self.pkg_root}: {e}",
+                    fg="yellow",
+                )
+
+            if self.version_method != AutoVersionMethod.none:
+                if self.version_method == AutoVersionMethod.directory:
+                    hash = hash_directory(self.pkg_root)[:6]
+                elif self.version_method == AutoVersionMethod.git:
+                    if self.git_commit_hash is None:
+                        click.secho(
+                            dedent(f"""
+                            Failed to extract git commit hash. Please ensure that
+                            the project is a git repository and that the git executable
+                            is available on your PATH.
+                            """),
+                            fg="red",
+                        )
+                        raise click.exceptions.Exit(1)
+
+                    hash = self.git_commit_hash[:6]
+                    if self.git_is_dirty:
+                        click.secho(
+                            dedent("""
+                            The git repository is dirty. The version will be suffixed
+                            with '-wip' until the changes are committed or removed.
+                            """),
+                            fg="yellow",
+                        )
+                        hash += "-wip"
+
+                self.version = f"{self.version}-{hash}"
 
             if self.nucleus_check_version(self.version, self.workflow_name):
                 click.secho(

latch_cli/main.py

@@ -25,7 +25,7 @@
     get_latest_package_version,
     get_local_package_version,
 )
-from latch_cli.workflow_config import BaseImageOptions
+from latch_cli.workflow_config import AutoVersionMethod, BaseImageOptions
 
 latch_cli.click_utils.patch()
 
@@ -455,6 +455,13 @@ def execute(
         " is called."
     ),
 )
+@click.option(
+    "--version-method",
+    help="Method to use when generating the workflow version",
+    type=EnumChoice(AutoVersionMethod, case_sensitive=False),
+    default="directory",
+    show_default=True,
+)
 @click.option(
     "--remote/--no-remote",
     is_flag=True,
@@ -526,6 +533,7 @@ def execute(
 def register(
     pkg_root: str,
     disable_auto_version: bool,
+    version_method: AutoVersionMethod,
     remote: bool,
     docker_progress: str,
     yes: bool,
@@ -554,11 +562,14 @@ def register(
         )
         raise click.exceptions.Exit(1)
 
+    if disable_auto_version:
+        version_method = AutoVersionMethod.none
+
     from latch_cli.services.register import register
 
     register(
         pkg_root,
-        disable_auto_version=disable_auto_version,
+        version_method=version_method,
         remote=remote,
         skip_confirmation=yes,
         open=open,

latch_cli/services/register/register.py

@@ -27,6 +27,7 @@
     upload_image,
 )
 from latch_cli.utils import WorkflowType
+from latch_cli.workflow_config import AutoVersionMethod
 
 
 def _delete_lines(num: int):
@@ -255,7 +256,7 @@ def _recursive_list(directory: Path) -> List[Path]:
 def register(
     pkg_root: str,
     *,
-    disable_auto_version: bool = False,
+    version_method: AutoVersionMethod = AutoVersionMethod.directory,
     remote: bool = False,
     open: bool = False,
     skip_confirmation: bool = False,
@@ -321,7 +322,7 @@ def register(
 
     with _CentromereCtx(
         Path(pkg_root),
-        disable_auto_version=disable_auto_version,
+        version_method=version_method,
         remote=remote,
         metadata_root=metadata_root,
         snakefile=snakefile,
@@ -395,7 +396,9 @@ def register(
             from ...snakemake.serialize import generate_jit_register_code
             from ...snakemake.workflow import build_jit_register_wrapper
 
-            sm_jit_wf = build_jit_register_wrapper(cache_tasks)
+            sm_jit_wf = build_jit_register_wrapper(
+                cache_tasks, ctx.git_commit_hash, ctx.git_is_dirty
+            )
             generate_jit_register_code(
                 sm_jit_wf,
                 ctx.pkg_root,

latch_cli/services/register/utils.py

@@ -19,6 +19,7 @@
 )
 
 import boto3
+import click
 import docker
 import requests
 from latch_sdk_config.latch import config
@@ -116,14 +117,24 @@ def upload_image(ctx: _CentromereCtx, image_name: str) -> List[str]:
 
 
 def serialize_pkg_in_container(
-    ctx: _CentromereCtx, image_name: str, serialize_dir: str, wf_name_override: Optional[str] = None
+    ctx: _CentromereCtx,
+    image_name: str,
+    serialize_dir: str,
+    wf_name_override: Optional[str] = None,
 ) -> Tuple[List[str], str]:
     assert ctx.dkr_client is not None
 
     _env = {"LATCH_DKR_REPO": ctx.dkr_repo, "LATCH_VERSION": ctx.version}
     if wf_name_override is not None:
         _env["LATCH_WF_NAME_OVERRIDE"] = wf_name_override
 
+    if ctx.git_commit_hash is not None:
+        click.secho(
+            f"Tagging workflow version with git commit {ctx.git_commit_hash}", fg="blue"
+        )
+        _env["GIT_COMMIT_HASH"] = ctx.git_commit_hash
+        _env["GIT_IS_DIRTY"] = str(ctx.git_is_dirty)
+
     _serialize_cmd = ["make", "serialize"]
     container = ctx.dkr_client.create_container(
         f"{ctx.dkr_repo}/{image_name}",

latch_cli/snakemake/workflow.py

@@ -311,7 +311,12 @@ def interface_to_parameters(
 class JITRegisterWorkflow(WorkflowBase, ClassStorageTaskResolver):
     out_parameter_name = "o0"  # must be "o0"
 
-    def __init__(self, cache_tasks: bool = False):
+    def __init__(
+        self,
+        cache_tasks: bool = False,
+        git_commit_hash: Optional[str] = None,
+        git_is_dirty: bool = False,
+    ):
         self.cache_tasks = cache_tasks
 
         assert metadata._snakemake_metadata is not None
@@ -336,6 +341,15 @@ def __init__(self, cache_tasks: bool = False):
 
             desc = about_page_path.read_text()
 
+        if git_commit_hash is not None:
+            click.secho(
+                f"Tagging workflow version with git commit {git_commit_hash}", fg="blue"
+            )
+            metadata._snakemake_metadata._non_standard["git_commit_hash"] = (
+                git_commit_hash
+            )
+            metadata._snakemake_metadata._non_standard["git_is_dirty"] = git_is_dirty
+
         docstring = Docstring(
             f"{display_name}\n\n{desc}\n\n" + str(metadata._snakemake_metadata)
         )
@@ -1014,8 +1028,12 @@ def execute(self, **kwargs):
         return exception_scopes.user_entry_point(self._workflow_function)(**kwargs)
 
 
-def build_jit_register_wrapper(cache_tasks: bool = False) -> JITRegisterWorkflow:
-    wrapper_wf = JITRegisterWorkflow(cache_tasks)
+def build_jit_register_wrapper(
+    cache_tasks: bool = False,
+    git_commit_hash: Optional[str] = None,
+    git_is_dirty: bool = False,
+) -> JITRegisterWorkflow:
+    wrapper_wf = JITRegisterWorkflow(cache_tasks, git_commit_hash, git_is_dirty)
     out_parameter_name = wrapper_wf.out_parameter_name
 
     python_interface = wrapper_wf.python_interface

latch_cli/workflow_config.py

@@ -10,6 +10,12 @@
 from latch_cli.constants import latch_constants
 
 
+class AutoVersionMethod(Enum):
+    directory = "directory"
+    git = "git"
+    none = "none"
+
+
 class BaseImageOptions(str, Enum):
     default = "default"
     cuda = "cuda"