register fixies

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

latch_cli/centromere/ctx.py

@@ -6,7 +6,7 @@
 from dataclasses import dataclass
 from pathlib import Path
 from textwrap import dedent
-from typing import Dict, Optional, Tuple
+from typing import Dict, Optional, Tuple, cast
 
 import click
 import docker
@@ -55,7 +55,7 @@ class _CentromereCtx:
     dkr_repo: Optional[str] = None
     dkr_client: Optional[docker.APIClient] = None
     ssh_client: Optional[paramiko.SSHClient] = None
-    pkg_root: Optional[Path] = None  # root
+    pkg_root: Path  # root
     disable_auto_version: bool = False
     image_full = None
     version = None
@@ -91,13 +91,13 @@ def __init__(
         self.use_new_centromere = use_new_centromere
         self.remote = remote
         self.disable_auto_version = disable_auto_version
+        self.pkg_root = pkg_root.resolve()
 
         try:
             self.token = retrieve_or_login()
             self.account_id = current_workspace()
 
             self.dkr_repo = config.dkr_repo
-            self.pkg_root = pkg_root.resolve()
 
             if snakefile and nf_script:
                 raise ValueError(
@@ -115,43 +115,48 @@ def __init__(
                 self.workflow_type = WorkflowType.latchbiosdk
 
             self.container_map: Dict[str, _Container] = {}
+
             if self.workflow_type == WorkflowType.latchbiosdk:
                 _import_flyte_objects([self.pkg_root])
                 for entity in FlyteEntities.entities:
                     if isinstance(entity, PythonFunctionWorkflow):
                         self.workflow_name = entity.name
 
                     if isinstance(entity, PythonTask):
-                        if (
-                            hasattr(entity, "dockerfile_path")
-                            and entity.dockerfile_path is not None
-                        ):
+                        dockerfile_path: Optional[Path] = getattr(
+                            entity, "dockerfile_path", None
+                        )
+
+                        if dockerfile_path is not None:
                             self.container_map[entity.name] = _Container(
-                                dockerfile=entity.dockerfile_path,
+                                dockerfile=dockerfile_path,
                                 image_name=self.task_image_name(entity.name),
-                                pkg_dir=entity.dockerfile_path.parent,
+                                pkg_dir=dockerfile_path.parent,
                             )
+
                 if not hasattr(self, "workflow_name"):
                     click.secho(
                         dedent("""
-                                     Unable to locate workflow code. If you
-                                     are a registering a Snakemake project,
-                                     make sure to pass the Snakefile path with
-                                     the --snakefile flag."""),
+                        Unable to locate workflow code.
+
+                        If you are a registering a Snakemake/Nextflow project, make sure to pass the appropriate
+                        script with either the --snakefile or --nf-script flag.
+                        """),
                         fg="red",
                     )
                     raise click.exceptions.Exit(1)
+
             elif self.workflow_type == WorkflowType.snakemake:
                 assert snakefile is not None
 
                 import latch.types.metadata as metadata
 
-                from ..services.register.utils import import_module_by_path
-                from ..snakemake.serialize import (
+                from ..extras.snakemake.serialize import (
                     get_snakemake_metadata_example,
                     snakemake_workflow_extractor,
                 )
-                from ..snakemake.utils import load_snakemake_metadata
+                from ..extras.snakemake.utils import load_snakemake_metadata
+                from ..services.register.utils import import_module_by_path
 
                 meta_file = load_snakemake_metadata(pkg_root)
                 if meta_file is not None:
@@ -166,18 +171,17 @@ def __init__(
                     except (ImportError, FileNotFoundError):
                         traceback.print_exc()
                         click.secho(
-                            "\n\n\n"
-                            + "The above error occured when reading "
-                            + "the Snakefile to extract workflow metadata.",
+                            "\n\n\nThe above error occurred when reading the Snakefile"
+                            " to extract workflow metadata.",
                             bold=True,
                             fg="red",
                         )
                         click.secho(
-                            "\nIt is possible to avoid including the Snakefile"
-                            " prior to registration by providing a"
-                            " `latch_metadata.py` file in the workflow root.\nThis"
-                            " way it is not necessary to install dependencies or"
-                            " ensure that Snakemake inputs locally.",
+                            "\nIt is possible to avoid including the Snakefile prior to"
+                            " registration by providing a `latch_metadata.py` file in"
+                            " the workflow root.\nThis way it is not necessary to"
+                            " install dependencies or ensure that Snakemake inputs"
+                            " locally.",
                             fg="red",
                         )
                         click.secho("\nExample ", fg="red", nl=False)
@@ -228,7 +232,8 @@ def __init__(
 
                                 if res is not None and res != 0:
                                     click.secho("Failed to open file", fg="red")
-                        sys.exit(1)
+
+                        raise click.exceptions.Exit(1)
 
                 if metadata._snakemake_metadata is None:
                     click.secho(
@@ -260,15 +265,17 @@ def __init__(
                 if metadata._nextflow_metadata is None:
                     click.secho(
                         dedent(
-                            """Make sure a `latch_metadata.py` file exists in the
-                        nextflow project root.""",
+                            "Make sure a `latch_metadata` package exists in the"
+                            " nextflow project root.",
                         ),
                         fg="red",
                     )
                     raise click.exceptions.Exit(1)
 
                 self.workflow_name = metadata._nextflow_metadata.name
 
+            assert self.workflow_name is not None
+
             version_file = self.pkg_root / "version"
             try:
                 self.version = version_file.read_text()
@@ -287,7 +294,7 @@ def __init__(
                 if os.environ.get("LATCH_NEW_VERSION_ALWAYS") is not None:
                     self.version = f"{self.version}-{int(time.monotonic())}"
 
-                click.echo(f"  {self.version}\n")
+            click.echo()
 
             if self.nucleus_check_version(self.version, self.workflow_name):
                 click.secho(
@@ -363,6 +370,8 @@ def image(self):
 
         from ..utils import identifier_suffix_from_str
 
+        assert self.workflow_name is not None
+
         wf_name = identifier_suffix_from_str(self.workflow_name).lower()
         wf_name = docker_image_name_illegal_pat.sub("_", wf_name)
 

latch_cli/click_utils.py

@@ -110,6 +110,9 @@ class AnsiCodes:
     bold = "\x1b[1m"
     reset_bold = "\x1b[22m"
 
+    italic = "\x1b[3m"
+    reset_italic = "\x1b[23m"
+
     underline = "\x1b[4m"
     no_underline = "\x1b[24m"
 
@@ -119,13 +122,17 @@ class AnsiCodes:
     url_end = "\x1b]8;;\x1b\\"
 
 
-def bold(s: str) -> str:
+def italic(s: object) -> str:
+    return f"{AnsiCodes.italic}{s}{AnsiCodes.reset_italic}"
+
+
+def bold(s: object) -> str:
     return f"{AnsiCodes.bold}{s}{AnsiCodes.reset_bold}"
 
 
-def underline(s: str) -> str:
+def underline(s: object) -> str:
     return f"{AnsiCodes.underline}{s}{AnsiCodes.no_underline}"
 
 
-def color(s: str, *, color: str = AnsiCodes.color):
+def color(s: object, *, color: str = AnsiCodes.color):
     return f"{color}{s}{AnsiCodes.reset_color}"

latch_cli/extras/nextflow/workflow.py

@@ -64,6 +64,7 @@
 from latch.types import metadata
 from latch.types.metadata import NextflowFileParameter, ParameterType, _IsDataclass
 
+from ...click_utils import italic
 from ...menus import select_tui
 from ...utils import identifier_from_str
 from ..common.serialize import binding_from_python
@@ -1290,16 +1291,33 @@ def build_nf_wf(pkg_root: Path, nf_script: Path) -> NextflowWorkflow:
     try:
         subprocess.run(
             [
-                str(pkg_root / ".latch/bin/nextflow"),
+                str(pkg_root / ".latch" / "bin" / "nextflow"),
+                "-quiet",
                 "run",
                 str(nf_script),
                 "-latchRegister",
             ],
+            env={
+                "NXF_HOME": str(pkg_root / ".latch" / ".nextflow"),
+                "NXF_DISABLE_CHECK_LATEST": "true",
+                "NXF_LOG_FILE": "/dev/null",
+            },
             check=True,
         )
-    except Exception as e:
-        print("\n\n\n[!] Failed\n\n\n")
-        raise e
+    except subprocess.CalledProcessError as e:
+        click.secho(
+            reindent(
+                f"""\
+                An error occurred while parsing your NF script ({italic(nf_script)})
+                Check your script for typos.
+
+                Contact [email protected] for help if the issue persists.
+                """,
+                0,
+            ),
+            fg="red",
+        )
+        raise click.exceptions.Exit(1) from e
 
     dags: Dict[str, DAG] = {}
 
@@ -1350,7 +1368,8 @@ def generate_nf_entrypoint(
     pkg_root: Path,
     nf_path: Path,
 ):
-    entrypoint_code_block = textwrap.dedent(r"""
+    entrypoint_code_block = reindent(
+        r"""
         import os
         from pathlib import Path
         from dataclasses import dataclass, fields
@@ -1383,7 +1402,9 @@ def generate_nf_entrypoint(
         sys.stdout.reconfigure(line_buffering=True)
         sys.stderr.reconfigure(line_buffering=True)
 
-    """).lstrip()
+        """,
+        0,
+    )
 
     # entrypoint_code_block += wf.main_task.get_fn_code(
     #     nf_path_in_container(nf_path, pkg_root)

latch_cli/main.py

@@ -109,9 +109,11 @@ def dockerfile(pkg_root: str, snakemake: bool = False):
         f"Dockerfile already exists at `{dest}`. Overwrite?"
     ):
         return
+
     workflow_type = WorkflowType.latchbiosdk
     if snakemake is True:
         workflow_type = WorkflowType.snakemake
+
     generate_dockerfile(source, dest, wf_type=workflow_type)
 
     click.secho(f"Successfully generated dockerfile `{dest}`", fg="green")