truefoundry 0.5.5rc1__py3-none-any.whl → 0.5.6rc1__py3-none-any.whl

truefoundry/deploy/__init__.py

@@ -9,6 +9,10 @@ from truefoundry.deploy.auto_gen.models import (
     Kustomize,
     ParamType,
     Protocol,
+    SparkDriverConfig,
+    SparkExecutorConfig,
+    SparkExecutorDynamicScaling,
+    SparkExecutorFixedInstances,
     WorkbenchImage,
 )
 from truefoundry.deploy.lib.dao.application import (
@@ -36,11 +40,12 @@ from truefoundry.deploy.v2.lib.deployable_patched_models import (
     Application,
     ApplicationSet,
     AsyncService,
-    Codeserver,
     Helm,
     Job,
     Notebook,
+    RStudio,
     Service,
+    SparkJob,
     SSHServer,
     Volume,
     Workflow,
@@ -110,7 +115,9 @@ from truefoundry.deploy.v2.lib.patched_models import (
     StaticVolumeConfig,
     StringDataMount,
     TPUType,
-    TruefoundryArtifactSource,
+    TrueFoundryArtifactSource,
+    TruefoundryArtifactSource,  # deprecated, kept for backwards compatibility
+    TrueFoundryInteractiveLogin,
     VolumeBrowser,
     VolumeMount,
     WorkerConfig,

truefoundry/deploy/auto_gen/models.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  application.json
-#   timestamp: 2024-12-16T11:59:07+00:00
+#   timestamp: 2025-01-29T10:40:43+00:00
 
 from __future__ import annotations
 
@@ -149,7 +149,7 @@ class BaseAutoscaling(BaseModel):
 
 class BasicAuthCreds(BaseModel):
     """
-    +label=Username and password for service auth
+    +label=Username and password
     """
 
     type: Literal["basic_auth"] = Field(..., description="+value=basic_auth")
@@ -659,6 +659,10 @@ class Profile(str, Enum):
 
 class NvidiaMIGGPU(BaseModel):
     type: Literal["nvidia_mig_gpu"] = Field(..., description="+value=nvidia_mig_gpu")
+    name: Optional[str] = Field(
+        None,
+        description="+label=GPU Name\n+usage=Name of the Nvidia GPU. One of [P4, P100, V100, T4, A10G, A100_40GB, A100_80GB]\nThis field is required for Node Selector and can be ignored in Nodepool Selector.\nOne instance of the card contains the following amount of memory -\nP4: 8 GB, P100: 16 GB, V100: 16 GB, T4: 16 GB, A10G: 24 GB, A100_40GB: 40GB, A100_80GB: 80 GB",
+    )
     profile: Profile = Field(
         ...,
         description="+label=MIG Profile\n+usage=Name of the MIG profile to use. One of [1g.5gb, 2g.10gb, 3g.20gb, 1g.10gb, 2g.20gb, 3g.40gb]",
@@ -669,6 +673,10 @@ class NvidiaTimeslicingGPU(BaseModel):
     type: Literal["nvidia_timeslicing_gpu"] = Field(
         ..., description="+value=nvidia_timeslicing_gpu"
     )
+    name: Optional[str] = Field(
+        None,
+        description="+label=GPU Name\n+usage=Name of the Nvidia GPU. One of [P4, P100, V100, T4, A10G, A100_40GB, A100_80GB]\nThis field is required for Node Selector and can be ignored in Nodepool Selector.\nOne instance of the card contains the following amount of memory -\nP4: 8 GB, P100: 16 GB, V100: 16 GB, T4: 16 GB, A10G: 24 GB, A100_40GB: 40GB, A100_80GB: 80 GB",
+    )
     gpu_memory: conint(ge=1, le=200000) = Field(
         ...,
         description="+label=GPU Memory (MB)\n+usage=Amount of GPU memory (in MB) to allocate. Please note, this limit is not being enforced today but will be in future. Applications are expected to operate in co-opertative mode",
@@ -735,45 +743,6 @@ class AppProtocol(str, Enum):
     tcp = "tcp"
 
 
-class Port(BaseModel):
-    """
-    +docs=Describes the ports the service should be exposed to.
-    """
-
-    port: conint(ge=1, le=65535) = Field(
-        80, description="+usage=Port number to expose."
-    )
-    protocol: Protocol = Field("TCP", description="+usage=Protocol for the port.")
-    expose: bool = Field(True, description="+usage=Expose the port")
-    app_protocol: AppProtocol = Field(
-        "http",
-        description="+label=Application Protocol\n+usage=Application Protocol for the port.\nSelect the application protocol used by your service. For most use cases, this should be `http`(HTTP/1.1).\nIf you are running a gRPC server, select the `grpc` option.\nThis is only applicable if `expose=true`.",
-    )
-    host: Optional[
-        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]?)))$"
-        )
-    ] = Field(
-        None,
-        description="+usage=Host e.g. ai.example.com, app.truefoundry.com\n+message=Upto 253 characters, each part of host should be at most 63 characters long, can contain alphabets, digits and hypen, must begin and end with an alphanumeric characters. Parts must be separated by periods (.)",
-    )
-    path: Optional[
-        constr(regex=r"^(/([a-zA-Z0-9]|[a-zA-Z0-9][a-zA-Z0-9\-_\.]*[a-zA-Z0-9]))*/$")
-    ] = Field(
-        None,
-        description="+usage=Path e.g. /v1/api/ml/, /v2/docs/\n+message=Should begin and end with a forward slash (/). Each part can can contain alphabets, digits and hypen, must begin and end with an alphanumeric characters. Parts should be separated by forward slashes (/)",
-    )
-    rewrite_path_to: Optional[
-        constr(regex=r"^(/([a-zA-Z0-9]|[a-zA-Z0-9][a-zA-Z0-9\-_\.]*[a-zA-Z0-9]))*/$")
-    ] = Field(
-        None,
-        description="+label=Rewrite Path to\n+usage=Rewrite the path prefix to a different path.\nIf `path` is `/v1/api` and `rewrite_path_to` is `/api`. The URI in the HTTP request `http://0.0.0.0:8080/v1/api/houses` will be rewritten to `http://0.0.0.0:8080/api/houses` before the request is forwarded your service.\nDefaults to `/`.\nThis is only applicable if `path` is given.\n+message=Should begin and end with a forward slash (/). Each part can can contain alphabets, digits and hypen, must begin and end with an alphanumeric characters. Parts should be separated by forward slashes (/)",
-    )
-    auth: Optional[Union[BasicAuthCreds, JwtAuthCreds]] = Field(
-        None, description="+usage=Username and Password for service auth"
-    )
-
-
 class PythonBuild(BaseModel):
     """
     +docs=Describes that we are using python to build a container image with a specific python version and pip packages installed.
@@ -1010,7 +979,7 @@ class SecretMount(BaseModel):
     )
     secret_fqn: constr(regex=r"^tfy-secret:\/\/.+:.+:.+$") = Field(
         ...,
-        description="+label=Secret\n+usage=The Truefoundry secret whose value will be the file content.",
+        description="+label=Secret\n+usage=The TrueFoundry secret whose value will be the file content.",
     )
 
 
@@ -1021,6 +990,42 @@ class ServiceAutoscaling(BaseAutoscaling):
     )
 
 
+class SparkDriverConfig(BaseModel):
+    """
+    +label=Driver Config
+    """
+
+    ui_endpoint: Endpoint
+    resources: Optional[Resources] = None
+
+
+class SparkExecutorDynamicScaling(BaseModel):
+    """
+    +label=Dynamic Scaling
+    """
+
+    type: Literal["dynamic"] = Field(..., description="+value=dynamic")
+    min: conint(ge=0, le=500) = Field(
+        1,
+        description="+label=Min Instances\n+usage=Minimum number of instances to start / scale down to\n+sort=100",
+    )
+    max: conint(ge=0, le=500) = Field(
+        1,
+        description="+label=Max Instances\n+usage=Maximum number of instances to scale up to\n+sort=200",
+    )
+
+
+class SparkExecutorFixedInstances(BaseModel):
+    """
+    +label=Fixed Instances
+    """
+
+    type: Literal["fixed"] = Field(..., description="+value=fixed")
+    count: conint(ge=0, le=500) = Field(
+        1, description="+label=Instances Count\n+usage=Number of instances to start"
+    )
+
+
 class StaticVolumeConfig(BaseModel):
     """
     +label=Static Volume Config
@@ -1106,10 +1111,10 @@ class TaskPythonBuild(BaseModel):
     )
 
 
-class TruefoundryArtifactSource(BaseModel):
+class TrueFoundryArtifactSource(BaseModel):
     """
-    +docs=Input for Artifact from Truefoundry Artifact Registry
-    +label=Truefoundry Artifact Source
+    +docs=Input for Artifact from TrueFoundry Artifact Registry
+    +label=TrueFoundry Artifact Source
     """
 
     type: Literal["truefoundry-artifact"] = Field(
@@ -1125,6 +1130,16 @@ class TruefoundryArtifactSource(BaseModel):
     )
 
 
+class TrueFoundryInteractiveLogin(BaseModel):
+    """
+    +label=Login with truefoundry
+    """
+
+    type: Literal["truefoundry_oauth"] = Field(
+        ..., description="+value=truefoundry_oauth"
+    )
+
+
 class VolumeBrowser(BaseModel):
     """
     +label=Volume Browser
@@ -1157,7 +1172,7 @@ class VolumeMount(BaseModel):
     )
     volume_fqn: constr(regex=r"^tfy-volume:\/\/.+:.+:.+$") = Field(
         ...,
-        description="+label=Volume\n+usage=The Truefoundry volume that needs to be mounted.",
+        description="+label=Volume\n+usage=The TrueFoundry volume that needs to be mounted.",
     )
 
 
@@ -1188,7 +1203,7 @@ class ArtifactsDownload(BaseModel):
     """
 
     cache_volume: Optional[ArtifactsCacheVolume] = None
-    artifacts: List[Union[TruefoundryArtifactSource, HuggingfaceArtifactSource]] = (
+    artifacts: List[Union[TrueFoundryArtifactSource, HuggingfaceArtifactSource]] = (
         Field(
             ..., description="+label=Artifacts\n+usage=List of artifacts to be cached"
         )
@@ -1213,7 +1228,7 @@ class BaseWorkbenchInput(BaseModel):
     +docs=Describes the configuration for the service
     """
 
-    name: constr(regex=r"^[a-z][a-z0-9\-]{1,30}[a-z0-9]$") = Field(
+    name: constr(regex=r"^[a-z](?:[a-z0-9]|-(?!-)){1,30}[a-z0-9]$") = Field(
         ...,
         description="+usage=Name of the workbench. This uniquely identifies this workbench in the workspace.\n> Name can only contain alphanumeric characters and '-' and can be atmost 25 characters long\n+sort=1\n+message=3 to 32 lower case characters long alphanumeric word, may contain - in between, cannot start with a number",
     )
@@ -1279,9 +1294,8 @@ class Codeserver(BaseWorkbenchInput):
     +docs=Describes the configuration for the code server
     """
 
-    type: Literal["codeserver"] = Field(..., description="+value=Code Server")
+    type: Literal["codeserver"] = Field(..., description="+value=codeserver")
     image: WorkbenchImage
-    auth: Optional[BasicAuthCreds] = None
 
 
 class ContainerTaskConfig(BaseModel):
@@ -1364,7 +1378,7 @@ class HealthProbe(BaseModel):
 
 class Helm(BaseModel):
     type: Literal["helm"] = Field(..., description="+value=helm")
-    name: constr(regex=r"^[a-z][a-z0-9\-]{1,30}[a-z0-9]$") = Field(
+    name: constr(regex=r"^[a-z](?:[a-z0-9]|-(?!-)){1,30}[a-z0-9]$") = Field(
         ...,
         description="+sort=1\n+message=3 to 32 lower case characters long alphanumeric word, may contain - in between, cannot start with a number\n+usage=Name of the Helm deployment. This will be set as the release name of the chart you are deploying.",
     )
@@ -1391,7 +1405,7 @@ class Job(BaseModel):
     """
 
     type: Literal["job"] = Field(..., description="+value=job")
-    name: constr(regex=r"^[a-z][a-z0-9\-]{1,30}[a-z0-9]$") = Field(
+    name: constr(regex=r"^[a-z](?:[a-z0-9]|-(?!-)){1,30}[a-z0-9]$") = Field(
         ...,
         description="+usage=Name of the job\n+sort=1\n+message=3 to 32 lower case characters long alphanumeric word, may contain - in between, cannot start with a number",
     )
@@ -1549,13 +1563,51 @@ class Notebook(BaseWorkbenchInput):
 
     type: Literal["notebook"] = Field(..., description="+value=notebook")
     image: WorkbenchImage
-    auth: Optional[BasicAuthCreds] = None
     cull_timeout: conint(ge=5) = Field(
         30,
         description="+label=Stop after (minutes of inactivity)\n+usage=Stop the notebook instance after this much time in minutes of inactivity.\nThe notebook instance will be stopped even if the notebook is open in your browser, but nothing is running on the notebook.\n+sort=5",
     )
 
 
+class Port(BaseModel):
+    """
+    +docs=Describes the ports the service should be exposed to.
+    """
+
+    port: conint(ge=1, le=65535) = Field(
+        80, description="+usage=Port number to expose."
+    )
+    protocol: Protocol = Field("TCP", description="+usage=Protocol for the port.")
+    expose: bool = Field(True, description="+usage=Expose the port")
+    app_protocol: AppProtocol = Field(
+        "http",
+        description="+label=Application Protocol\n+usage=Application Protocol for the port.\nSelect the application protocol used by your service. For most use cases, this should be `http`(HTTP/1.1).\nIf you are running a gRPC server, select the `grpc` option.\nThis is only applicable if `expose=true`.",
+    )
+    host: Optional[
+        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]?)))$"
+        )
+    ] = Field(
+        None,
+        description="+usage=Host e.g. ai.example.com, app.truefoundry.com\n+message=Upto 253 characters, each part of host should be at most 63 characters long, can contain alphabets, digits and hypen, must begin and end with an alphanumeric characters. Parts must be separated by periods (.)",
+    )
+    path: Optional[
+        constr(regex=r"^(/([a-zA-Z0-9]|[a-zA-Z0-9][a-zA-Z0-9\-_\.]*[a-zA-Z0-9]))*/$")
+    ] = Field(
+        None,
+        description="+usage=Path e.g. /v1/api/ml/, /v2/docs/\n+message=Should begin and end with a forward slash (/). Each part can can contain alphabets, digits and hypen, must begin and end with an alphanumeric characters. Parts should be separated by forward slashes (/)",
+    )
+    rewrite_path_to: Optional[
+        constr(regex=r"^(/([a-zA-Z0-9]|[a-zA-Z0-9][a-zA-Z0-9\-_\.]*[a-zA-Z0-9]))*/$")
+    ] = Field(
+        None,
+        description="+label=Rewrite Path to\n+usage=Rewrite the path prefix to a different path.\nIf `path` is `/v1/api` and `rewrite_path_to` is `/api`. The URI in the HTTP request `http://0.0.0.0:8080/v1/api/houses` will be rewritten to `http://0.0.0.0:8080/api/houses` before the request is forwarded your service.\nDefaults to `/`.\nThis is only applicable if `path` is given.\n+message=Should begin and end with a forward slash (/). Each part can can contain alphabets, digits and hypen, must begin and end with an alphanumeric characters. Parts should be separated by forward slashes (/)",
+    )
+    auth: Optional[Union[BasicAuthCreds, JwtAuthCreds, TrueFoundryInteractiveLogin]] = (
+        Field(None, description="+usage=Authentication method for inbound traffic")
+    )
+
+
 class PythonTaskConfig(BaseModel):
     """
     +docs=Describes the configuration for the python function task
@@ -1582,22 +1634,88 @@ class PythonTaskConfig(BaseModel):
     )
 
 
+class RStudio(BaseWorkbenchInput):
+    """
+    +docs=Describes the configuration for the Rstudio server
+    """
+
+    type: Literal["rstudio"] = Field(..., description="+value=rstudio")
+    image: WorkbenchImage
+
+
 class SSHServer(BaseWorkbenchInput):
     """
     +docs=Describes the configuration for the ssh server
     """
 
-    type: Literal["ssh-server"] = Field(..., description="+value=SSH Server")
+    type: Literal["ssh-server"] = Field(..., description="+value=ssh-server")
     image: WorkbenchImage
     ssh_public_key: str = Field(
         ...,
         description="+label: SSH Public Key\n+usage=Add Your SSH Public Key, this will be used to authenticate you to the SSH Server.  \\\nYou can find it using `cat ~/.ssh/id_rsa.pub` in Mac/Linux or `type $home\\.ssh\\id_rsa.pub` in Windows Powershell.  \\\nYou can also generate a new SSH key pair using `ssh-keygen -t rsa` in your local terminal. (same for both Mac/Linux and Windows Powershell)\n+uiType=TextArea\n+sort=4",
     )
+    cull_timeout: Optional[conint(ge=5)] = Field(
+        None,
+        description="+label=Stop after (minutes of inactivity)\n+usage=Stop the SSH Server instance after this much time in minutes of inactivity. \\\nThe instance is considered active if there is at least one active SSH connection (a client connected to the SSH server), \\\nor if a background job is running using tmux or screen, or if the pod has restarted.\n+sort=5",
+    )
+
+
+class SparkExecutorConfig(BaseModel):
+    """
+    +label=Executor Config
+    """
+
+    instances: Union[SparkExecutorFixedInstances, SparkExecutorDynamicScaling] = Field(
+        {"type": "fixed", "count": 1}, description="+label=Executor Instances"
+    )
+    resources: Optional[Resources] = None
+
+
+class SparkJob(BaseModel):
+    type: Literal["spark-job"] = Field(..., description="+value=spark-job\n+sort=1")
+    name: constr(regex=r"^[a-z](?:[a-z0-9]|-(?!-)){1,30}[a-z0-9]$") = Field(
+        ...,
+        description="+label=Name\n+usage=Name of the job\n+message=3 to 32 lower case characters long alphanumeric word, may contain - in between, cannot start with a number\n+sort=2",
+    )
+    image: Image
+    spark_version: str = Field(
+        "3.5.2",
+        description="+label=Spark Version\n+usage=Spark version should match the spark version installed in the image.\n+sort=2000",
+    )
+    main_application_file: str = Field(
+        ...,
+        description="+label=Main Application File\n+usage=The main application file to be executed by the spark job.\n+sort=3000",
+    )
+    arguments: Optional[str] = Field(
+        None,
+        description="+label=Arguments\n+usage=Arguments to be passed to the main application file.\n+sort=4000",
+    )
+    driver_config: SparkDriverConfig
+    executor_config: SparkExecutorConfig
+    env: Optional[Dict[str, Any]] = Field(
+        None,
+        description="+label=Environment Variables\n+usage=Configure environment variables to be injected in the service either as plain text. [Docs](https://docs.truefoundry.com/docs/env-variables)\n+icon=fa-globe\n+sort=21000",
+    )
+    mounts: Optional[List[VolumeMount]] = Field(
+        None,
+        description="+label=Mounts\n+usage=Configure volumes to be mounted to driver and executors. [Docs](https://docs.truefoundry.com/docs/mounting-volumes-job)\n+sort=22000\n+uiType=Mounts",
+    )
+    retries: conint(ge=0, le=10) = Field(
+        0,
+        description="+label=Retries\n+usage=Specify the maximum number of attempts to retry a job before it is marked as failed.\n+icon=fa-repeat\n+sort=23000",
+    )
+    service_account: Optional[str] = Field(
+        None, description="+label=Service Account\n+sort=24000"
+    )
+    workspace_fqn: Optional[str] = Field(
+        None,
+        description="+label=Workspace FQN\n+docs=Fully qualified name of the workspace\n+uiType=Hidden",
+    )
 
 
 class Volume(BaseModel):
     type: Literal["volume"] = Field(..., description="+value=volume")
-    name: constr(regex=r"^[a-z][a-z0-9\-]{1,30}[a-z0-9]$") = Field(
+    name: constr(regex=r"^[a-z](?:[a-z0-9]|-(?!-)){1,30}[a-z0-9]$") = Field(
         ...,
         description="+sort=1\n+message=3 to 32 lower case characters long alphanumeric word, may contain - in between, cannot start with a number\n+usage=Name of the Volume. This will be set as the volume name.",
     )
@@ -1632,7 +1750,7 @@ class WorkerConfig(BaseModel):
 
 
 class BaseService(BaseModel):
-    name: constr(regex=r"^[a-z][a-z0-9\-]{1,30}[a-z0-9]$") = Field(
+    name: constr(regex=r"^[a-z](?:[a-z0-9]|-(?!-)){1,30}[a-z0-9]$") = Field(
         ...,
         description="+usage=Name of the service. This uniquely identifies this service in the workspace.\n> Name can only contain alphanumeric characters and '-' and can be atmost 25 characters long\n+sort=1\n+message=3 to 32 lower case characters long alphanumeric word, may contain - in between, cannot start with a number",
     )
@@ -1727,7 +1845,7 @@ class Workflow(BaseModel):
     """
 
     type: Literal["workflow"] = Field(..., description="+value=workflow")
-    name: constr(regex=r"^[a-z][a-z0-9\-]{1,30}[a-z0-9]$") = Field(
+    name: constr(regex=r"^[a-z](?:[a-z0-9]|-(?!-)){1,30}[a-z0-9]$") = Field(
         ...,
         description="+usage=Name of the workflow\n+sort=1\n+message=3 to 32 lower case characters long alphanumeric word, may contain - in between, cannot start with a number",
     )
@@ -1780,8 +1898,10 @@ class Application(BaseModel):
         Notebook,
         Codeserver,
         SSHServer,
+        RStudio,
         Helm,
         Volume,
         ApplicationSet,
         Workflow,
+        SparkJob,
     ]

truefoundry/deploy/cli/commands/__init__.py

@@ -1,6 +1,5 @@
 from truefoundry.deploy.cli.commands.apply_command import get_apply_command
 from truefoundry.deploy.cli.commands.build_command import get_build_command
-from truefoundry.deploy.cli.commands.build_logs_command import get_build_logs_command
 from truefoundry.deploy.cli.commands.create_command import get_create_command
 from truefoundry.deploy.cli.commands.delete_command import get_delete_command
 from truefoundry.deploy.cli.commands.deploy_command import get_deploy_command

truefoundry/deploy/lib/clients/servicefoundry_client.py

@@ -273,20 +273,33 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         subscribe_message: str,
         socketio_path: str = "socket.io",
         callback=None,
-        wait=True,
-    ):
+    ) -> socketio.Client:
         callback = callback or OutputCallBack()
         sio = socketio.Client(request_timeout=60)
         callback.print_line("Waiting for the task to start...")
+        next_log_start_timestamp = query_dict.get("startTs")
 
         @sio.on(subscribe_message)
         def logs(data):
             try:
                 _log = json.loads(data)
-                callback.print_line(self._get_log_print_line(_log["body"]))
+                body = _log["body"]
+                callback.print_line(self._get_log_print_line(body))
+                nonlocal next_log_start_timestamp
+                next_log_start_timestamp = body["time"]
             except Exception:
                 logger.exception(f"Error while parsing log line, {data!r}")
 
+        @sio.on("connect")
+        def on_connect():
+            # TODO: We should have have a timeout here. `emit` does
+            #   not support timeout. Explore `sio.call`.
+            query_dict["startTs"] = next_log_start_timestamp
+            sio.emit(
+                subscribe_message,
+                json.dumps(query_dict),
+            )
+
         def sio_disconnect_no_exception():
             try:
                 sio.disconnect()
@@ -300,14 +313,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
                 headers=self._get_header(),
                 socketio_path=socketio_path,
             )
-            # TODO: We should have have a timeout here. `emit` does
-            #   not support timeout. Explore `sio.call`.
-            sio.emit(
-                subscribe_message,
-                json.dumps(query_dict),
-            )
-            if wait:
-                sio.wait()
+            return sio
 
     @check_min_cli_version
     def get_deployment(self, application_id: str, deployment_id: str) -> Deployment:
@@ -372,11 +378,10 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         self,
         build_response: BuildResponse,
         callback=None,
-        wait: bool = True,
-    ):
+    ) -> socketio.Client:
         callback = callback or OutputCallBack()
         tail_logs_obj = json.loads(build_response.tailLogsUrl)
-        self._tail_logs(
+        socket = self._tail_logs(
             tail_logs_url=urljoin(
                 tail_logs_obj["uri"], f"/?type={BUILD_LOGS_SUBSCRIBE_MESSAGE}"
             ),
@@ -386,9 +391,9 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
                 "startTs": build_response.logsStartTs,
             },
             callback=callback,
-            wait=wait,
             subscribe_message=BUILD_LOGS_SUBSCRIBE_MESSAGE,
         )
+        return socket
 
     @check_min_cli_version
     def tail_logs_for_deployment(
@@ -416,7 +421,6 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
                 },
             },
             callback=callback,
-            wait=wait,
             subscribe_message=DEPLOYMENT_LOGS_SUBSCRIBE_MESSAGE,
         )
 

truefoundry/deploy/v2/lib/deploy.py

@@ -1,7 +1,8 @@
 import sys
 import time
-from typing import List, Optional, TypeVar
+from typing import Optional, TypeVar
 
+import socketio
 from rich.status import Status
 
 from truefoundry.common.utils import poll_for_function
@@ -96,18 +97,15 @@ def _log_application_dashboard_url(deployment: Deployment, log_message: str):
     logger.info(log_message, url)
 
 
-def _tail_build_logs(build_responses: List[BuildResponse]) -> None:
+def _tail_build_logs(build_response: BuildResponse) -> socketio.Client:
     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)
+    logger.info("Tailing build logs for '%s'", build_response.componentName)
+    socket = client.tail_build_logs(build_response=build_response)
+    return socket
 
 
-def _deploy_wait_handler(
+def _deploy_wait_handler(  # noqa: C901
     deployment: Deployment,
 ) -> Optional[DeploymentTransitionStatus]:
     _log_application_dashboard_url(
@@ -122,9 +120,10 @@ def _deploy_wait_handler(
         last_status_printed = None
         client = ServiceFoundryServiceClient()
         start_time = time.monotonic()
-        total_timeout_time: int = 300
+        total_timeout_time: int = 600
         poll_interval_seconds = 5
         time_elapsed = 0
+        socket: socketio.Client = None
 
         for deployment_statuses in poll_for_function(
             client.get_deployment_statuses,
@@ -150,20 +149,29 @@ def _deploy_wait_handler(
                 last_status_printed = status_to_print
 
             if latest_deployment_status.state.isTerminalState:
+                if socket and socket.connected:
+                    socket.disconnect()
                 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)
+                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])
 
             time_elapsed = time.monotonic() - start_time
             if time_elapsed > total_timeout_time:
-                logger.warning("Polled server for %s secs.", int(time_elapsed))
+                logger.warning(
+                    "Polled server for %s secs. Disconnecting from server, the deployment will still continue.",
+                    int(time_elapsed),
+                )
+                if socket and socket.connected:
+                    socket.disconnect()
                 break
 
     return last_status_printed

truefoundry/deploy/v2/lib/deployable_patched_models.py

@@ -30,6 +30,10 @@ class Job(models.Job, DeployablePatchedModelBase):
     resources: models.Resources = Field(default_factory=models.Resources)
 
 
+class SparkJob(models.SparkJob, DeployablePatchedModelBase):
+    type: Literal["spark-job"] = "spark-job"
+
+
 class Notebook(models.Notebook, DeployablePatchedModelBase):
     type: Literal["notebook"] = "notebook"
     resources: models.Resources = Field(default_factory=models.Resources)
@@ -37,6 +41,12 @@ class Notebook(models.Notebook, DeployablePatchedModelBase):
 
 class Codeserver(models.Codeserver, DeployablePatchedModelBase):
     type: Literal["codeserver"] = "codeserver"
+    resources: models.Resources = Field(default_factory=models.Resources)
+
+
+class RStudio(models.RStudio, DeployablePatchedModelBase):
+    type: Literal["rstudio"] = "rstudio"
+    resources: models.Resources = Field(default_factory=models.Resources)
 
 
 class Helm(models.Helm, DeployablePatchedModelBase):

truefoundry/deploy/v2/lib/patched_models.py

@@ -256,6 +256,10 @@ class JwtAuthCreds(models.JwtAuthCreds, PatchedModelBase):
     type: Literal["jwt_auth"] = "jwt_auth"
 
 
+class TrueFoundryInteractiveLogin(models.TrueFoundryInteractiveLogin, PatchedModelBase):
+    type: Literal["truefoundry_oauth"] = "truefoundry_oauth"
+
+
 class HealthProbe(models.HealthProbe, PatchedModelBase):
     pass
 
@@ -500,10 +504,14 @@ class HuggingfaceArtifactSource(models.HuggingfaceArtifactSource, PatchedModelBa
     type: Literal["huggingface-hub"] = "huggingface-hub"
 
 
-class TruefoundryArtifactSource(models.TruefoundryArtifactSource, PatchedModelBase):
+class TrueFoundryArtifactSource(models.TrueFoundryArtifactSource, PatchedModelBase):
     type: Literal["truefoundry-artifact"] = "truefoundry-artifact"
 
 
+# Deprecated alias, kept for backward compatibility
+TruefoundryArtifactSource = TrueFoundryArtifactSource
+
+
 class ArtifactsDownload(models.ArtifactsDownload, PatchedModelBase):
     pass
 

truefoundry/ml/autogen/client/__init__.py

@@ -175,8 +175,8 @@ from truefoundry.ml.autogen.client.models.experiment_tag_dto import ExperimentTa
 from truefoundry.ml.autogen.client.models.export_deployment_files_request_dto import (
     ExportDeploymentFilesRequestDto,
 )
-from truefoundry.ml.autogen.client.models.external_artifact_source import (
-    ExternalArtifactSource,
+from truefoundry.ml.autogen.client.models.external_blob_storage_source import (
+    ExternalBlobStorageSource,
 )
 from truefoundry.ml.autogen.client.models.fast_ai_framework import FastAIFramework
 from truefoundry.ml.autogen.client.models.file_info_dto import FileInfoDto
@@ -403,8 +403,8 @@ from truefoundry.ml.autogen.client.models.trigger_job_run_config_request_dto imp
 from truefoundry.ml.autogen.client.models.trigger_job_run_config_response_dto import (
     TriggerJobRunConfigResponseDto,
 )
-from truefoundry.ml.autogen.client.models.true_foundry_artifact_source import (
-    TrueFoundryArtifactSource,
+from truefoundry.ml.autogen.client.models.true_foundry_managed_source import (
+    TrueFoundryManagedSource,
 )
 from truefoundry.ml.autogen.client.models.update_artifact_version_request_dto import (
     UpdateArtifactVersionRequestDto,