truefoundry 0.6.2__py3-none-any.whl → 0.6.3__py3-none-any.whl

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/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/deploy/_autogen/models.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  application.json
-#   timestamp: 2025-03-10T17:05:18+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]
@@ -389,6 +381,10 @@ class JwtAuthConfig(BaseModel):
     claims: Optional[List[Claim]] = Field(
         None, description="List of key-value pairs of claims to verify in the JWT token"
     )
+    bypass_auth_paths: Optional[List[constr(regex=r"^/[^*]*")]] = Field(
+        None,
+        description="List of paths that will bypass auth.\nneeds to start with a forward slash(/) and should not contain wildcards(*)",
+    )
 
 
 class KafkaMetricConfig(BaseModel):
@@ -510,15 +506,36 @@ class NvidiaGPU(BaseModel):
 
 class Profile(str, Enum):
     """
-    Name of the MIG profile to use. One of [1g.5gb, 2g.10gb, 3g.20gb, 1g.10gb, 2g.20gb, 3g.40gb]
+    Name of the MIG profile to use. One of the following based on gpu type
+    Please refer to https://docs.nvidia.com/datacenter/tesla/mig-user-guide/#supported-mig-profiles for more details
+    A100 40 GB - [1g.5gb, 1g.10gb, 2g.10gb, 3g.20gb, 4g.20gb]
+    A100 80 GB / H100 80 GB - [1g.10gb, 1g.20gb, 2g.20gb, 3g.40gb, 4g.40gb]
+    H100 94 GB - [1g.12gb, 1g.24gb, 2g.24gb, 3g.47gb, 4g.47gb]
+    H100 96 GB - [1g.12gb, 1g.24gb, 2g.24gb, 3g.48gb, 4g.48gb]
+    H200 141 GB - [1g.18gb, 1g.35gb, 2g.35gb, 3g.71gb, 4g.71gb]
     """
 
     field_1g_5gb = "1g.5gb"
-    field_2g_10gb = "2g.10gb"
-    field_3g_20gb = "3g.20gb"
     field_1g_10gb = "1g.10gb"
+    field_1g_12gb = "1g.12gb"
+    field_1g_18gb = "1g.18gb"
+    field_1g_20gb = "1g.20gb"
+    field_1g_24gb = "1g.24gb"
+    field_1g_35gb = "1g.35gb"
+    field_2g_10gb = "2g.10gb"
     field_2g_20gb = "2g.20gb"
+    field_2g_24gb = "2g.24gb"
+    field_2g_35gb = "2g.35gb"
+    field_3g_20gb = "3g.20gb"
     field_3g_40gb = "3g.40gb"
+    field_3g_47gb = "3g.47gb"
+    field_3g_48gb = "3g.48gb"
+    field_3g_71gb = "3g.71gb"
+    field_4g_20gb = "4g.20gb"
+    field_4g_40gb = "4g.40gb"
+    field_4g_47gb = "4g.47gb"
+    field_4g_48gb = "4g.48gb"
+    field_4g_71gb = "4g.71gb"
 
 
 class NvidiaMIGGPU(BaseModel):
@@ -529,7 +546,7 @@ class NvidiaMIGGPU(BaseModel):
     )
     profile: Profile = Field(
         ...,
-        description="Name of the MIG profile to use. One of [1g.5gb, 2g.10gb, 3g.20gb, 1g.10gb, 2g.20gb, 3g.40gb]",
+        description="Name of the MIG profile to use. One of the following based on gpu type\nPlease refer to https://docs.nvidia.com/datacenter/tesla/mig-user-guide/#supported-mig-profiles for more details\nA100 40 GB - [1g.5gb, 1g.10gb, 2g.10gb, 3g.20gb, 4g.20gb]\nA100 80 GB / H100 80 GB - [1g.10gb, 1g.20gb, 2g.20gb, 3g.40gb, 4g.40gb]\nH100 94 GB - [1g.12gb, 1g.24gb, 2g.24gb, 3g.47gb, 4g.47gb]\nH100 96 GB - [1g.12gb, 1g.24gb, 2g.24gb, 3g.48gb, 4g.48gb]\nH200 141 GB - [1g.18gb, 1g.35gb, 2g.35gb, 3g.71gb, 4g.71gb]",
     )
 
 
@@ -812,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
@@ -918,6 +942,10 @@ class TrueFoundryArtifactSource(BaseModel):
 
 class TrueFoundryInteractiveLogin(BaseModel):
     type: Literal["truefoundry_oauth"] = Field(..., description="")
+    bypass_auth_paths: Optional[List[constr(regex=r"^/[^*]*")]] = Field(
+        None,
+        description="List of paths that will bypass auth.\nneeds to start with a forward slash(/) and should not contain wildcards(*)",
+    )
 
 
 class VolumeBrowser(BaseModel):
@@ -1145,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
@@ -1439,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(
         ...,
@@ -1487,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
@@ -1528,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
@@ -1546,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/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/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/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,
 )
 
 
@@ -132,6 +134,7 @@ def _validate_workflow_entities(  # noqa: C901
         `from truefoundry.workflow import PythonTaskConfig, ContainerTaskConfig`
          """
     tasks_without_truefoundry_worflow_package = []
+    task_names = set()
     for task in tasks:
         if _is_dynamic_task(task):
             raise ValueError("Dynamic workflows are not supported yet.")
@@ -141,6 +144,12 @@ def _validate_workflow_entities(  # noqa: C901
                     task.template.id.name
                 )
             )
+        task_name = task.template.id.name
+        if task_name in task_names:
+            raise ValueError(
+                f"Task name should be unique, task with name {task_name} is repeated in the workflow"
+            )
+        task_names.add(task_name)
         task_image_spec = task.template.custom["truefoundry"]["image"]
         if task_image_spec["type"] == "task-python-build":
             is_tfy_wf_present_in_task_python_build = (
@@ -166,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
@@ -229,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
@@ -255,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()})
 
@@ -281,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
 
@@ -307,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/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/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

{truefoundry-0.6.2.dist-info → truefoundry-0.6.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: truefoundry
-Version: 0.6.2
+Version: 0.6.3
 Summary: TrueFoundry CLI
 Author-email: TrueFoundry Team <abhishek@truefoundry.com>
 Requires-Python: <3.14,>=3.8.1
@@ -30,7 +30,7 @@ Requires-Dist: requirements-parser<0.12.0,>=0.11.0
 Requires-Dist: rich-click<2.0.0,>=1.2.1
 Requires-Dist: rich<14.0.0,>=13.7.1
 Requires-Dist: tqdm<5.0.0,>=4.0.0
-Requires-Dist: truefoundry-sdk==0.0.6
+Requires-Dist: truefoundry-sdk==0.0.8
 Requires-Dist: typing-extensions>=4.0
 Requires-Dist: urllib3<3,>=1.26.18
 Requires-Dist: yq<4.0.0,>=3.1.0

{truefoundry-0.6.2.dist-info → truefoundry-0.6.3.dist-info}/RECORD

@@ -46,11 +46,11 @@ truefoundry/common/servicefoundry_client.py,sha256=2fYhdVPSvLXz5C5tosOq86JD8WM3I
 truefoundry/common/session.py,sha256=xeBAPUNEJv2XVFQCRUGeBDTePh5zrKNSok8vmSxBjPw,2813
 truefoundry/common/storage_provider_utils.py,sha256=yURhMw8k0FLFvaviRHDiifhvc6GnuQwGMC9Qd2uM440,10934
 truefoundry/common/types.py,sha256=BMJFCsR1lPJAw66IQBSvLyV4I6o_x5oj78gVsUa9si8,188
-truefoundry/common/utils.py,sha256=Ev4lATS0jT9BDNPmQEP3aDn8IIeh9pmahv5pcNIjuuw,6366
+truefoundry/common/utils.py,sha256=yEQtJW2fT9xbNpRhfRoD9hvhGw-FgGS3agh1oZptmjg,6379
 truefoundry/common/warnings.py,sha256=rs6BHwk7imQYedo07iwh3TWEOywAR3Lqhj0AY4khByg,504
-truefoundry/deploy/__init__.py,sha256=lso6tFZVNyB3KmjEZICq_Sd6R_Yip-67XAPXYAQCvGw,2597
+truefoundry/deploy/__init__.py,sha256=zNsyzJAOPCPqlhVEEq6sBpfPsa4XkChaXvE-EBCmfzs,2645
 truefoundry/deploy/python_deploy_codegen.py,sha256=AainOFR20XvhNeztJkLPWGZ40lAT_nwc-ZmG77Kum4o,6525
-truefoundry/deploy/_autogen/models.py,sha256=oi0Q80dQsXCMVL3SJudogxaBVdIOvBQDEaOi791Dj-4,66926
+truefoundry/deploy/_autogen/models.py,sha256=0m43YJVG0fBhLp0fPkCod_vIcMFob0jTQEGnjk6nOUQ,69804
 truefoundry/deploy/builder/__init__.py,sha256=nGQiR3r16iumRy7xbVQ6q-k0EApmijspsfVpXDE-9po,4953
 truefoundry/deploy/builder/constants.py,sha256=amUkHoHvVKzGv0v_knfiioRuKiJM0V0xW0diERgWiI0,508
 truefoundry/deploy/builder/docker_service.py,sha256=sm7GWeIqyrKaZpxskdLejZlsxcZnM3BTDJr6orvPN4E,3948
@@ -93,7 +93,7 @@ truefoundry/deploy/lib/session.py,sha256=-FX4gOtiGlc6Jk56JPVZpDqXR9xQza77AIlBvNJ
 truefoundry/deploy/lib/util.py,sha256=J7r8San2wKo48A7-BlH2-OKTlBO67zlPjLEhMsL8os0,1059
 truefoundry/deploy/lib/win32.py,sha256=1RcvPTdlOAJ48rt8rCbE2Ufha2ztRqBAE9dueNXArrY,5009
 truefoundry/deploy/lib/clients/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-truefoundry/deploy/lib/clients/servicefoundry_client.py,sha256=DO2uSfrlX_udIAQ6YM_l59OGBA75mhp3-6phbdLVSPc,26475
+truefoundry/deploy/lib/clients/servicefoundry_client.py,sha256=NjNGmBsmrIfDZFfqJz12B4TVDHMZeAZDf-5eKSnS9zs,26697
 truefoundry/deploy/lib/dao/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truefoundry/deploy/lib/dao/application.py,sha256=oMszpueXPUfTUuN_XdKwoRjQyqAgWHhZ-10cbprCVdM,9226
 truefoundry/deploy/lib/dao/apply.py,sha256=5IFERe5sLmZGlavaKTIxL4xPHAme4ZS2Ww0a2rKTyT0,3029
@@ -104,11 +104,11 @@ truefoundry/deploy/lib/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NM
 truefoundry/deploy/lib/model/entity.py,sha256=Up-DDOezkwM2tdqibfLdZO6jmT2pVq6SShB5sobBIGI,8531
 truefoundry/deploy/v2/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truefoundry/deploy/v2/lib/__init__.py,sha256=WEiVMZXOVljzEE3tpGJil14liIn_PCDoACJ6b3tZ6sI,188
-truefoundry/deploy/v2/lib/deploy.py,sha256=23Sc2OxtAMvDzKfNmVv5vuaB6GqP-e1WAZWIezMzm3c,12413
-truefoundry/deploy/v2/lib/deploy_workflow.py,sha256=7NlL5eASwsDLcmsWhhh-YR-yBxS1CPwM-ceP87lNm_M,12860
+truefoundry/deploy/v2/lib/deploy.py,sha256=FWmqX4N6RUN4-S68i-7xaIifYWnf0wc-zfINdpiFu-4,12901
+truefoundry/deploy/v2/lib/deploy_workflow.py,sha256=G5BzMIbap8pgDX1eY-TITruUxQdkKhYtBmRwLL6lDeY,14342
 truefoundry/deploy/v2/lib/deployable_patched_models.py,sha256=xbHFD3pURflvCm8EODPvjfvRrv67mlSrjPUknY8SMB8,4060
 truefoundry/deploy/v2/lib/models.py,sha256=ogc1UYs1Z2nBdGSKCrde9sk8d0GxFKMkem99uqO5CmM,1148
-truefoundry/deploy/v2/lib/patched_models.py,sha256=Lowm-tyCGGv6-4iY9vfgmHrr4zPFy2u9wZEY3xQBv4E,16359
+truefoundry/deploy/v2/lib/patched_models.py,sha256=vVjYs1Gm7mpuTx3C0l40RQ_zvMdQ1S0s6J-f00qO0nA,16557
 truefoundry/deploy/v2/lib/source.py,sha256=d6-8_6Zn5koBglqrBrY6ZLG_7yyPuLdyEmK4iZTw6xY,9405
 truefoundry/ml/__init__.py,sha256=EEEHV7w58Krpo_W9Chd8Y3TdItfFO3LI6j6Izqc4-P8,2219
 truefoundry/ml/constants.py,sha256=vDq72d4C9FSWqr9MMdjgTF4TuyNFApvo_6RVsSeAjB4,2837
@@ -366,12 +366,12 @@ truefoundry/workflow/container_task.py,sha256=8arieePsX4__OnG337hOtCiNgJwtKJJCsZ
 truefoundry/workflow/map_task.py,sha256=f9vcAPRQy0Ttw6bvdZBKUVJMSm4eGQrbE1GHWhepHIU,1864
 truefoundry/workflow/python_task.py,sha256=SRXRLC4vdBqGjhkwuaY39LEWN6iPCpJAuW17URRdWTY,1128
 truefoundry/workflow/task.py,sha256=34m55mALXx6ko9o5HkK6FDtMajdvJzBhOsHwDM2RcBA,1779
-truefoundry/workflow/workflow.py,sha256=WaTqUjhwfAXDWu4E5ehuwAxrCbDJkoAf1oWmR2E9Qy0,4575
+truefoundry/workflow/workflow.py,sha256=OjKBwEArxTzNDpfJWgnIqkXDQrYQRLXjheRwpOCu3LE,4861
 truefoundry/workflow/remote_filesystem/__init__.py,sha256=LQ95ViEjJ7Ts4JcCGOxMPs7NZmQdZ4bTiq6qXtsjUhE,206
 truefoundry/workflow/remote_filesystem/logger.py,sha256=em2l7D6sw7xTLDP0kQSLpgfRRCLpN14Qw85TN7ujQcE,1022
 truefoundry/workflow/remote_filesystem/tfy_signed_url_client.py,sha256=xcT0wQmQlgzcj0nP3tJopyFSVWT1uv3nhiTIuwfXYeg,12342
 truefoundry/workflow/remote_filesystem/tfy_signed_url_fs.py,sha256=nSGPZu0Gyd_jz0KsEE-7w_BmnTD8CVF1S8cUJoxaCbc,13305
-truefoundry-0.6.2.dist-info/METADATA,sha256=dooURHmZMoy9716p3g3tFJkt_LwqIqMVPU1vJgvp1VQ,2349
-truefoundry-0.6.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-truefoundry-0.6.2.dist-info/entry_points.txt,sha256=xVjn7RMN-MW2-9f7YU-bBdlZSvvrwzhpX1zmmRmsNPU,98
-truefoundry-0.6.2.dist-info/RECORD,,
+truefoundry-0.6.3.dist-info/METADATA,sha256=i3AnyI6S08k481_Ht6nvQvl4SXeDkV_9Sq8h_cONET4,2349
+truefoundry-0.6.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+truefoundry-0.6.3.dist-info/entry_points.txt,sha256=xVjn7RMN-MW2-9f7YU-bBdlZSvvrwzhpX1zmmRmsNPU,98
+truefoundry-0.6.3.dist-info/RECORD,,