diff --git a/latch/cli/main.py b/latch/cli/main.py index 6af5d662..82d4ee67 100644 --- a/latch/cli/main.py +++ b/latch/cli/main.py @@ -40,13 +40,23 @@ def main(): default=None, help="The name of your workflow package (the folder with __init__.py). This is a mandatory option if --dockerfile is provided", ) -def register(pkg_root: str, dockerfile: Union[str, None], pkg_name: Union[str, None]): +@click.option( + "--remote", + default=None, + help="The ssh url of a remote git repository where registered workflow code will reside.", +) +def register( + pkg_root: str, + dockerfile: Union[str, None], + pkg_name: Union[str, None], + remote: Union[str, None], +): """Register local workflow code to Latch. Visit docs.latch.bio to learn more. """ try: - _register(pkg_root, dockerfile, pkg_name) + _register(pkg_root, dockerfile, pkg_name, remote) click.secho( "Successfully registered workflow. View @ console.latch.bio.", fg="green" ) diff --git a/latch/config/latch.py b/latch/config/latch.py index b7922ff6..948c2567 100644 --- a/latch/config/latch.py +++ b/latch/config/latch.py @@ -14,6 +14,7 @@ "list-files": "/sdk/list", "initiate-image-upload": "/sdk/initiate-image-upload", "register-workflow": "/sdk/register-workflow", + "commit-workflow": "/sdk/commit-workflow", "get-workflow-interface": "/sdk/wf-interface", "access-jwt": "/sdk/access-jwt", "execute-workflow": "/sdk/wf", diff --git a/latch/services/register/models.py b/latch/services/register/models.py index e30c893c..aefe5099 100644 --- a/latch/services/register/models.py +++ b/latch/services/register/models.py @@ -56,11 +56,13 @@ class RegisterCtx: dkr_repo: Optional[str] = None dkr_client: docker.APIClient = None pkg_root: Path = None # root + remote: Optional[str] = None image_full = None token = None version = None serialize_dir = None latch_register_api_url = endpoints["register-workflow"] + latch_commit_api_url = endpoints["commit-workflow"] latch_image_api_url = endpoints["initiate-image-upload"] def __init__(self, pkg_root: Path, token: Optional[str] = None): diff --git a/latch/services/register/register.py b/latch/services/register/register.py index 5daabffb..d282a117 100644 --- a/latch/services/register/register.py +++ b/latch/services/register/register.py @@ -13,6 +13,7 @@ import requests from latch.services.register import RegisterCtx, RegisterOutput +from latch.utils import retrieve_or_login def _print_build_logs(build_logs, image): @@ -73,6 +74,7 @@ def register( pkg_root: str, dockerfile: Union[str, None] = None, requirements: Union[str, None] = None, + remote: Union[str, None] = None, ) -> RegisterOutput: """Registers a workflow, defined as python code, with Latch. @@ -137,6 +139,7 @@ def register( """ ctx = RegisterCtx(pkg_root) + ctx.remote = remote print(f"Initializing registration for {pkg_root}") if dockerfile is not None: @@ -296,18 +299,18 @@ def _build_file_list(root: str): dockerfile = textwrap.dedent( f""" - FROM {ctx.dkr_repo}/wf-base:fbe8-main + FROM {ctx.dkr_repo}/wf-base:fbe8-main - COPY flytekit.config /root - COPY {ctx.pkg_root.name} /root/{ctx.pkg_root.name} - WORKDIR /root - RUN python3 -m pip install --upgrade latch + COPY flytekit.config /root + COPY {ctx.pkg_root.name} /root/{ctx.pkg_root.name} + WORKDIR /root + RUN python3 -m pip install --upgrade latch - {requirements_cmds} + {requirements_cmds} - ARG tag - ENV FLYTE_INTERNAL_IMAGE $tag - """ + ARG tag + ENV FLYTE_INTERNAL_IMAGE $tag + """ ) dockerfile = BytesIO(dockerfile.encode("utf-8")) dfinfo = tarfile.TarInfo("Dockerfile") @@ -359,13 +362,43 @@ def _upload_pkg_image(ctx: RegisterCtx) -> List[str]: def _register_serialized_pkg(ctx: RegisterCtx, serialize_dir: Path) -> dict: + headers = {"Authorization": f"Bearer {ctx.token}"} - files = {"version": ctx.version.encode("utf-8")} + serialize_files = {"version": ctx.version.encode("utf-8")} for dirname, dirnames, fnames in os.walk(serialize_dir): for filename in fnames + dirnames: file = Path(dirname).resolve().joinpath(filename) - files[file.name] = open(file, "rb") + serialize_files[file.name] = open(file, "rb") + + response = requests.post( + ctx.latch_register_api_url, + headers=headers, + files=serialize_files, + ) + + commit_files = {".workflow_name": ctx.pkg_root.name.encode("utf-8")} + + if not (ctx.remote is None): + commit_files[".remote_name"] = ctx.remote.encode("utf-8") + + for dirname, dirnames, fnames in os.walk(ctx.pkg_root): + for filename in fnames: + file = Path(dirname).resolve().joinpath(filename) + # ignore data folder + if ctx.pkg_root.joinpath("data") in file.parents: + continue + key = str(file.relative_to(ctx.pkg_root)) + commit_files[key] = open(file, "rb") + + commit_response = requests.post( + url=ctx.latch_commit_api_url, + headers=headers, + files=commit_files, + ) + + if not commit_response.json()["success"]: + raise ValueError( + "Issue committing: please make sure the specified remote exists, and that Latch can push to it." + ) - headers = {"Authorization": f"Bearer {ctx.token}"} - response = requests.post(ctx.latch_register_api_url, headers=headers, files=files) return response.json()