truefoundry 0.11.0__py3-none-any.whl → 0.11.2__py3-none-any.whl

truefoundry/common/utils.py

@@ -16,6 +16,7 @@ from truefoundry.common.constants import (
     API_SERVER_RELATIVE_PATH,
     ENV_VARS,
     MLFOUNDRY_SERVER_RELATIVE_PATH,
+    TFY_API_KEY_ENV_KEY,
     TFY_DEBUG_ENV_KEY,
     TFY_HOST_ENV_KEY,
     TFY_INTERNAL_ENV_KEY,
@@ -113,9 +114,14 @@ def validate_tfy_host(tfy_host: str) -> None:
 def resolve_tfy_host(tfy_host: Optional[str] = None) -> str:
     tfy_host = tfy_host or ENV_VARS.TFY_HOST
     if not tfy_host:
-        raise ValueError(
-            f"Either `host` should be provided using `--host <value>`, or `{TFY_HOST_ENV_KEY}` env must be set"
-        )
+        if ENV_VARS.TFY_API_KEY:
+            raise ValueError(
+                f"TFY_HOST` env must be set since `{TFY_API_KEY_ENV_KEY}` env is set. Either set `{TFY_HOST_ENV_KEY}` or unset `{TFY_API_KEY_ENV_KEY}` and login"
+            )
+        else:
+            raise ValueError(
+                f"Either `host` should be provided using `--host <value>`, or `{TFY_HOST_ENV_KEY}` env must be set"
+            )
     tfy_host = tfy_host.strip("/")
     validate_tfy_host(tfy_host)
     return tfy_host

truefoundry/deploy/builder/builders/tfy_spark_buildpack/tfy_execute_notebook.py

@@ -45,16 +45,9 @@ def upload_file_to_s3(file_path, bucket_name, s3_key):
     # Use s3proxy for pushing data to s3
     # The JWT token is already available in the pod
     aws_access_key_id = os.environ.get("SPARK_APPLICATION_EVENT_LOG_JWT_TOKEN")
-    aws_secret_access_key = "__token__"
+    aws_secret_access_key = os.environ.get("TFY_NOTEBOOK_OUTPUT_S3_SECRET_KEY")
     s3_endpoint_url = os.environ.get("S3_PROXY_URL")
 
-    if not aws_access_key_id:
-        raise ValueError(
-            "SPARK_APPLICATION_EVENT_LOG_JWT_TOKEN environment variable is not set"
-        )
-    if not s3_endpoint_url:
-        raise ValueError("S3_PROXY_URL environment variable is not set")
-
     # Needed for the issue https://github.com/gaul/s3proxy/issues/765
     s3_config = Config(
         request_checksum_calculation="when_required",
@@ -106,7 +99,22 @@ def execute_notebook(notebook_path, output_path="/tmp/output.ipynb", parameters=
         stderr_file=sys.stderr,
     )
     print(f"Successfully executed notebook: {notebook_path}")
-    return output_path
+
+
+def validate_env_vars():
+    keys = [
+        "TFY_NOTEBOOK_OUTPUT_S3_KEY",
+        "TFY_NOTEBOOK_OUTPUT_S3_BUCKET",
+        "SPARK_APPLICATION_EVENT_LOG_JWT_TOKEN",
+        "TFY_NOTEBOOK_OUTPUT_S3_SECRET_KEY",
+    ]
+    unset_keys = [key for key in keys if not os.environ.get(key)]
+    if unset_keys:
+        raise ValueError(
+            f"Environment variables {unset_keys} are not set."
+            f"Contact you tenant-admin to configure storage bucket on the control plane "
+            f"to enable uploading spark notebook outputs."
+        )
 
 
 if __name__ == "__main__":
@@ -116,35 +124,31 @@ if __name__ == "__main__":
     parser.add_argument("notebook_path", help="Path to the notebook file to execute")
     args = parser.parse_args()
 
+    # Since failure to upload is considered a job failure, fail the job even before it run if uploads cannot happen
+    validate_env_vars()
+
     output_notebook_path = "/tmp/output.ipynb"
 
     # This would be the same as the default bucket used by servicefoundry-server
     s3_bucket = os.environ.get("TFY_NOTEBOOK_OUTPUT_S3_BUCKET")
-    # This would be something like sparkjob-events/<tenant-id>
-    s3_key_prefix = os.environ.get("TFY_NOTEBOOK_OUTPUT_S3_KEY_PREFIX")
+    # This would be something like sparkjob-events/<tenant-id>/output-notebooks/<application-id>/<jobrun-name>/output.html
+    s3_key = os.environ.get("TFY_NOTEBOOK_OUTPUT_S3_KEY")
 
     try:
-        executed_notebook_path = execute_notebook(
-            args.notebook_path, output_path=output_notebook_path
-        )
+        execute_notebook(args.notebook_path, output_path=output_notebook_path)
 
         # The following may also be modeled as an entrypoint
         # https://papermill.readthedocs.io/en/latest/extending-entry-points.html
         # Will take that up with next iteration where we save the executed notebook periodically
-        if s3_bucket and s3_key_prefix:
-            print("Converting notebook to HTML and uploading to S3...")
-            html_output_path = "/tmp/output.html"
-            convert_notebook_to_html(
-                notebook_path=executed_notebook_path, output_html_path=html_output_path
-            )
-
-            # Construct S3 key: use the original notebook name for the HTML file
-            notebook_name = os.path.basename(args.notebook_path)
-            s3_html_key = f"{s3_key_prefix}/output.html"
-            upload_file_to_s3(
-                file_path=html_output_path, bucket_name=s3_bucket, s3_key=s3_html_key
-            )
-            print(f"Successfully uploaded HTML to s3://{s3_bucket}/{s3_html_key}")
+        print("Converting notebook to HTML and uploading to S3...")
+        html_output_path = "/tmp/output.html"
+        convert_notebook_to_html(
+            notebook_path=output_notebook_path, output_html_path=html_output_path
+        )
+        upload_file_to_s3(
+            file_path=html_output_path, bucket_name=s3_bucket, s3_key=s3_key
+        )
+        print(f"Successfully uploaded HTML to s3://{s3_bucket}/{s3_key}")
 
     except Exception as e:
         print(f"Error executing notebook {args.notebook_path}: {e}")

truefoundry/ml/artifact/truefoundry_artifact_repo.py

@@ -6,17 +6,7 @@ import uuid
 from concurrent.futures import FIRST_EXCEPTION, Future, ThreadPoolExecutor, wait
 from shutil import rmtree
 from threading import Event
-from typing import (
-    Any,
-    Callable,
-    Dict,
-    Iterator,
-    List,
-    Optional,
-    Sequence,
-    Tuple,
-    Union,
-)
+from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Tuple
 from urllib.parse import unquote
 from urllib.request import pathname2url
 
@@ -31,6 +21,7 @@ from rich.progress import (
 )
 from tqdm.utils import CallbackIOWrapper
 from truefoundry_sdk import (
+    FileInfo,
     MultiPartUploadResponse,
     MultiPartUploadStorageProvider,
     Operation,
@@ -52,11 +43,6 @@ from truefoundry.common.storage_provider_utils import (
 )
 from truefoundry.ml._autogen.client import (  # type: ignore[attr-defined]
     ApiClient,
-    FileInfoDto,
-    ListFilesForArtifactVersionRequestDto,
-    ListFilesForArtifactVersionsResponseDto,
-    ListFilesForDatasetRequestDto,
-    ListFilesForDatasetResponseDto,
     MlfoundryArtifactsApi,
     RunArtifactsApi,
 )
@@ -592,44 +578,23 @@ class MlFoundryArtifactsRepository:
             progress=progress,
         )
 
-    def _list_files(
-        self, artifact_identifier: ArtifactIdentifier, path, page_size, page_token
-    ) -> Union[ListFilesForDatasetResponseDto, ListFilesForArtifactVersionsResponseDto]:
-        if artifact_identifier.dataset_fqn:
-            return self._mlfoundry_artifacts_api.list_files_for_dataset_post(
-                list_files_for_dataset_request_dto=ListFilesForDatasetRequestDto(
-                    dataset_fqn=artifact_identifier.dataset_fqn,
-                    path=path,
-                    max_results=page_size,
-                    page_token=page_token,
-                )
-            )
-        else:
-            return self._mlfoundry_artifacts_api.list_files_for_artifact_version_post(
-                list_files_for_artifact_version_request_dto=ListFilesForArtifactVersionRequestDto(
-                    id=str(artifact_identifier.artifact_version_id),
-                    path=path,
-                    max_results=page_size,
-                    page_token=page_token,
-                )
-            )
-
     def list_artifacts(
         self, path=None, page_size=_LIST_FILES_PAGE_SIZE, **kwargs
-    ) -> Iterator[FileInfoDto]:
-        page_token = None
-        started = False
-        while not started or page_token is not None:
-            started = True
-            page = self._list_files(
-                artifact_identifier=self.artifact_identifier,
+    ) -> Iterator[FileInfo]:
+        if self.artifact_identifier.dataset_id:
+            for file_info in client.data_directories.list_files(
+                id=str(self.artifact_identifier.dataset_id),
                 path=path,
-                page_size=page_size,
-                page_token=page_token,
-            )
-            for file_info in page.files:
+                limit=page_size,
+            ):
+                yield file_info
+        else:
+            for file_info in client.artifact_versions.list_files(
+                id=str(self.artifact_identifier.artifact_version_id),
+                path=path,
+                limit=page_size,
+            ):
                 yield file_info
-            page_token = page.next_page_token
 
     def _is_directory(self, artifact_path):
         # TODO: Ideally server should return a flag to indicate if it is a directory

truefoundry/ml/mlfoundry_api.py

@@ -46,7 +46,6 @@ from truefoundry.ml._autogen.client.exceptions import (
     ApiException,
     NotFoundException,
 )
-from truefoundry.ml.clients.servicefoundry_client import ServiceFoundryServiceClient
 from truefoundry.ml.enums import ModelFramework, ViewType
 from truefoundry.ml.exceptions import MlFoundryException
 from truefoundry.ml.internal_namespace import NAMESPACE
@@ -139,7 +138,8 @@ class MlFoundry:
         except (IndexError, NotFoundError):
             err_msg = (
                 f"ML Repo Does Not Exist for name: {ml_repo}. You may either "
-                "create it from the dashboard or using client.create_ml_repo('<ml_repo_name>')"
+                f"create it from the dashboard or using "
+                f"`client.create_ml_repo(ml_repo='{ml_repo}', storage_integration_fqn='<storage_integration_fqn>')`"
             )
             raise MlFoundryException(err_msg) from None
         except ApiError as e:
@@ -191,7 +191,12 @@ class MlFoundry:
 
             client = get_client()
 
-            client.create_ml_repo(ml_repo="my-repo")
+            client.create_ml_repo(
+                ml_repo="my-repo",
+                # This controls which bucket is used.
+                # You can get this from Platform > Integrations > Copy FQN of any Blob Storage type integration.
+                storage_integration_fqn="..."
+            )
             ```
         """
         _validate_ml_repo_name(ml_repo_name=name)
@@ -235,28 +240,11 @@ class MlFoundry:
                 raise MlFoundryException(err_msg) from e
             return
 
-        servicefoundry_client = ServiceFoundryServiceClient(
-            tfy_host=self._api_client.tfy_host,
-            token=self._api_client.access_token,
-        )
-
-        assert existing_ml_repo.storage_integration_id is not None
-        try:
-            existing_storage_integration = (
-                servicefoundry_client.get_integration_from_id(
-                    existing_ml_repo.storage_integration_id
-                )
-            )
-        except Exception as e:
-            raise MlFoundryException(
-                "Error in getting storage integration for ML Repo"
-            ) from e
-
-        if existing_storage_integration["fqn"] != storage_integration_fqn:
+        if existing_ml_repo.manifest.storage_integration_fqn != storage_integration_fqn:
             raise MlFoundryException(
                 f"ML Repo with same name already exists with storage integration:"
-                f"{existing_storage_integration['fqn']}. Cannot update the storage integration to: "
-                f"{storage_integration_fqn}"
+                f"{existing_ml_repo.manifest.storage_integration_fqn!r}. "
+                f"Cannot update the storage integration to: {storage_integration_fqn!r}"
             )
 
     def create_run(
@@ -562,7 +550,8 @@ class MlFoundry:
         except (IndexError, NotFoundError):
             raise MlFoundryException(
                 f"ML Repo with name {ml_repo} does not exist. "
-                "You may either create it from the dashboard or using client.create_ml_repo('<ml_repo_name>')"
+                f"You may either create it from the dashboard or using "
+                f"`client.create_ml_repo(ml_repo='{ml_repo}', storage_integration_fqn='<storage_integration_fqn>')`"
             ) from None
         except ApiError as e:
             raise MlFoundryException(
@@ -1202,7 +1191,12 @@ class MlFoundry:
             client = get_client()
             ml_repo = "sample-repo"
 
-            client.create_ml_repo(ml_repo=ml_repo)
+            client.create_ml_repo(
+                ml_repo=ml_repo,
+                # This controls which bucket is used.
+                # You can get this from Platform > Integrations > Copy FQN of any Blob Storage type integration.
+                storage_integration_fqn="..."
+            )
             client.log_artifact(
                 ml_repo=ml_repo,
                 name="hello-world-file",
@@ -1301,8 +1295,8 @@ class MlFoundry:
             client.create_ml_repo(  # This is only required once
                 ml_repo="my-classification-project",
                 # This controls which bucket is used.
-                # You can get this from Integrations > Blob Storage. `None` picks the default
-                storage_integration_fqn=None
+                # You can get this from Platform > Integrations > Copy FQN of any Blob Storage type integration.
+                storage_integration_fqn="..."
             )
             model_version = client.log_model(
                 ml_repo="my-classification-project",

{truefoundry-0.11.0.dist-info → truefoundry-0.11.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: truefoundry
-Version: 0.11.0
+Version: 0.11.2
 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.2.0,>=0.1.5
+Requires-Dist: truefoundry-sdk<0.2.0,>=0.1.7
 Requires-Dist: typing-extensions>=4.0
 Requires-Dist: urllib3<3,>=1.26.18
 Requires-Dist: yq<4.0.0,>=3.1.0

{truefoundry-0.11.0.dist-info → truefoundry-0.11.2.dist-info}/RECORD

@@ -50,7 +50,7 @@ truefoundry/common/servicefoundry_client.py,sha256=2fYhdVPSvLXz5C5tosOq86JD8WM3I
 truefoundry/common/session.py,sha256=d9l3TEBpqVP4mr4mTGY1qVxc815skzMlNNdw14otg34,2923
 truefoundry/common/storage_provider_utils.py,sha256=yURhMw8k0FLFvaviRHDiifhvc6GnuQwGMC9Qd2uM440,10934
 truefoundry/common/types.py,sha256=BMJFCsR1lPJAw66IQBSvLyV4I6o_x5oj78gVsUa9si8,188
-truefoundry/common/utils.py,sha256=j3QP0uOsaGD_VmDDR68JTwoYE1okkAq6OqpVkzVf48Q,6424
+truefoundry/common/utils.py,sha256=P0FuAadoJGdpieUORLSN-PiFnkyoGO-K2cS4OPITBWg,6714
 truefoundry/common/warnings.py,sha256=xDMhR_-ZGC40Ycaj6nlFb5MYPexn8WbKCHd4FlflTXQ,705
 truefoundry/deploy/__init__.py,sha256=PVbGPU9S3-dTFn5LvLwaEnfsp2RrGT9iiM7_15kOV84,2837
 truefoundry/deploy/python_deploy_codegen.py,sha256=k19_m5DGsUyjOUCSKwIVP8vDna2sq01tHABsUfoVpW4,8019
@@ -67,7 +67,7 @@ truefoundry/deploy/builder/builders/tfy_python_buildpack/__init__.py,sha256=_fjq
 truefoundry/deploy/builder/builders/tfy_python_buildpack/dockerfile_template.py,sha256=f4l3fH21E2b8W3-JotMKc0AdPcCxV7LRPxxYJa7z_UQ,9134
 truefoundry/deploy/builder/builders/tfy_spark_buildpack/__init__.py,sha256=NEPlM6_vTVxp4ITa18B8DBbgYCn1q5d8be21lbgu5oY,2888
 truefoundry/deploy/builder/builders/tfy_spark_buildpack/dockerfile_template.py,sha256=2zohUaW8Yw_QREHlpRW7Pooomt19HJh44fHjlsiDmwM,6064
-truefoundry/deploy/builder/builders/tfy_spark_buildpack/tfy_execute_notebook.py,sha256=tTx-GZDVf5iB1Pyz2z5c2LH1yrb7lErFbJcr-giAIuI,5734
+truefoundry/deploy/builder/builders/tfy_spark_buildpack/tfy_execute_notebook.py,sha256=QvawKw30dcHROJ05XQU2KgwH3gtUmEGSkuLxiuPNJ2c,5899
 truefoundry/deploy/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truefoundry/deploy/cli/commands/__init__.py,sha256=qv818jxqSAygJ3h-6Ul8t-5VOgR_UrSgsVtNCl3e5G0,1408
 truefoundry/deploy/cli/commands/apply_command.py,sha256=DmXmKVokkauyKIiJDtErTwbJ5_LvQeJbTQsG5BjyKpo,2427
@@ -127,7 +127,7 @@ truefoundry/ml/exceptions.py,sha256=QpDJSWmF7dIsByS0qOQbQZ_jytdNTzkHDDO3BxhTSo0,
 truefoundry/ml/git_info.py,sha256=jvAVm9ilqivnGq8qJdUvYdd8Siv0PLtqurB-PXsS5ho,2023
 truefoundry/ml/internal_namespace.py,sha256=QcqMHp6-C2im2H_02hlhi01EIcr1HhNaZprszs13EMU,1790
 truefoundry/ml/logger.py,sha256=VT-BF3BnBYTWVq87O58F0c8uXMu94gYzsiFlGY3_7Ao,458
-truefoundry/ml/mlfoundry_api.py,sha256=hA6XjaONGEUS4beewZd5HEA0mGRew_wMvUJaQW7CThc,60168
+truefoundry/ml/mlfoundry_api.py,sha256=EZvFs-HiUDV2H0tDrgRdeotplWTZm3T5Uw-0Zq3WuHo,60174
 truefoundry/ml/mlfoundry_run.py,sha256=b0wm08OUAiO1UpzdDNBeJh-Gsm_O5kEdbY86Ugv-26Y,44181
 truefoundry/ml/model_framework.py,sha256=nVbKTtKDRBdLzt7Wrg5_vJKZ-awHbISGvL73s-V9egU,18975
 truefoundry/ml/prompt_utils.py,sha256=8FueyElVTXLnLtC3O6hKsW_snocArr_B8KG3Qv6eFIQ,2651
@@ -347,7 +347,7 @@ truefoundry/ml/_autogen/models/schema.py,sha256=a_bp42MMPUbwO3407m0UW2W8EOhnxZXf
 truefoundry/ml/_autogen/models/signature.py,sha256=rBjpxUIsEeWM0sIyYG5uCJB18DKHR4k5yZw8TzuoP48,4987
 truefoundry/ml/_autogen/models/utils.py,sha256=c7RtSLXhOLcP8rjuUtfnMdaKVTZvvbsmw98gPAkAFrs,24371
 truefoundry/ml/artifact/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-truefoundry/ml/artifact/truefoundry_artifact_repo.py,sha256=ocX5EIcLQa8Uc_C3NxxgNorpxc-z1Yp4TLvmzSORPpw,36862
+truefoundry/ml/artifact/truefoundry_artifact_repo.py,sha256=hbgLxSoihkLVuICzRueuh8iAIc-yruCW5TuMXYQ-aCU,35692
 truefoundry/ml/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truefoundry/ml/cli/cli.py,sha256=MwpY7z_NEeJE_XIP7XbZELjNeu2vpMmohttHCKDRk54,335
 truefoundry/ml/cli/utils.py,sha256=j6_mZ4Spn114mz3P4QQ8jx0tmorXIuyQnHXVUSDvZi4,1035
@@ -355,7 +355,6 @@ truefoundry/ml/cli/commands/__init__.py,sha256=diDUiRUX4l6TtNLI4iF-ZblczkELM7FRV
 truefoundry/ml/cli/commands/download.py,sha256=N9MhsEQ3U24v_OmnMZT8Q4SoAi38Sm7a21unrACOSDw,2573
 truefoundry/ml/cli/commands/model_init.py,sha256=INyUAU6hiFClI8cZqX5hgnrtNbeKxlZxrjFrjzStU18,2664
 truefoundry/ml/clients/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-truefoundry/ml/clients/servicefoundry_client.py,sha256=GiB_SmETI2uo8E0Q43cEsOG6qYb80nB6gM7l3rL7zEs,2422
 truefoundry/ml/log_types/__init__.py,sha256=g4u4D4Jaj0aBK5GtrLV88-qThKZR9pSZ17vFEkN-LmM,125
 truefoundry/ml/log_types/plot.py,sha256=LDh4uy6z2P_a2oPM2lc85c0lt8utVvunohzeMawFjZw,7572
 truefoundry/ml/log_types/pydantic_base.py,sha256=eBlw_AEyAz4iJKDP4zgJOCFWcldwQqpf7FADW1jzIQY,272
@@ -381,7 +380,7 @@ truefoundry/workflow/remote_filesystem/__init__.py,sha256=LQ95ViEjJ7Ts4JcCGOxMPs
 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.11.0.dist-info/METADATA,sha256=nxfBh-eMI7_N9k5Baa3I_H3UwKMNOUH2jWPuE5-l4TE,2505
-truefoundry-0.11.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-truefoundry-0.11.0.dist-info/entry_points.txt,sha256=xVjn7RMN-MW2-9f7YU-bBdlZSvvrwzhpX1zmmRmsNPU,98
-truefoundry-0.11.0.dist-info/RECORD,,
+truefoundry-0.11.2.dist-info/METADATA,sha256=7n9RCZ9gGKy-q0nVRgja-FIyJEBcgo4eXtbrCPpfIFk,2505
+truefoundry-0.11.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+truefoundry-0.11.2.dist-info/entry_points.txt,sha256=xVjn7RMN-MW2-9f7YU-bBdlZSvvrwzhpX1zmmRmsNPU,98
+truefoundry-0.11.2.dist-info/RECORD,,

truefoundry/ml/clients/servicefoundry_client.py

@@ -1,65 +0,0 @@
-import functools
-
-from truefoundry.common.constants import (
-    SERVICEFOUNDRY_CLIENT_MAX_RETRIES,
-    VERSION_PREFIX,
-)
-from truefoundry.common.exceptions import HttpRequestException
-from truefoundry.common.request_utils import (
-    http_request,
-    request_handling,
-    requests_retry_session,
-)
-from truefoundry.common.servicefoundry_client import (
-    ServiceFoundryServiceClient as BaseServiceFoundryServiceClient,
-)
-from truefoundry.common.utils import get_user_agent
-from truefoundry.ml.exceptions import MlFoundryException
-
-
-class ServiceFoundryServiceClient(BaseServiceFoundryServiceClient):
-    def __init__(self, tfy_host: str, token: str):
-        super().__init__(tfy_host=tfy_host)
-        self._token = token
-
-    @functools.cached_property
-    def _min_cli_version_required(self) -> str:
-        # TODO (chiragjn): read the mlfoundry min cli version from the config?
-        return self.python_sdk_config.truefoundry_cli_min_version
-
-    def get_integration_from_id(self, integration_id: str):
-        integration_id = integration_id or ""
-        session = requests_retry_session(retries=SERVICEFOUNDRY_CLIENT_MAX_RETRIES)
-        response = http_request(
-            method="get",
-            url=f"{self._api_server_url}/{VERSION_PREFIX}/provider-accounts/provider-integrations",
-            token=self._token,
-            timeout=3,
-            params={"id": integration_id, "type": "blob-storage"},
-            session=session,
-            headers={
-                "User-Agent": get_user_agent(),
-            },
-        )
-
-        try:
-            result = request_handling(response)
-            assert isinstance(result, dict)
-        except HttpRequestException as he:
-            raise MlFoundryException(
-                f"Failed to get storage integration from id: {integration_id}. Error: {he.message}",
-                status_code=he.status_code,
-            ) from None
-        except Exception as e:
-            raise MlFoundryException(
-                f"Failed to get storage integration from id: {integration_id}. Error: {str(e)}"
-            ) from None
-
-        data = result.get("data", result.get("providerAccounts"))
-        # TODO (chiragjn): Parse this using Pydantic
-        if data and len(data) > 0 and data[0]:
-            return data[0]
-        else:
-            raise MlFoundryException(
-                f"Invalid storage integration id: {integration_id}"
-            )