truefoundry 0.6.3rc2__tar.gz → 0.6.3rc3__tar.gz

{truefoundry-0.6.3rc2 → truefoundry-0.6.3rc3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: truefoundry
-Version: 0.6.3rc2
+Version: 0.6.3rc3
 Summary: TrueFoundry CLI
 Author-email: TrueFoundry Team <abhishek@truefoundry.com>
 Requires-Python: <3.14,>=3.8.1

{truefoundry-0.6.3rc2 → truefoundry-0.6.3rc3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "truefoundry"
-version = "0.6.3rc2"  # This is a placeholder version, it will be replaced when releasing
+version = "0.6.3rc3"  # This is a placeholder version, it will be replaced when releasing
 requires-python = ">=3.8.1,<3.14"
 description = "TrueFoundry CLI"
 authors = [{ name = "TrueFoundry Team", email = "abhishek@truefoundry.com" }]

{truefoundry-0.6.3rc2 → truefoundry-0.6.3rc3}/truefoundry/common/utils.py

@@ -93,7 +93,10 @@ def timed_lru_cache(
 
 
 def poll_for_function(
-    func: Callable[..., T], poll_after_secs: int = 5, *args, **kwargs
+    func: Callable[..., T],
+    poll_after_secs: int = 5,
+    *args,
+    **kwargs,
 ) -> Generator[T, None, None]:
     while True:
         yield func(*args, **kwargs)

{truefoundry-0.6.3rc2 → truefoundry-0.6.3rc3}/truefoundry/deploy/__init__.py

@@ -14,6 +14,7 @@ from truefoundry.deploy._autogen.models import (
     SparkExecutorDynamicScaling,
     SparkExecutorFixedInstances,
     WorkbenchImage,
+    WorkflowAlert,
 )
 from truefoundry.deploy.lib.dao.application import (
     delete_application,
@@ -72,6 +73,7 @@ from truefoundry.deploy.v2.lib.patched_models import (
     CUDAVersion,
     DockerFileBuild,
     DynamicVolumeConfig,
+    Email,
     Endpoint,
     GcpTPU,
     GitHelmRepo,
@@ -109,6 +111,7 @@ from truefoundry.deploy.v2.lib.patched_models import (
     Schedule,
     SecretMount,
     ServiceAutoscaling,
+    SlackWebhook,
     SQSInputConfig,
     SQSOutputConfig,
     SQSQueueMetricConfig,

{truefoundry-0.6.3rc2 → truefoundry-0.6.3rc3}/truefoundry/deploy/_autogen/models.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  application.json
-#   timestamp: 2025-03-20T05:09:43+00:00
+#   timestamp: 2025-03-27T11:40:57+00:00
 
 from __future__ import annotations
 
@@ -212,6 +212,17 @@ class DynamicVolumeConfig(BaseModel):
     size: conint(ge=1, le=64000) = Field(..., description="Size of volume in Gi")
 
 
+class Email(BaseModel):
+    type: Literal["email"] = Field(..., description="")
+    notification_channel: constr(min_length=1) = Field(
+        ..., description="Specify the notification channel to send alerts to"
+    )
+    to_emails: List[constr(min_length=1)] = Field(
+        ...,
+        description="List of recipients' email addresses if the notification channel is Email.",
+    )
+
+
 class Endpoint(BaseModel):
     host: constr(
         regex=r"^((([a-zA-Z0-9\-]{1,63}\.)([a-zA-Z0-9\-]{1,63}\.)*([A-Za-z]{1,63}))|(((25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)))$"
@@ -349,25 +360,6 @@ class Image(BaseModel):
     )
 
 
-class JobAlert(BaseModel):
-    """
-    Describes the configuration for the job alerts
-    """
-
-    notification_channel: constr(min_length=1) = Field(
-        ..., description="Specify the notification channel to send alerts to"
-    )
-    to_emails: Optional[List[constr(min_length=1)]] = Field(
-        None,
-        description="List of recipients' email addresses if the notification channel is Email.",
-    )
-    on_start: bool = Field(False, description="Send an alert when the job starts")
-    on_completion: bool = Field(
-        False, description="Send an alert when the job completes"
-    )
-    on_failure: bool = Field(True, description="Send an alert when the job fails")
-
-
 class Claim(BaseModel):
     key: str
     values: List[str]
@@ -837,6 +829,13 @@ class ServiceAutoscaling(BaseAutoscaling):
     )
 
 
+class SlackWebhook(BaseModel):
+    type: Literal["slack-webhook"] = Field(..., description="")
+    notification_channel: constr(min_length=1) = Field(
+        ..., description="Specify the notification channel to send alerts to"
+    )
+
+
 class SparkDriverConfig(BaseModel):
     ui_endpoint: Endpoint
     resources: Optional[Resources] = None
@@ -1174,59 +1173,6 @@ class Helm(BaseModel):
     )
 
 
-class Job(BaseModel):
-    """
-    Describes the configuration for the job
-    """
-
-    type: Literal["job"] = Field(..., description="")
-    name: constr(regex=r"^[a-z](?:[a-z0-9]|-(?!-)){1,30}[a-z0-9]$") = Field(
-        ..., description="Name of the job"
-    )
-    image: Union[Build, Image] = Field(
-        ...,
-        description="Specify whether you want to deploy a Docker image or build and deploy from source code",
-    )
-    trigger: Union[Manual, Schedule] = Field(
-        {"type": "manual"}, description="Specify the trigger"
-    )
-    trigger_on_deploy: bool = Field(
-        False, description="Trigger the job after deploy immediately"
-    )
-    params: Optional[List[Param]] = Field(
-        None, description="Configure params and pass it to create different job runs"
-    )
-    env: Optional[Dict[str, str]] = Field(
-        None,
-        description="Configure environment variables to be injected in the service either as plain text or secrets. [Docs](https://docs.truefoundry.com/docs/env-variables)",
-    )
-    resources: Optional[Resources] = None
-    alerts: Optional[List[JobAlert]] = Field(
-        None,
-        description="Configure alerts to be sent when the job starts/fails/completes",
-    )
-    retries: conint(ge=0, le=10) = Field(
-        0,
-        description="Specify the maximum number of attempts to retry a job before it is marked as failed.",
-    )
-    timeout: Optional[conint(le=432000, gt=0)] = Field(
-        None, description="Job timeout in seconds."
-    )
-    concurrency_limit: Optional[PositiveInt] = Field(
-        None, description="Number of runs that can run concurrently"
-    )
-    service_account: Optional[str] = Field(None, description="")
-    mounts: Optional[List[Union[SecretMount, StringDataMount, VolumeMount]]] = Field(
-        None,
-        description="Configure data to be mounted to job pod(s) as a string, secret or volume. [Docs](https://docs.truefoundry.com/docs/mounting-volumes-job)",
-    )
-    labels: Optional[Dict[str, str]] = Field(None, description="")
-    kustomize: Optional[Kustomize] = None
-    workspace_fqn: Optional[str] = Field(
-        None, description="Fully qualified name of the workspace"
-    )
-
-
 class KafkaInputConfig(BaseModel):
     """
     Describes the configuration for the input Kafka worker
@@ -1468,6 +1414,18 @@ class WorkerConfig(BaseModel):
     )
 
 
+class WorkflowAlert(BaseModel):
+    """
+    Describes the configuration for the workflow alerts
+    """
+
+    notification_target: Optional[Union[Email, SlackWebhook]] = None
+    on_completion: bool = Field(
+        False, description="Send an alert when the job completes"
+    )
+    on_failure: bool = Field(True, description="Send an alert when the job fails")
+
+
 class BaseService(BaseModel):
     name: constr(regex=r"^[a-z](?:[a-z0-9]|-(?!-)){1,30}[a-z0-9]$") = Field(
         ...,
@@ -1516,6 +1474,26 @@ class FlyteTaskTemplate(BaseModel):
     custom: FlyteTaskCustom
 
 
+class JobAlert(BaseModel):
+    """
+    Describes the configuration for the job alerts
+    """
+
+    notification_channel: Optional[constr(min_length=1)] = Field(
+        None, description="Specify the notification channel to send alerts to"
+    )
+    to_emails: Optional[List[constr(min_length=1)]] = Field(
+        None,
+        description="List of recipients' email addresses if the notification channel is Email.",
+    )
+    notification_target: Optional[Union[Email, SlackWebhook]] = None
+    on_start: bool = Field(False, description="Send an alert when the job starts")
+    on_completion: bool = Field(
+        False, description="Send an alert when the job completes"
+    )
+    on_failure: bool = Field(True, description="Send an alert when the job fails")
+
+
 class Service(BaseService):
     """
     Describes the configuration for the service
@@ -1557,6 +1535,59 @@ class FlyteTask(BaseModel):
     description: Optional[Any] = None
 
 
+class Job(BaseModel):
+    """
+    Describes the configuration for the job
+    """
+
+    type: Literal["job"] = Field(..., description="")
+    name: constr(regex=r"^[a-z](?:[a-z0-9]|-(?!-)){1,30}[a-z0-9]$") = Field(
+        ..., description="Name of the job"
+    )
+    image: Union[Build, Image] = Field(
+        ...,
+        description="Specify whether you want to deploy a Docker image or build and deploy from source code",
+    )
+    trigger: Union[Manual, Schedule] = Field(
+        {"type": "manual"}, description="Specify the trigger"
+    )
+    trigger_on_deploy: bool = Field(
+        False, description="Trigger the job after deploy immediately"
+    )
+    params: Optional[List[Param]] = Field(
+        None, description="Configure params and pass it to create different job runs"
+    )
+    env: Optional[Dict[str, str]] = Field(
+        None,
+        description="Configure environment variables to be injected in the service either as plain text or secrets. [Docs](https://docs.truefoundry.com/docs/env-variables)",
+    )
+    resources: Optional[Resources] = None
+    alerts: Optional[List[JobAlert]] = Field(
+        None,
+        description="Configure alerts to be sent when the job starts/fails/completes",
+    )
+    retries: conint(ge=0, le=10) = Field(
+        0,
+        description="Specify the maximum number of attempts to retry a job before it is marked as failed.",
+    )
+    timeout: Optional[conint(le=432000, gt=0)] = Field(
+        None, description="Job timeout in seconds."
+    )
+    concurrency_limit: Optional[PositiveInt] = Field(
+        None, description="Number of runs that can run concurrently"
+    )
+    service_account: Optional[str] = Field(None, description="")
+    mounts: Optional[List[Union[SecretMount, StringDataMount, VolumeMount]]] = Field(
+        None,
+        description="Configure data to be mounted to job pod(s) as a string, secret or volume. [Docs](https://docs.truefoundry.com/docs/mounting-volumes-job)",
+    )
+    labels: Optional[Dict[str, str]] = Field(None, description="")
+    kustomize: Optional[Kustomize] = None
+    workspace_fqn: Optional[str] = Field(
+        None, description="Fully qualified name of the workspace"
+    )
+
+
 class Workflow(BaseModel):
     """
     Describes the configuration for the worflow
@@ -1575,6 +1606,7 @@ class Workflow(BaseModel):
     flyte_entities: Optional[List[Union[FlyteTask, FlyteWorkflow, FlyteLaunchPlan]]] = (
         Field(None, description="")
     )
+    alerts: Optional[List[WorkflowAlert]] = Field(None, description="")
 
 
 class ApplicationSet(BaseModel):

{truefoundry-0.6.3rc2 → truefoundry-0.6.3rc3}/truefoundry/deploy/lib/clients/servicefoundry_client.py

@@ -14,7 +14,10 @@ from rich.status import Status
 from tqdm import tqdm
 from tqdm.utils import CallbackIOWrapper
 
-from truefoundry.common.constants import VERSION_PREFIX
+from truefoundry.common.constants import (
+    SERVICEFOUNDRY_CLIENT_MAX_RETRIES,
+    VERSION_PREFIX,
+)
 from truefoundry.common.request_utils import request_handling
 from truefoundry.common.servicefoundry_client import (
     ServiceFoundryServiceClient as BaseServiceFoundryServiceClient,
@@ -247,7 +250,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         callback: Optional[OutputCallBack] = None,
     ) -> socketio.Client:
         callback = callback or OutputCallBack()
-        sio = socketio.Client(request_timeout=60)
+        sio = socketio.Client(request_timeout=60, reconnection_attempts=10)
         callback.print_line("Waiting for the task to start...")
         next_log_start_timestamp = query_dict.get("startTs")
 
@@ -286,6 +289,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
                 transports="websocket",
                 headers=self._get_headers(),
                 socketio_path=socketio_path,
+                retry=True,
             )
             return sio
 
@@ -298,10 +302,15 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
 
     @check_min_cli_version
     def get_deployment_statuses(
-        self, application_id: str, deployment_id: str
+        self,
+        application_id: str,
+        deployment_id: str,
+        retry_count: int = SERVICEFOUNDRY_CLIENT_MAX_RETRIES,
     ) -> List[AppDeploymentStatusResponse]:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/app/{application_id}/deployments/{deployment_id}/statuses"
-        response = session_with_retries().get(url, headers=self._get_headers())
+        response = session_with_retries(retries=retry_count).get(
+            url, headers=self._get_headers()
+        )
         response_data = request_handling(response)
         return parse_obj_as(List[AppDeploymentStatusResponse], response_data)
 

{truefoundry-0.6.3rc2 → truefoundry-0.6.3rc3}/truefoundry/deploy/v2/lib/deploy.py

@@ -107,12 +107,20 @@ def _tail_build_logs(build_response: BuildResponse) -> socketio.Client:
 
 def _deploy_wait_handler(  # noqa: C901
     deployment: Deployment,
+    tail_logs: bool = True,
 ) -> Optional[DeploymentTransitionStatus]:
+    tail_logs_or_polling_status_message = (
+        "You can press Ctrl + C to exit the tailing of build logs "
+    )
+    if not tail_logs:
+        tail_logs_or_polling_status_message = (
+            "you can press Ctrl + C to exit the polling of deployment status "
+        )
     _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 "
+            f"{tail_logs_or_polling_status_message}"
             "and deployment will continue on the server"
         ),
     )
@@ -130,6 +138,7 @@ def _deploy_wait_handler(  # noqa: C901
             poll_after_secs=poll_interval_seconds,
             application_id=deployment.applicationId,
             deployment_id=deployment.id,
+            retry_count=10,
         ):
             if len(deployment_statuses) == 0:
                 logger.warning("Did not receive any deployment status")
@@ -157,12 +166,15 @@ def _deploy_wait_handler(  # noqa: C901
                 latest_deployment_status.transition
                 == DeploymentTransitionStatus.BUILDING
             ):
-                if not socket:
-                    build_responses = client.get_deployment_build_response(
-                        application_id=deployment.applicationId,
-                        deployment_id=deployment.id,
-                    )
-                    socket = _tail_build_logs(build_responses[0])
+                if tail_logs and not socket:
+                    try:
+                        build_responses = client.get_deployment_build_response(
+                            application_id=deployment.applicationId,
+                            deployment_id=deployment.id,
+                        )
+                        socket = _tail_build_logs(build_responses[0])
+                    except Exception as e:
+                        logger.error("Error tailing build logs: %s", e)
 
             time_elapsed = time.monotonic() - start_time
             if time_elapsed > total_timeout_time:

{truefoundry-0.6.3rc2 → truefoundry-0.6.3rc3}/truefoundry/deploy/v2/lib/deploy_workflow.py

@@ -25,15 +25,17 @@ from truefoundry.deploy.lib.clients.servicefoundry_client import (
     ServiceFoundryServiceClient,
 )
 from truefoundry.deploy.lib.dao.workspace import get_workspace_by_fqn
-from truefoundry.deploy.lib.model.entity import Deployment
+from truefoundry.deploy.lib.model.entity import Deployment, DeploymentTransitionStatus
+from truefoundry.deploy.v2.lib.deploy import _deploy_wait_handler
 from truefoundry.deploy.v2.lib.source import (
     local_source_to_remote_source,
 )
 from truefoundry.logger import logger
 from truefoundry.pydantic_v1 import ValidationError
 from truefoundry.workflow.workflow import (
+    TRUEFOUNDRY_ALERTS_CONFIG,
     TRUEFOUNDRY_LAUNCH_PLAN_NAME,
-    execution_config_store,
+    truefoundry_config_store,
 )
 
 
@@ -173,7 +175,7 @@ def _validate_workflow_entities(  # noqa: C901
     # validate that all inputs have default values for cron workflows
     for launch_plan in launch_plans:
         if (
-            execution_config_store.get(launch_plan.spec.workflow_id.name)
+            truefoundry_config_store.get(launch_plan.spec.workflow_id.name)
             and launch_plan.id.name == TRUEFOUNDRY_LAUNCH_PLAN_NAME
         ):
             workflow_inputs = launch_plan.spec.default_inputs.parameters
@@ -236,7 +238,7 @@ def _generate_manifest_for_workflow(
 
             # this is the case when someone has a cron schedule. and this line is for handling default launch plan in this case.
             if (
-                execution_config_store.get(workflow_name)
+                truefoundry_config_store.get(workflow_name)
                 and workflow_name == entity.id.name
             ):
                 continue
@@ -262,6 +264,14 @@ def _generate_manifest_for_workflow(
 
         workflow.flyte_entities.append(message_dict)
 
+    alerts_config = truefoundry_config_store.get(TRUEFOUNDRY_ALERTS_CONFIG)
+    if alerts_config:
+        if not workflow.alerts:
+            workflow.alerts = alerts_config
+        else:
+            logger.warning(
+                "Alerts are configured in both workflow decorator as well as in deployment config. Alerts configured in workflow decorator will be ignored."
+            )
     # this step is just to verify if pydantic model is still valid after adding flyte_entities
     autogen_models.Workflow.validate({**workflow.dict()})
 
@@ -288,7 +298,7 @@ def deploy_workflow(
 
     # we need to rest the execution config store as it is a global variable and we don't want to keep the cron execution config for next workflow
     # this is only needed for notebook environment
-    execution_config_store.reset()
+    truefoundry_config_store.reset()
 
     workspace_id = get_workspace_by_fqn(workspace_fqn).id
 
@@ -314,4 +324,17 @@ def deploy_workflow(
         deployment.fqn,
     )
     deployment_url = f"{client.tfy_host.strip('/')}/applications/{deployment.applicationId}?tab=deployments"
+    if wait:
+        try:
+            last_status_printed = _deploy_wait_handler(
+                deployment=deployment, tail_logs=False
+            )
+            if not last_status_printed or DeploymentTransitionStatus.is_failure_state(
+                last_status_printed
+            ):
+                deployment_tab_url = f"{client.tfy_host.strip('/')}/applications/{deployment.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.")
     logger.info("You can find the application on the dashboard:- '%s'", deployment_url)

{truefoundry-0.6.3rc2 → truefoundry-0.6.3rc3}/truefoundry/deploy/v2/lib/patched_models.py

@@ -563,3 +563,11 @@ class TaskPythonBuild(models.TaskPythonBuild, PatchedModelBase):
 
 # Deprecated aliases, kept for backward compatibility
 TruefoundryArtifactSource = TrueFoundryArtifactSource
+
+
+class Email(models.Email, PatchedModelBase):
+    type: Literal["email"] = "email"
+
+
+class SlackWebhook(models.SlackWebhook, PatchedModelBase):
+    type: Literal["slack-webhook"] = "slack-webhook"

{truefoundry-0.6.3rc2 → truefoundry-0.6.3rc3}/truefoundry/workflow/workflow.py

@@ -10,9 +10,11 @@ from flytekit.core.workflow import (
 )
 from flytekit.core.workflow import workflow as flytekit_workflow
 
+from truefoundry.deploy._autogen.models import WorkflowAlert
 from truefoundry.pydantic_v1 import BaseModel
 
 TRUEFOUNDRY_LAUNCH_PLAN_NAME = "default"
+TRUEFOUNDRY_ALERTS_CONFIG = "tfy_alerts_config"
 
 
 class ExecutionConfig(BaseModel):
@@ -22,7 +24,7 @@ class ExecutionConfig(BaseModel):
     schedule: str
 
 
-class ExecutionConfigStore(BaseModel):
+class TruefoundryConfigStore(BaseModel):
     execution_config_map: Dict[str, ExecutionConfig] = {}
 
     def reset(self):
@@ -35,7 +37,7 @@ class ExecutionConfigStore(BaseModel):
         self.execution_config_map[key] = value
 
 
-execution_config_store = ExecutionConfigStore()
+truefoundry_config_store = TruefoundryConfigStore()
 
 
 def workflow(
@@ -43,6 +45,7 @@ def workflow(
     failure_policy: Optional[WorkflowFailurePolicy] = None,
     on_failure: Optional[Task] = None,
     execution_configs: Optional[List[ExecutionConfig]] = None,
+    alerts: Optional[List[WorkflowAlert]] = None,
 ) -> Union[
     Callable[[Callable[..., FuncOut]], PythonFunctionWorkflow],
     PythonFunctionWorkflow,
@@ -85,6 +88,7 @@ def workflow(
                 failure_policy=failure_policy,
                 on_failure=on_failure,
                 execution_configs=execution_configs,
+                alerts=alerts,
             )
 
         return wrapper
@@ -109,6 +113,9 @@ def workflow(
         function_module = _workflow_function.__module__
         function_name = _workflow_function.__name__
         function_path = f"{function_module}.{function_name}"
-        execution_config_store.set(function_path, execution_config)
+        truefoundry_config_store.set(function_path, execution_config)
+
+    if alerts:
+        truefoundry_config_store.set(TRUEFOUNDRY_ALERTS_CONFIG, alerts)
 
     return workflow_object