truefoundry 0.5.10rc2__py3-none-any.whl → 0.5.11rc1__py3-none-any.whl

truefoundry/deploy/builder/__init__.py

@@ -1,6 +1,6 @@
 from typing import Any, Dict, List, Optional, Union
 
-from truefoundry.deploy.auto_gen.models import (
+from truefoundry.deploy._autogen.models import (
     DockerFileBuild,
     PythonBuild,
     TaskDockerFileBuild,

truefoundry/deploy/builder/builders/dockerfile.py

@@ -1,7 +1,7 @@
 import os
 from typing import Dict, List, Optional
 
-from truefoundry.deploy.auto_gen.models import DockerFileBuild
+from truefoundry.deploy._autogen.models import DockerFileBuild
 from truefoundry.deploy.builder.docker_service import build_docker_image
 from truefoundry.logger import logger
 

truefoundry/deploy/builder/builders/tfy_notebook_buildpack/__init__.py

@@ -2,7 +2,7 @@ import os
 from tempfile import TemporaryDirectory
 from typing import List, Optional
 
-from truefoundry.deploy.auto_gen.models import DockerFileBuild
+from truefoundry.deploy._autogen.models import DockerFileBuild
 from truefoundry.deploy.builder.builders import dockerfile
 from truefoundry.deploy.builder.builders.tfy_notebook_buildpack.dockerfile_template import (
     NotebookImageBuild,

truefoundry/deploy/builder/builders/tfy_python_buildpack/__init__.py

@@ -2,7 +2,7 @@ import os
 from tempfile import TemporaryDirectory
 from typing import List, Optional
 
-from truefoundry.deploy.auto_gen.models import DockerFileBuild, PythonBuild
+from truefoundry.deploy._autogen.models import DockerFileBuild, PythonBuild
 from truefoundry.deploy.builder.builders import dockerfile
 from truefoundry.deploy.builder.builders.tfy_python_buildpack.dockerfile_template import (
     generate_dockerfile_content,

truefoundry/deploy/builder/builders/tfy_python_buildpack/dockerfile_template.py

@@ -4,7 +4,7 @@ from typing import Dict, List, Optional
 from mako.template import Template
 
 from truefoundry.common.constants import ENV_VARS, PythonPackageManager
-from truefoundry.deploy.auto_gen.models import PythonBuild
+from truefoundry.deploy._autogen.models import PythonBuild
 from truefoundry.deploy.builder.constants import (
     PIP_CONF_BUILDKIT_SECRET_MOUNT,
     PIP_CONF_SECRET_MOUNT_AS_ENV,

truefoundry/deploy/lib/clients/servicefoundry_client.py

@@ -24,7 +24,8 @@ from truefoundry.common.servicefoundry_client import (
     session_with_retries,
 )
 from truefoundry.common.session import Session
-from truefoundry.deploy.auto_gen import models as auto_gen_models
+from truefoundry.common.utils import get_user_agent
+from truefoundry.deploy._autogen import models as autogen_models
 from truefoundry.deploy.io.output_callback import OutputCallBack
 from truefoundry.deploy.lib.model.entity import (
     Application,
@@ -89,16 +90,20 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
             raise Exception("Neither session, not tfy_host provided")
         super().__init__(tfy_host=tfy_host)
 
-    def _get_header(self) -> Dict[str, str]:
+    def _get_headers(self) -> Dict[str, str]:
+        headers = {"User-Agent": get_user_agent()}
         if not self._session:
-            return {}
-        return {"Authorization": f"Bearer {self._session.access_token}"}
+            return headers
+        return {
+            **headers,
+            "Authorization": f"Bearer {self._session.access_token}",
+        }
 
     @check_min_cli_version
     def get_id_from_fqn(self, fqn_type: str, fqn: str) -> Dict[str, Any]:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/fqn/{fqn_type}"
         response = session_with_retries().get(
-            url, headers=self._get_header(), params={"fqn": fqn}
+            url, headers=self._get_headers(), params={"fqn": fqn}
         )
         return request_handling(response)
 
@@ -118,7 +123,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         if workspace_fqn:
             params["workspaceFqn"] = workspace_fqn
         response = session_with_retries().get(
-            url, params=params, headers=self._get_header()
+            url, params=params, headers=self._get_headers()
         )
         response = request_handling(response)
         return parse_obj_as(List[Workspace], response)
@@ -140,7 +145,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
                     "resources": resources.dict(exclude_none=True),
                 }
             },
-            headers=self._get_header(),
+            headers=self._get_headers(),
         )
         response_data = request_handling(response)
         return Workspace.parse_obj(response_data)
@@ -152,7 +157,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
             force
         )  # this dumb conversion is required because `params` just casts as str
         response = session_with_retries().delete(
-            url, headers=self._get_header(), params={"force": force_str}
+            url, headers=self._get_headers(), params={"force": force_str}
         )
         response = cast(Dict[str, Any], request_handling(response))
         return Workspace.parse_obj(response["workspace"])
@@ -162,7 +167,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         url = f"{self._api_server_url}/{VERSION_PREFIX}/workspace"
         response = session_with_retries().get(
             url,
-            headers=self._get_header(),
+            headers=self._get_headers(),
             params={"fqn": workspace_fqn},
         )
         response = request_handling(response)
@@ -171,7 +176,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
     @check_min_cli_version
     def get_cluster(self, cluster_id: str) -> Dict[str, Any]:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/cluster/{cluster_id}"
-        response = session_with_retries().get(url, headers=self._get_header())
+        response = session_with_retries().get(url, headers=self._get_headers())
         return cast(Dict[str, Any], request_handling(response))
 
     @check_min_cli_version
@@ -186,7 +191,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
                 "serviceName": service_name,
                 "stage": env,
             },
-            headers=self._get_header(),
+            headers=self._get_headers(),
         )
         return cast(Dict[str, Any], request_handling(response))
 
@@ -204,7 +209,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
     def deploy_application(
         self,
         workspace_id: str,
-        application: auto_gen_models.Workflow,
+        application: autogen_models.Workflow,
         force: bool = False,
     ) -> Deployment:
         data = {
@@ -216,7 +221,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         logger.debug(json.dumps(data))
         url = f"{self._api_server_url}/{VERSION_PREFIX}/deployment"
         response = session_with_retries().post(
-            url, json=data, headers=self._get_header()
+            url, json=data, headers=self._get_headers()
         )
         response_data = cast(Dict[str, Any], request_handling(response))
         return Deployment.parse_obj(response_data["deployment"])
@@ -279,7 +284,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
             sio.connect(
                 tail_logs_url,
                 transports="websocket",
-                headers=self._get_header(),
+                headers=self._get_headers(),
                 socketio_path=socketio_path,
             )
             return sio
@@ -287,7 +292,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
     @check_min_cli_version
     def get_deployment(self, application_id: str, deployment_id: str) -> Deployment:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/app/{application_id}/deployments/{deployment_id}"
-        response = session_with_retries().get(url, headers=self._get_header())
+        response = session_with_retries().get(url, headers=self._get_headers())
         response_data = request_handling(response)
         return Deployment.parse_obj(response_data)
 
@@ -296,7 +301,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         self, application_id: str, deployment_id: str
     ) -> 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_header())
+        response = session_with_retries().get(url, headers=self._get_headers())
         response_data = request_handling(response)
         return parse_obj_as(List[AppDeploymentStatusResponse], response_data)
 
@@ -305,7 +310,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         self, application_id: str, deployment_id: str
     ) -> List[BuildResponse]:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/app/{application_id}/deployments/{deployment_id}/builds"
-        response = session_with_retries().get(url, headers=self._get_header())
+        response = session_with_retries().get(url, headers=self._get_headers())
         response_data = request_handling(response)
         return parse_obj_as(List[BuildResponse], response_data)
 
@@ -337,7 +342,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
 
         url = f"{self._api_server_url}/{VERSION_PREFIX}/logs/{workspace_id}"
         response = session_with_retries().get(
-            url=url, params=data, headers=self._get_header()
+            url=url, params=data, headers=self._get_headers()
         )
         response_data = cast(Dict[str, Any], request_handling(response))
         return parse_obj_as(List[LogBody], response_data["data"])
@@ -352,7 +357,8 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         tail_logs_obj = json.loads(build_response.tailLogsUrl)
         socket = self._tail_logs(
             tail_logs_url=urljoin(
-                tail_logs_obj["uri"], f"/?type={BUILD_LOGS_SUBSCRIBE_MESSAGE}"
+                tail_logs_obj["uri"],
+                f"/?type={BUILD_LOGS_SUBSCRIBE_MESSAGE}",
             ),
             socketio_path=tail_logs_obj["path"],
             query_dict={
@@ -474,7 +480,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
     ) -> None:
         callback = callback or OutputCallBack()
         url = build_response.getLogsUrl
-        response = session_with_retries().get(url=url, headers=self._get_header())
+        response = session_with_retries().get(url=url, headers=self._get_headers())
         logs_response = cast(Dict[str, Any], request_handling(response))
         for _log_body in logs_response["logs"]:
             try:
@@ -487,7 +493,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
     def get_deployment_info_by_fqn(self, deployment_fqn: str) -> DeploymentFqnResponse:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/fqn/deployment"
         response = session_with_retries().get(
-            url, headers=self._get_header(), params={"fqn": deployment_fqn}
+            url, headers=self._get_headers(), params={"fqn": deployment_fqn}
         )
         response_data = request_handling(response)
         return DeploymentFqnResponse.parse_obj(response_data)
@@ -498,7 +504,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
     ) -> ApplicationFqnResponse:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/fqn/app"
         response = session_with_retries().get(
-            url, headers=self._get_header(), params={"fqn": application_fqn}
+            url, headers=self._get_headers(), params={"fqn": application_fqn}
         )
         response_data = request_handling(response)
         return ApplicationFqnResponse.parse_obj(response_data)
@@ -506,7 +512,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
     @check_min_cli_version
     def remove_application(self, application_id: str) -> Dict[str, Any]:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/app/{application_id}"
-        response = session_with_retries().delete(url, headers=self._get_header())
+        response = session_with_retries().delete(url, headers=self._get_headers())
         response = cast(Dict[str, Any], request_handling(response))
         # TODO: Add pydantic here.
         return response
@@ -514,7 +520,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
     @check_min_cli_version
     def get_application_info(self, application_id: str) -> Application:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/app/{application_id}"
-        response = session_with_retries().get(url, headers=self._get_header())
+        response = session_with_retries().get(url, headers=self._get_headers())
         response = request_handling(response)
         return Application.parse_obj(response)
 
@@ -534,7 +540,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         if search_prefix:
             params["searchPrefix"] = search_prefix
         response = session_with_retries().get(
-            url, headers=self._get_header(), params=params
+            url, headers=self._get_headers(), params=params
         )
         response_data = cast(Dict[str, Any], request_handling(response))
         return parse_obj_as(List[JobRun], response_data["data"])
@@ -545,7 +551,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         job_run_name: str,
     ) -> JobRun:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/jobs/{application_id}/runs/{job_run_name}"
-        response = session_with_retries().get(url, headers=self._get_header())
+        response = session_with_retries().get(url, headers=self._get_headers())
         response_data = request_handling(response)
         return parse_obj_as(JobRun, response_data)
 
@@ -565,7 +571,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         if params:
             body["input"]["params"] = params
         response = session_with_retries().post(
-            url, json=body, headers=self._get_header()
+            url, json=body, headers=self._get_headers()
         )
         response = request_handling(response)
         return TriggerJobResult.parse_obj(response)
@@ -576,7 +582,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         url = f"{self._api_server_url}/{VERSION_PREFIX}/workflow/{application_id}/executions"
         body = {"inputs": inputs}
         response = session_with_retries(retries=MAX_RETRIES_WORKFLOW_TRIGGER).post(
-            url, json=body, headers=self._get_header()
+            url, json=body, headers=self._get_headers()
         )
         response = cast(Dict[str, Any], request_handling(response))
         return response
@@ -588,7 +594,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         url = f"{self._api_server_url}/{VERSION_PREFIX}/docker-registry/creds"
         response = session_with_retries().get(
             url,
-            headers=self._get_header(),
+            headers=self._get_headers(),
             params={
                 "fqn": docker_registry_fqn,
                 "clusterId": cluster_id,
@@ -604,7 +610,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         url = f"{self._api_server_url}/{VERSION_PREFIX}/docker-registry/create-repo"
         response = session_with_retries().post(
             url,
-            headers=self._get_header(),
+            headers=self._get_headers(),
             data={
                 "fqn": docker_registry_fqn,
                 "workspaceFqn": workspace_fqn,
@@ -630,7 +636,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         if application_name:
             params["applicationName"] = application_name
         response = session_with_retries().get(
-            url, params=params, headers=self._get_header()
+            url, params=params, headers=self._get_headers()
         )
         response = request_handling(response)
         return parse_obj_as(List[Application], response)
@@ -651,7 +657,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         if deployment_id:
             params["deploymentId"] = deployment_id
         response = session_with_retries().get(
-            url, params=params, headers=self._get_header()
+            url, params=params, headers=self._get_headers()
         )
         response = request_handling(response)
         return parse_obj_as(List[Deployment], response)
@@ -661,7 +667,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         url = f"{self._api_server_url}/{VERSION_PREFIX}/apply"
         body = {"manifest": manifest, "dryRun": dry_run}
         response = session_with_retries().put(
-            url, headers=self._get_header(), json=body
+            url, headers=self._get_headers(), json=body
         )
         response_data = cast(Dict[str, Any], request_handling(response))
         return response_data
@@ -671,7 +677,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
         url = f"{self._api_server_url}/{VERSION_PREFIX}/delete"
         body = {"manifest": manifest}
         response = session_with_retries().post(
-            url, headers=self._get_header(), json=body
+            url, headers=self._get_headers(), json=body
         )
         response_data = cast(Dict[str, Any], request_handling(response))
         return response_data
@@ -691,7 +697,7 @@ class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
             # TODO (chiragjn): Check if this is supposed to be params or json
             params=body,
             json=body,
-            headers=self._get_header(),
+            headers=self._get_headers(),
         )
         response_data = cast(Dict[str, Any], request_handling(response))
         return response_data

truefoundry/deploy/v2/lib/deploy.py

@@ -6,7 +6,7 @@ import socketio
 from rich.status import Status
 
 from truefoundry.common.utils import poll_for_function
-from truefoundry.deploy.auto_gen import models as auto_gen_models
+from truefoundry.deploy._autogen import models as autogen_models
 from truefoundry.deploy.builder.docker_service import env_has_docker
 from truefoundry.deploy.lib.clients.servicefoundry_client import (
     ServiceFoundryServiceClient,
@@ -28,8 +28,8 @@ 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)
+        and isinstance(component.image, autogen_models.Build)
+        and isinstance(component.image.build_source, autogen_models.LocalSource)
     ):
         new_component = component.copy(deep=True)
 
@@ -180,13 +180,13 @@ def _deploy_wait_handler(  # noqa: C901
 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 isinstance(component.image, autogen_models.Build)
+        and isinstance(component.image.build_spec, autogen_models.PythonBuild)
         and not component.image.build_spec.cuda_version
     )
     uses_gpu = (
         hasattr(component, "resources")
-        and isinstance(component, auto_gen_models.Resources)
+        and isinstance(component, autogen_models.Resources)
         and component.resources.gpu_count > 0
     )
     if is_python_build_without_cuda and uses_gpu:
@@ -260,7 +260,7 @@ def deploy_component(
     )
     component.workspace_fqn = workspace_fqn
     workspace_id = get_workspace_by_fqn(workspace_fqn).id
-    if isinstance(component, auto_gen_models.ApplicationSet):
+    if isinstance(component, autogen_models.ApplicationSet):
         updated_component = component.copy(deep=True)
         for i in range(len(updated_component.components)):
             updated_component.components[i] = _handle_if_local_source(

truefoundry/deploy/v2/lib/deploy_workflow.py

@@ -20,7 +20,7 @@ from flytekit.tools.translator import WorkflowSpec as FlyteWorkflowSpec
 from google.protobuf.json_format import MessageToDict
 
 from truefoundry.common.types import UploadCodePackageCallable
-from truefoundry.deploy.auto_gen import models as auto_gen_models
+from truefoundry.deploy._autogen import models as autogen_models
 from truefoundry.deploy.lib.clients.servicefoundry_client import (
     ServiceFoundryServiceClient,
 )
@@ -38,10 +38,10 @@ from truefoundry.workflow.workflow import (
 
 
 def _handle_code_upload_for_workflow(
-    workflow: auto_gen_models.Workflow,
+    workflow: autogen_models.Workflow,
     workspace_fqn: str,
     upload_code_package: UploadCodePackageCallable,
-) -> auto_gen_models.Workflow:
+) -> autogen_models.Workflow:
     new_workflow = workflow.copy(deep=True)
     new_workflow.source = local_source_to_remote_source(
         local_source=workflow.source,
@@ -151,7 +151,7 @@ def _validate_workflow_entities(  # noqa: C901
             if not is_tfy_wf_present_in_task_python_build:
                 tasks_without_truefoundry_worflow_package.append(task.template.id.name)
         try:
-            auto_gen_models.FlyteTaskCustom.validate(task.template.custom)
+            autogen_models.FlyteTaskCustom.validate(task.template.custom)
         except ValidationError:
             raise ValueError(
                 error_message_to_use_truefoundry_decorators.format(
@@ -193,7 +193,7 @@ def _get_relative_package_path_from_filepath(
 
 
 def _generate_manifest_for_workflow(
-    workflow: auto_gen_models.Workflow,
+    workflow: autogen_models.Workflow,
 ):
     settings = SerializationSettings(
         # We are adding these defaults to avoid validation errors in flyte objects
@@ -248,7 +248,7 @@ def _generate_manifest_for_workflow(
             and message_dict["template"].get("custom")
             and message_dict["template"]["custom"].get("truefoundry")
         ):
-            parsed_model = auto_gen_models.FlyteTaskCustom.parse_obj(
+            parsed_model = autogen_models.FlyteTaskCustom.parse_obj(
                 message_dict["template"]["custom"]
             )
             message_dict["template"]["custom"]["truefoundry"] = parsed_model.truefoundry
@@ -256,11 +256,11 @@ def _generate_manifest_for_workflow(
         workflow.flyte_entities.append(message_dict)
 
     # this step is just to verify if pydantic model is still valid after adding flyte_entities
-    auto_gen_models.Workflow.validate({**workflow.dict()})
+    autogen_models.Workflow.validate({**workflow.dict()})
 
 
 def _validate_workspace_fqn(
-    workflow: auto_gen_models.Workflow,
+    workflow: autogen_models.Workflow,
     workspace_fqn: Optional[str] = None,
 ):
     if not workspace_fqn:
@@ -271,7 +271,7 @@ def _validate_workspace_fqn(
 
 
 def deploy_workflow(
-    workflow: auto_gen_models.Workflow,
+    workflow: autogen_models.Workflow,
     workspace_fqn: str,
     wait: bool = True,
     force: bool = False,

truefoundry/deploy/v2/lib/deployable_patched_models.py

@@ -1,6 +1,6 @@
 from typing import Literal, Union
 
-from truefoundry.deploy.auto_gen import models
+from truefoundry.deploy._autogen import models
 from truefoundry.deploy.lib.model.entity import Deployment
 from truefoundry.deploy.v2.lib.deploy import deploy_component
 from truefoundry.deploy.v2.lib.patched_models import LocalSource

truefoundry/deploy/v2/lib/patched_models.py

@@ -6,7 +6,7 @@ from collections.abc import Mapping
 from typing import Literal, Optional, Union
 
 from truefoundry.common.warnings import TrueFoundryDeprecationWarning
-from truefoundry.deploy.auto_gen import models
+from truefoundry.deploy._autogen import models
 from truefoundry.pydantic_v1 import (
     BaseModel,
     Field,
@@ -72,7 +72,6 @@ resources:
 ---------
 """
 
-
 AUTO_DISCOVERED_REQUIREMENTS_TXT_WARNING_MESSAGE_TEMPLATE = """\
 Using automatically discovered {requirements_txt_path} as the requirements file.
 This auto discovery behavior is deprecated and will be removed in a future release.
@@ -163,6 +162,7 @@ class GPUType(str, enum.Enum):
     L4 = "L4"
     H100_80GB = "H100_80GB"
     H100_94GB = "H100_94GB"
+    H200 = "H200"
 
 
 class TPUType(str, enum.Enum):
@@ -252,7 +252,7 @@ class BasicAuthCreds(models.BasicAuthCreds, PatchedModelBase):
     type: Literal["basic_auth"] = "basic_auth"
 
 
-class JwtAuthCreds(models.JwtAuthCreds, PatchedModelBase):
+class JwtAuthConfig(models.JwtAuthConfig, PatchedModelBase):
     type: Literal["jwt_auth"] = "jwt_auth"
 
 
@@ -508,10 +508,6 @@ class TrueFoundryArtifactSource(models.TrueFoundryArtifactSource, PatchedModelBa
     type: Literal["truefoundry-artifact"] = "truefoundry-artifact"
 
 
-# Deprecated alias, kept for backward compatibility
-TruefoundryArtifactSource = TrueFoundryArtifactSource
-
-
 class ArtifactsDownload(models.ArtifactsDownload, PatchedModelBase):
     pass
 
@@ -563,3 +559,7 @@ class TaskDockerFileBuild(models.TaskDockerFileBuild, PatchedModelBase):
 
 class TaskPythonBuild(models.TaskPythonBuild, PatchedModelBase):
     type: Literal["task-python-build"] = "task-python-build"
+
+
+# Deprecated aliases, kept for backward compatibility
+TruefoundryArtifactSource = TrueFoundryArtifactSource

truefoundry/deploy/v2/lib/source.py

@@ -11,7 +11,7 @@ from tqdm import tqdm
 from truefoundry.common.types import UploadCodePackageCallable
 from truefoundry.common.warnings import TrueFoundryDeprecationWarning
 from truefoundry.deploy import builder
-from truefoundry.deploy.auto_gen import models
+from truefoundry.deploy._autogen import models
 from truefoundry.deploy.builder.docker_service import (
     pull_docker_image,
     push_docker_image,

truefoundry/ml/__init__.py

@@ -1,10 +1,10 @@
-from truefoundry.ml.autogen.client.models import (  # type: ignore[attr-defined]
+from truefoundry.ml._autogen.client.models import (  # type: ignore[attr-defined]
     InferMethodName,
     ModelVersionEnvironment,
     SklearnModelSchema,
     XGBoostModelSchema,
 )
-from truefoundry.ml.autogen.entities.artifacts import LibraryName
+from truefoundry.ml._autogen.entities.artifacts import LibraryName
 from truefoundry.ml.enums import (
     DataSlice,
     FileFormat,