truefoundry 0.2.10__py3-none-any.whl → 0.3.0rc2__py3-none-any.whl

truefoundry/deploy/lib/win32.py

@@ -0,0 +1,129 @@
+"""
+Win32 compatibility utilities.
+https://github.com/zeromq/pyzmq/blob/3e7e83a841829fcb46b7e1a881d968fb5471ef62/zmq/utils/win32.py
+"""
+
+# -----------------------------------------------------------------------------
+# Copyright (C) PyZMQ Developers
+# Distributed under the terms of the Modified BSD License.
+# -----------------------------------------------------------------------------
+
+import os
+from typing import Any, Callable, Optional
+
+
+class allow_interrupt:
+    """Utility for fixing CTRL-C events on Windows.
+
+    On Windows, the Python interpreter intercepts CTRL-C events in order to
+    translate them into ``KeyboardInterrupt`` exceptions.  It (presumably)
+    does this by setting a flag in its "console control handler" and
+    checking it later at a convenient location in the interpreter.
+
+    However, when the Python interpreter is blocked waiting for the ZMQ
+    poll operation to complete, it must wait for ZMQ's ``select()``
+    operation to complete before translating the CTRL-C event into the
+    ``KeyboardInterrupt`` exception.
+
+    The only way to fix this seems to be to add our own "console control
+    handler" and perform some application-defined operation that will
+    unblock the ZMQ polling operation in order to force ZMQ to pass control
+    back to the Python interpreter.
+
+    This context manager performs all that Windows-y stuff, providing you
+    with a hook that is called when a CTRL-C event is intercepted.  This
+    hook allows you to unblock your ZMQ poll operation immediately, which
+    will then result in the expected ``KeyboardInterrupt`` exception.
+
+    Without this context manager, your ZMQ-based application will not
+    respond normally to CTRL-C events on Windows.  If a CTRL-C event occurs
+    while blocked on ZMQ socket polling, the translation to a
+    ``KeyboardInterrupt`` exception will be delayed until the I/O completes
+    and control returns to the Python interpreter (this may never happen if
+    you use an infinite timeout).
+
+    A no-op implementation is provided on non-Win32 systems to avoid the
+    application from having to conditionally use it.
+
+    Example usage:
+
+    .. sourcecode:: python
+
+       def stop_my_application():
+           # ...
+
+       with allow_interrupt(stop_my_application):
+           # main polling loop.
+
+    In a typical ZMQ application, you would use the "self pipe trick" to
+    send message to a ``PAIR`` socket in order to interrupt your blocking
+    socket polling operation.
+
+    In a Tornado event loop, you can use the ``IOLoop.stop`` method to
+    unblock your I/O loop.
+    """
+
+    def __init__(self, action: Optional[Callable[[], Any]] = None) -> None:
+        """Translate ``action`` into a CTRL-C handler.
+
+        ``action`` is a callable that takes no arguments and returns no
+        value (returned value is ignored).  It must *NEVER* raise an
+        exception.
+
+        If unspecified, a no-op will be used.
+        """
+        if os.name != "nt":
+            return
+        self._init_action(action)
+
+    def _init_action(self, action):
+        from ctypes import WINFUNCTYPE, windll
+        from ctypes.wintypes import BOOL, DWORD
+
+        kernel32 = windll.LoadLibrary("kernel32")
+
+        # <http://msdn.microsoft.com/en-us/library/ms686016.aspx>
+        PHANDLER_ROUTINE = WINFUNCTYPE(BOOL, DWORD)
+        SetConsoleCtrlHandler = self._SetConsoleCtrlHandler = (
+            kernel32.SetConsoleCtrlHandler
+        )
+        SetConsoleCtrlHandler.argtypes = (PHANDLER_ROUTINE, BOOL)
+        SetConsoleCtrlHandler.restype = BOOL
+
+        if action is None:
+            action = lambda: None  # noqa: E731
+        self.action = action
+
+        @PHANDLER_ROUTINE
+        def handle(event):
+            if event == 0:  # CTRL_C_EVENT
+                action()
+                # Typical C implementations would return 1 to indicate that
+                # the event was processed and other control handlers in the
+                # stack should not be executed.  However, that would
+                # prevent the Python interpreter's handler from translating
+                # CTRL-C to a `KeyboardInterrupt` exception, so we pretend
+                # that we didn't handle it.
+            return 0
+
+        self.handle = handle
+
+    def __enter__(self):
+        """Install the custom CTRL-C handler."""
+        if os.name != "nt":
+            return
+        result = self._SetConsoleCtrlHandler(self.handle, 1)
+        if result == 0:
+            # Have standard library automatically call `GetLastError()` and
+            # `FormatMessage()` into a nice exception object :-)
+            raise OSError()
+
+    def __exit__(self, *args):
+        """Remove the custom CTRL-C handler."""
+        if os.name != "nt":
+            return
+        result = self._SetConsoleCtrlHandler(self.handle, 0)
+        if result == 0:
+            # Have standard library automatically call `GetLastError()` and
+            # `FormatMessage()` into a nice exception object :-)
+            raise OSError()

truefoundry/deploy/v2/lib/__init__.py

@@ -0,0 +1,3 @@
+# NOTE: Later all of these can be under a `models` or some other module
+# I am not entirely sure about the structure of `lib` module yet
+# This can go through another round of refactoring

truefoundry/deploy/v2/lib/deploy.py

@@ -0,0 +1,232 @@
+import sys
+import time
+from typing import List, Optional, TypeVar
+
+from rich.status import Status
+
+from truefoundry.deploy.auto_gen import models as auto_gen_models
+from truefoundry.deploy.builder.docker_service import env_has_docker
+from truefoundry.deploy.lib.clients.servicefoundry_client import (
+    ServiceFoundryServiceClient,
+)
+from truefoundry.deploy.lib.clients.utils import poll_for_function
+from truefoundry.deploy.lib.dao.workspace import get_workspace_by_fqn
+from truefoundry.deploy.lib.model.entity import Deployment, DeploymentTransitionStatus
+from truefoundry.deploy.lib.util import get_application_fqn_from_deployment_fqn
+from truefoundry.deploy.v2.lib.models import BuildResponse
+from truefoundry.deploy.v2.lib.source import (
+    local_source_to_image,
+    local_source_to_remote_source,
+)
+from truefoundry.logger import logger
+from truefoundry.pydantic_v1 import BaseModel
+
+Component = TypeVar("Component", bound=BaseModel)
+
+
+def _handle_if_local_source(component: Component, workspace_fqn: str) -> Component:
+    if (
+        hasattr(component, "image")
+        and isinstance(component.image, auto_gen_models.Build)
+        and isinstance(component.image.build_source, auto_gen_models.LocalSource)
+    ):
+        new_component = component.copy(deep=True)
+
+        if component.image.build_source.local_build:
+            if not env_has_docker():
+                logger.warning(
+                    "Did not find Docker locally installed on this system, image will be built remotely. "
+                    "For faster builds it is recommended to install Docker locally. "
+                    "If you always want to build remotely, "
+                    "please set `image.build_source.local_build` to `false` in your YAML spec or equivalently set "
+                    "`image=Build(build_source=LocalSource(local_build=False, ...))` in your "
+                    "`Service` or `Job` definition code."
+                )
+                local_build = False
+            else:
+                logger.info(
+                    "Found locally installed docker, image will be built locally and then pushed. "
+                    "If you want to always build remotely instead of locally, "
+                    "please set `image.build_source.local_build` to `false` in your YAML spec or equivalently set "
+                    "`image=Build(build_source=LocalSource(local_build=False, ...))` in your "
+                    "`Service` or `Job` definition code."
+                )
+                local_build = True
+        else:
+            logger.info(
+                "Image will be built remotely because `image.build_source.local_build` is set to `false`. "
+                "For faster builds it is recommended to install Docker locally and "
+                "set `image.build_source.local_build` to `true` in your YAML spec "
+                "or equivalently set `image=Build(build_source=LocalSource(local_build=True, ...))` "
+                "in your `Service` or `Job` definition code."
+            )
+            local_build = False
+
+        if local_build:
+            # We are to build the image locally, push and update `image` in spec
+            logger.info("Building image for %s '%s'", component.type, component.name)
+            new_component.image = local_source_to_image(
+                build=component.image,
+                docker_registry_fqn=component.image.docker_registry,
+                workspace_fqn=workspace_fqn,
+                component_name=component.name,
+            )
+        else:
+            # We'll build image on Truefoundry servers, upload the source and update image.build_source
+            logger.info("Uploading code for %s '%s'", component.type, component.name)
+            new_component.image.build_source = local_source_to_remote_source(
+                local_source=component.image.build_source,
+                workspace_fqn=workspace_fqn,
+                component_name=component.name,
+            )
+            logger.debug("Uploaded code for %s '%s'", component.type, component.name)
+        return new_component
+    return component
+
+
+def _log_application_dashboard_url(deployment: Deployment, log_message: str):
+    application_id = deployment.applicationId
+
+    # TODO: is there any simpler way to get this? :cry
+    client = ServiceFoundryServiceClient()
+
+    url = f"{client.base_url.strip('/')}/applications/{application_id}?tab=deployments"
+    logger.info(log_message, url)
+
+
+def _tail_build_logs(build_responses: List[BuildResponse]) -> None:
+    client = ServiceFoundryServiceClient()
+
+    # TODO: Explore other options like,
+    # https://rich.readthedocs.io/en/stable/live.html#live-display
+    # How does docker/compose does multiple build logs?
+    for build_response in build_responses:
+        logger.info("Tailing build logs for '%s'", build_response.componentName)
+        client.tail_build_logs(build_response=build_response, wait=True)
+
+
+def _deploy_wait_handler(
+    deployment: Deployment,
+) -> Optional[DeploymentTransitionStatus]:
+    _log_application_dashboard_url(
+        deployment=deployment,
+        log_message=(
+            "You can track the progress below or on the dashboard:- '%s'\n"
+            "You can press Ctrl + C to exit the tailing of build logs "
+            "and deployment will continue on the server"
+        ),
+    )
+    with Status(status="Polling for deployment status") as spinner:
+        last_status_printed = None
+        client = ServiceFoundryServiceClient()
+        start_time = time.monotonic()
+        total_timeout_time: int = 300
+        poll_interval_seconds = 5
+        time_elapsed = 0
+
+        for deployment_statuses in poll_for_function(
+            client.get_deployment_statuses,
+            poll_after_secs=poll_interval_seconds,
+            application_id=deployment.applicationId,
+            deployment_id=deployment.id,
+        ):
+            if len(deployment_statuses) == 0:
+                logger.warning("Did not receive any deployment status")
+                continue
+
+            latest_deployment_status = deployment_statuses[-1]
+
+            status_to_print = (
+                latest_deployment_status.transition or latest_deployment_status.status
+            )
+            spinner.update(status=f"Current state: {status_to_print}")
+            if status_to_print != last_status_printed:
+                if DeploymentTransitionStatus.is_failure_state(status_to_print):
+                    logger.error("State: %r", status_to_print)
+                else:
+                    logger.info("State: %r", status_to_print)
+                last_status_printed = status_to_print
+
+            if latest_deployment_status.state.isTerminalState:
+                break
+
+            if (
+                latest_deployment_status.transition
+                == DeploymentTransitionStatus.BUILDING
+            ):
+                build_responses = client.get_deployment_build_response(
+                    application_id=deployment.applicationId, deployment_id=deployment.id
+                )
+                _tail_build_logs(build_responses)
+
+            time_elapsed = time.monotonic() - start_time
+            if time_elapsed > total_timeout_time:
+                logger.warning("Polled server for %s secs.", int(time_elapsed))
+                break
+
+    return last_status_printed
+
+
+def _warn_when_gpu_selected_without_cuda(component: Component):
+    is_python_build_without_cuda = (
+        hasattr(component, "image")
+        and isinstance(component.image, auto_gen_models.Build)
+        and isinstance(component.image.build_spec, auto_gen_models.PythonBuild)
+        and not component.image.build_spec.cuda_version
+    )
+    uses_gpu = (
+        hasattr(component, "resources")
+        and isinstance(component, auto_gen_models.Resources)
+        and component.resources.gpu_count > 0
+    )
+    if is_python_build_without_cuda and uses_gpu:
+        logger.warning(
+            "Warning: `gpu_count` is greater than 0 in `Resources` (i.e. `resources.gpu_count`) "
+            "but no `cuda_version` was passed to `PythonBuild` "
+            "(i.e. `image.build_spec.cuda_version`). "
+            "Your application might optionally need CUDA toolkit installed "
+            "to utilize the GPU. You can choose one by passing one of "
+            "`truefoundry.deploy.CUDAVersion` in `PythonBuild` instance. "
+            "\n\nE.g.\n```\nPythonBuild(..., cuda_version=CUDAVersion.CUDA_11_3_CUDNN8)\n```"
+        )
+
+
+def deploy_component(
+    component: Component, workspace_fqn: str, wait: bool = True
+) -> Deployment:
+    _warn_when_gpu_selected_without_cuda(component=component)
+    workspace_id = get_workspace_by_fqn(workspace_fqn).id
+    updated_component = _handle_if_local_source(
+        component=component, workspace_fqn=workspace_fqn
+    )
+    client = ServiceFoundryServiceClient()
+    response = client.deploy_application(
+        workspace_id=workspace_id, application=updated_component
+    )
+    logger.info(
+        "🚀 Deployment started for application '%s'. Deployment FQN is '%s'.",
+        updated_component.name,
+        response.fqn,
+    )
+    if wait:
+        try:
+            last_status_printed = _deploy_wait_handler(deployment=response)
+            if not last_status_printed or DeploymentTransitionStatus.is_failure_state(
+                last_status_printed
+            ):
+                deployment_tab_url = f"{client.base_url.strip('/')}/applications/{response.applicationId}?tab=deployments"
+                message = f"Deployment Failed. Please refer to the logs for additional details - {deployment_tab_url}"
+                sys.exit(message)
+        except KeyboardInterrupt:
+            logger.info("Ctrl-c executed. The deployment will still continue.")
+
+    deployment_fqn = response.fqn
+    application_fqn = get_application_fqn_from_deployment_fqn(deployment_fqn)
+    logger.info("Deployment FQN: %s", deployment_fqn)
+    logger.info("Application FQN: %s", application_fqn)
+
+    _log_application_dashboard_url(
+        deployment=response,
+        log_message="You can find the application on the dashboard:- '%s'",
+    )
+    return response

truefoundry/deploy/v2/lib/deployable_patched_models.py

@@ -0,0 +1,72 @@
+from typing import Literal, Union
+
+from truefoundry.deploy.auto_gen import models
+from truefoundry.deploy.lib.model.entity import Deployment
+from truefoundry.deploy.v2.lib.deploy import deploy_component
+from truefoundry.pydantic_v1 import BaseModel, Field, conint
+
+
+class DeployablePatchedModelBase(BaseModel):
+    class Config:
+        extra = "forbid"
+
+    def deploy(self, workspace_fqn: str, wait: bool = True) -> Deployment:
+        return deploy_component(component=self, workspace_fqn=workspace_fqn, wait=wait)
+
+
+class Service(models.Service, DeployablePatchedModelBase):
+    type: Literal["service"] = "service"
+    resources: models.Resources = Field(default_factory=models.Resources)
+    # This is being patched because cue export marks this as a "number"
+    replicas: Union[conint(ge=0, le=100), models.ServiceAutoscaling] = Field(
+        1,
+        description="+label=Replicas\n+usage=Replicas of service you want to run\n+icon=fa-clone\n+sort=3",
+    )
+
+
+class Job(models.Job, DeployablePatchedModelBase):
+    type: Literal["job"] = "job"
+    resources: models.Resources = Field(default_factory=models.Resources)
+
+
+class Notebook(models.Notebook, DeployablePatchedModelBase):
+    type: Literal["notebook"] = "notebook"
+    resources: models.Resources = Field(default_factory=models.Resources)
+
+
+class Codeserver(models.Codeserver, DeployablePatchedModelBase):
+    type: Literal["codeserver"] = "codeserver"
+
+
+class Helm(models.Helm, DeployablePatchedModelBase):
+    type: Literal["helm"] = "helm"
+
+
+class Volume(models.Volume, DeployablePatchedModelBase):
+    type: Literal["volume"] = "volume"
+
+
+class ApplicationSet(models.ApplicationSet, DeployablePatchedModelBase):
+    type: Literal["application-set"] = "application-set"
+
+
+class AsyncService(models.AsyncService, DeployablePatchedModelBase):
+    type: Literal["async-service"] = "async-service"
+    replicas: Union[conint(ge=0, le=100), models.AsyncServiceAutoscaling] = 1
+    resources: models.Resources = Field(default_factory=models.Resources)
+
+
+class SSHServer(models.SSHServer, DeployablePatchedModelBase):
+    type: Literal["ssh-server"] = "ssh-server"
+    resources: models.Resources = Field(default_factory=models.Resources)
+
+
+class Application(models.Application, DeployablePatchedModelBase):
+    def deploy(self, workspace_fqn: str, wait: bool = True) -> Deployment:
+        return deploy_component(
+            component=self.__root__, workspace_fqn=workspace_fqn, wait=wait
+        )
+
+
+class Workflow(models.Workflow, DeployablePatchedModelBase):
+    type: Literal["workflow"] = "workflow"

truefoundry/deploy/v2/lib/models.py

@@ -0,0 +1,53 @@
+import datetime
+from typing import List, Optional
+
+from truefoundry.pydantic_v1 import BaseModel, Extra, create_model
+
+
+class BuildResponse(BaseModel):
+    id: str
+    name: str
+    # TODO: make status an enum
+    status: str
+    # TODO: should we just make these fields
+    # snake-case and add camelCase aliases?
+    deploymentId: str
+    componentName: str
+    createdAt: datetime.datetime
+    updatedAt: datetime.datetime
+    imageUri: Optional[str]
+    failureReason: Optional[str]
+    getLogsUrl: str
+    tailLogsUrl: str
+    logsStartTs: int
+
+    class Config:
+        extra = Extra.allow
+
+
+class AppDeploymentStatusResponse(BaseModel):
+    state: create_model(
+        "State",
+        isTerminalState=(bool, ...),
+        type=(str, ...),
+        transitions=(List[str], ...),
+    )
+
+    id: str
+    status: str
+    message: Optional[str]
+    transition: Optional[str]
+
+    class Config:
+        extra = Extra.allow
+
+
+class DeploymentFqnResponse(BaseModel):
+    deploymentId: str
+    applicationId: str
+    workspaceId: str
+
+
+class ApplicationFqnResponse(BaseModel):
+    applicationId: str
+    workspaceId: str