truefoundry 0.10.7rc1__py3-none-any.whl → 0.11.1rc1__py3-none-any.whl

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

@@ -47,7 +47,12 @@ RUN apt update && \
     + _POST_USER_TEMPLATE
 )
 
-ADDITIONAL_PIP_PACKAGES = ["papermill>=2.6.0,<2.7.0", "ipykernel>=6.0.0,<7.0.0"]
+ADDITIONAL_PIP_PACKAGES = [
+    "papermill>=2.6.0,<2.7.0",
+    "ipykernel>=6.0.0,<7.0.0",
+    "nbconvert>=7.16.6,<7.17.0",
+    "boto3>=1.38.43,<1.40.0",
+]
 
 
 def generate_pip_install_command(

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

@@ -4,9 +4,77 @@
 
 
 import argparse
+import os
 import sys
 
+import boto3
+import nbformat
 import papermill as pm
+from botocore.client import Config
+from nbconvert import HTMLExporter
+
+
+def convert_notebook_to_html(notebook_path, output_html_path):
+    """
+    Convert a Jupyter notebook to an HTML file.
+
+    Args:
+        notebook_path: Path to the input notebook (.ipynb)
+        output_html_path: Path for the output HTML file (.html)
+    """
+    print(f"Converting notebook {notebook_path} to HTML...")
+    try:
+        with open(notebook_path, "r", encoding="utf-8") as f:
+            notebook_content = nbformat.read(f, as_version=4)
+
+        html_exporter = HTMLExporter()
+        # Use lab for https://nbconvert.readthedocs.io/en/latest/customizing.html#where-are-nbconvert-templates-installed
+        html_exporter.template_name = "lab"
+        (body, _) = html_exporter.from_notebook_node(notebook_content)
+
+        with open(output_html_path, "w", encoding="utf-8") as f:
+            f.write(body)
+        print(f"Successfully converted notebook to HTML: {output_html_path}")
+    except Exception as e:
+        print(f"Error converting notebook to HTML: {e}")
+        raise
+
+
+def upload_file_to_s3(file_path, bucket_name, s3_key):
+    print(f"Uploading {file_path} to s3://{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__"
+    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",
+        response_checksum_validation="when_required",
+    )
+    try:
+        client = boto3.client(
+            "s3",
+            aws_access_key_id=aws_access_key_id,
+            aws_secret_access_key=aws_secret_access_key,
+            endpoint_url=s3_endpoint_url,
+            config=s3_config,
+        )
+
+        with open(file_path, "rb") as data:
+            client.put_object(Bucket=bucket_name, Key=s3_key, Body=data)
+        print(f"Successfully uploaded {file_path} to s3://{bucket_name}/{s3_key}")
+    except Exception as e:
+        print(f"Error uploading file to S3: {e}")
+        raise
 
 
 def execute_notebook(notebook_path, output_path="/tmp/output.ipynb", parameters=None):
@@ -26,8 +94,8 @@ def execute_notebook(notebook_path, output_path="/tmp/output.ipynb", parameters=
 
     print(f"Starting execution of notebook: {notebook_path}")
     pm.execute_notebook(
-        notebook_path,
-        output_path,
+        input_path=notebook_path,
+        output_path=output_path,
         parameters=parameters,
         # TODO(gw): Replace with kernel name for venv
         kernel_name="python3",
@@ -38,6 +106,7 @@ 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
 
 
 if __name__ == "__main__":
@@ -45,17 +114,41 @@ if __name__ == "__main__":
         description="Execute a Jupyter notebook using papermill for Spark applications"
     )
     parser.add_argument("notebook_path", help="Path to the notebook file to execute")
-
     args = parser.parse_args()
 
-    # TODO(gw): Add support for passing parameters to the notebook
+    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")
+
     try:
-        execute_notebook(args.notebook_path)
+        executed_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}")
+
     except Exception as e:
         print(f"Error executing notebook {args.notebook_path}: {e}")
         print(
             "Exiting with status code 1 to signal failure to parent process/orchestrator"
         )
         sys.exit(1)
-
-    # TODO(gw): Publish the output notebook to blob storage from where it could be rendered

truefoundry/ml/artifact/truefoundry_artifact_repo.py

@@ -30,7 +30,14 @@ from rich.progress import (
     TransferSpeedColumn,
 )
 from tqdm.utils import CallbackIOWrapper
+from truefoundry_sdk import (
+    MultiPartUploadResponse,
+    MultiPartUploadStorageProvider,
+    Operation,
+    SignedUrl,
+)
 
+from truefoundry import client
 from truefoundry.common.constants import ENV_VARS
 from truefoundry.common.request_utils import (
     augmented_raise_for_status,
@@ -45,23 +52,13 @@ from truefoundry.common.storage_provider_utils import (
 )
 from truefoundry.ml._autogen.client import (  # type: ignore[attr-defined]
     ApiClient,
-    CreateMultiPartUploadForDatasetRequestDto,
-    CreateMultiPartUploadRequestDto,
     FileInfoDto,
-    GetSignedURLForDatasetWriteRequestDto,
-    GetSignedURLsForArtifactVersionReadRequestDto,
-    GetSignedURLsForArtifactVersionWriteRequestDto,
-    GetSignedURLsForDatasetReadRequestDto,
     ListFilesForArtifactVersionRequestDto,
     ListFilesForArtifactVersionsResponseDto,
     ListFilesForDatasetRequestDto,
     ListFilesForDatasetResponseDto,
     MlfoundryArtifactsApi,
-    MultiPartUploadDto,
-    MultiPartUploadResponseDto,
-    MultiPartUploadStorageProvider,
     RunArtifactsApi,
-    SignedURLDto,
 )
 from truefoundry.ml.exceptions import MlFoundryException
 from truefoundry.ml.logger import logger
@@ -127,7 +124,7 @@ def verify_artifact_path(artifact_path):
 
 
 def _signed_url_upload_file(
-    signed_url: SignedURLDto,
+    signed_url: SignedUrl,
     local_file: str,
     progress_bar: Progress,
     abort_event: Optional[Event] = None,
@@ -206,21 +203,20 @@ def _any_future_has_failed(futures) -> bool:
 
 class ArtifactIdentifier(BaseModel):
     artifact_version_id: Optional[uuid.UUID] = None
+    dataset_id: Optional[str] = None
     dataset_fqn: Optional[str] = None
 
     @root_validator
     def _check_identifier_type(cls, values: Dict[str, Any]):
         if not values.get("artifact_version_id", False) and not values.get(
-            "dataset_fqn", False
+            "dataset_id", False
         ):
             raise MlFoundryException(
-                "One of the version_id or dataset_fqn should be passed"
+                "One of the version_id or dataset_id should be passed"
             )
-        if values.get("artifact_version_id", False) and values.get(
-            "dataset_fqn", False
-        ):
+        if values.get("artifact_version_id", False) and values.get("dataset_id", False):
             raise MlFoundryException(
-                "Exactly one of version_id or dataset_fqn should be passed"
+                "Exactly one of version_id or dataset_id should be passed"
             )
         return values
 
@@ -243,24 +239,25 @@ class MlFoundryArtifactsRepository:
         self,
         artifact_identifier: ArtifactIdentifier,
         paths,
-    ) -> List[SignedURLDto]:
+    ) -> List[SignedUrl]:
         if artifact_identifier.artifact_version_id:
-            signed_urls_response = self._mlfoundry_artifacts_api.get_signed_urls_for_read_post(
-                get_signed_urls_for_artifact_version_read_request_dto=GetSignedURLsForArtifactVersionReadRequestDto(
-                    id=str(artifact_identifier.artifact_version_id), paths=paths
-                )
+            signed_urls_response = client.artifact_versions.get_signed_urls(
+                id=str(artifact_identifier.artifact_version_id),
+                paths=paths,
+                operation=Operation.READ,
             )
-            signed_urls = signed_urls_response.signed_urls
-        elif artifact_identifier.dataset_fqn:
-            signed_urls_dataset_response = self._mlfoundry_artifacts_api.get_signed_urls_dataset_read_post(
-                get_signed_urls_for_dataset_read_request_dto=GetSignedURLsForDatasetReadRequestDto(
-                    dataset_fqn=artifact_identifier.dataset_fqn, paths=paths
-                )
+            signed_urls = signed_urls_response.data
+        elif artifact_identifier.dataset_id:
+            signed_urls_dataset_response = client.data_directories.get_signed_urls(
+                id=str(artifact_identifier.dataset_id),
+                paths=paths,
+                operation=Operation.READ,
             )
-            signed_urls = signed_urls_dataset_response.signed_urls
+            signed_urls = signed_urls_dataset_response.data
+
         else:
             raise ValueError(
-                "Invalid artifact type - both `artifact_version_id` and `dataset_fqn` both are None"
+                "Invalid artifact type - both `artifact_version_id` and `dataset_id` both are None"
             )
         return signed_urls
 
@@ -268,24 +265,24 @@ class MlFoundryArtifactsRepository:
         self,
         artifact_identifier: ArtifactIdentifier,
         paths: List[str],
-    ) -> List[SignedURLDto]:
+    ) -> List[SignedUrl]:
         if artifact_identifier.artifact_version_id:
-            signed_urls_response = self._mlfoundry_artifacts_api.get_signed_urls_for_write_post(
-                get_signed_urls_for_artifact_version_write_request_dto=GetSignedURLsForArtifactVersionWriteRequestDto(
-                    id=str(artifact_identifier.artifact_version_id), paths=paths
-                )
+            signed_urls_response = client.artifact_versions.get_signed_urls(
+                id=str(artifact_identifier.artifact_version_id),
+                paths=paths,
+                operation=Operation.WRITE,
             )
-            signed_urls = signed_urls_response.signed_urls
-        elif artifact_identifier.dataset_fqn:
-            signed_urls_dataset_response = self._mlfoundry_artifacts_api.get_signed_urls_for_dataset_write_post(
-                get_signed_url_for_dataset_write_request_dto=GetSignedURLForDatasetWriteRequestDto(
-                    dataset_fqn=artifact_identifier.dataset_fqn, paths=paths
-                )
+            signed_urls = signed_urls_response.data
+        elif artifact_identifier.dataset_id:
+            signed_urls_dataset_response = client.data_directories.get_signed_urls(
+                id=str(artifact_identifier.dataset_id),
+                paths=paths,
+                operation=Operation.WRITE,
             )
-            signed_urls = signed_urls_dataset_response.signed_urls
+            signed_urls = signed_urls_dataset_response.data
         else:
             raise ValueError(
-                "Invalid artifact type - both `artifact_version_id` and `dataset_fqn` both are None"
+                "Invalid artifact type - both `artifact_version_id` and `dataset_id` both are None"
             )
         return signed_urls
 
@@ -293,7 +290,7 @@ class MlFoundryArtifactsRepository:
         self,
         local_file: str,
         artifact_path: str,
-        signed_url: Optional[SignedURLDto],
+        signed_url: Optional[SignedUrl],
         progress_bar: Progress,
         abort_event: Optional[Event] = None,
     ):
@@ -322,28 +319,30 @@ class MlFoundryArtifactsRepository:
         artifact_identifier: ArtifactIdentifier,
         path,
         num_parts,
-    ) -> MultiPartUploadDto:
+    ) -> MultiPartUpload:
         if artifact_identifier.artifact_version_id:
-            create_multipart_response: MultiPartUploadResponseDto = self._mlfoundry_artifacts_api.create_multi_part_upload_post(
-                create_multi_part_upload_request_dto=CreateMultiPartUploadRequestDto(
-                    artifact_version_id=str(artifact_identifier.artifact_version_id),
+            create_multipart_response: MultiPartUploadResponse = (
+                client.artifact_versions.create_multi_part_upload(
+                    id=str(artifact_identifier.artifact_version_id),
                     path=path,
                     num_parts=num_parts,
                 )
             )
-            multipart_upload = create_multipart_response.multipart_upload
-        elif artifact_identifier.dataset_fqn:
-            create_multipart_for_dataset_response = self._mlfoundry_artifacts_api.create_multipart_upload_for_dataset_post(
-                create_multi_part_upload_for_dataset_request_dto=CreateMultiPartUploadForDatasetRequestDto(
-                    dataset_fqn=artifact_identifier.dataset_fqn,
+            multipart_upload = create_multipart_response.data
+
+        elif artifact_identifier.dataset_id:
+            create_multipart_for_dataset_response: MultiPartUploadResponse = (
+                client.data_directories.create_multipart_upload(
+                    id=str(artifact_identifier.dataset_id),
                     path=path,
                     num_parts=num_parts,
                 )
             )
-            multipart_upload = create_multipart_for_dataset_response.multipart_upload
+            multipart_upload = create_multipart_for_dataset_response.data
+
         else:
             raise ValueError(
-                "Invalid artifact type - both `artifact_version_id` and `dataset_fqn` both are None"
+                "Invalid artifact type - both `artifact_version_id` and `dataset_id` both are None"
             )
         return multipart_upload
 
@@ -370,10 +369,10 @@ class MlFoundryArtifactsRepository:
         )
         if (
             multipart_upload.storage_provider
-            is MultiPartUploadStorageProvider.S3_COMPATIBLE
+            is MultiPartUploadStorageProvider.S3COMPATIBLE
         ):
             s3_compatible_multipart_upload(
-                multipart_upload=MultiPartUpload.parse_obj(multipart_upload.to_dict()),
+                multipart_upload=MultiPartUpload.parse_obj(multipart_upload.dict()),
                 local_file=local_file,
                 executor=executor,
                 multipart_info=multipart_info,
@@ -403,7 +402,7 @@ class MlFoundryArtifactsRepository:
         artifact_path: str,
         multipart_info: _FileMultiPartInfo,
         progress_bar: Progress,
-        signed_url: Optional[SignedURLDto] = None,
+        signed_url: Optional[SignedUrl] = None,
         abort_event: Optional[Event] = None,
         executor_for_multipart_upload: Optional[ThreadPoolExecutor] = None,
     ):
@@ -678,7 +677,7 @@ class MlFoundryArtifactsRepository:
         remote_file_path: str,
         local_path: str,
         progress_bar: Optional[Progress],
-        signed_url: Optional[SignedURLDto],
+        signed_url: Optional[SignedUrl],
         abort_event: Optional[Event] = None,
     ):
         if not remote_file_path:
@@ -719,7 +718,7 @@ class MlFoundryArtifactsRepository:
         self,
         src_artifact_path: str,
         dst_local_dir_path: str,
-        signed_url: Optional[SignedURLDto],
+        signed_url: Optional[SignedUrl],
         progress_bar: Optional[Progress] = None,
         abort_event=None,
     ) -> str:

truefoundry/ml/log_types/artifacts/artifact.py

@@ -8,6 +8,7 @@ import warnings
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, NamedTuple, Optional, Union
 
+from truefoundry import client
 from truefoundry.common.warnings import TrueFoundryDeprecationWarning
 from truefoundry.ml._autogen.client import (  # type: ignore[attr-defined]
     ArtifactDto,
@@ -16,7 +17,6 @@ from truefoundry.ml._autogen.client import (  # type: ignore[attr-defined]
     ArtifactVersionDto,
     CreateArtifactVersionRequestDto,
     DeleteArtifactVersionsRequestDto,
-    ExperimentsApi,
     ExternalBlobStorageSource,
     FinalizeArtifactVersionRequestDto,
     MlfoundryArtifactsApi,
@@ -444,19 +444,15 @@ def _log_artifact_version_helper(
 
     if run:
         mlfoundry_artifacts_api = run._mlfoundry_artifacts_api
-        experiments_api = run._experiments_api
-        ml_repo = experiments_api.get_experiment_get(
-            experiment_id=str(run._experiment_id)
-        ).experiment.name
+        repos = client.ml_repos.get(id=run._experiment_id)
+        ml_repo = repos.data.manifest.name
         ml_repo_id = run._experiment_id
     else:
         assert ml_repo is not None
         api_client = _get_api_client()
         mlfoundry_artifacts_api = MlfoundryArtifactsApi(api_client=api_client)
-        experiments_api = ExperimentsApi(api_client=api_client)
-        ml_repo_id = experiments_api.get_experiment_by_name_get(
-            experiment_name=ml_repo
-        ).experiment.experiment_id
+        result = list(client.ml_repos.list(name=ml_repo, limit=1))[0]
+        ml_repo_id = result.id
 
     metadata = metadata or {}
     if ARTIFACT_METADATA_TRUEFOUNDRY_KEY not in metadata:

truefoundry/ml/log_types/artifacts/dataset.py

@@ -84,7 +84,9 @@ class DataDirectory:
 
     def _get_artifacts_repo(self):
         return MlFoundryArtifactsRepository(
-            artifact_identifier=ArtifactIdentifier(dataset_fqn=self._dataset.fqn),
+            artifact_identifier=ArtifactIdentifier(
+                dataset_id=self._dataset.id, dataset_fqn=self._dataset.fqn
+            ),
             api_client=self._api_client,
         )
 

truefoundry/ml/log_types/artifacts/model.py

@@ -9,12 +9,12 @@ import warnings
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
+from truefoundry import client
 from truefoundry.common.warnings import TrueFoundryDeprecationWarning
 from truefoundry.ml._autogen.client import (  # type: ignore[attr-defined]
     ArtifactType,
     CreateArtifactVersionRequestDto,
     DeleteArtifactVersionsRequestDto,
-    ExperimentsApi,
     ExternalBlobStorageSource,
     FinalizeArtifactVersionRequestDto,
     Framework,
@@ -531,19 +531,15 @@ def _log_model_version(  # noqa: C901
 
     if run:
         mlfoundry_artifacts_api = run._mlfoundry_artifacts_api
-        experiments_api = run._experiments_api
-        ml_repo = experiments_api.get_experiment_get(
-            experiment_id=str(run._experiment_id)
-        ).experiment.name
+        repos = client.ml_repos.get(id=run._experiment_id)
+        ml_repo = repos.data.manifest.name
         ml_repo_id = run._experiment_id
     else:
         assert ml_repo is not None
         api_client = _get_api_client()
         mlfoundry_artifacts_api = MlfoundryArtifactsApi(api_client=api_client)
-        experiments_api = ExperimentsApi(api_client=api_client)
-        ml_repo_id = experiments_api.get_experiment_by_name_get(
-            experiment_name=ml_repo
-        ).experiment.experiment_id
+        result = list(client.ml_repos.list(name=ml_repo, limit=1))[0]
+        ml_repo_id = result.id
 
     step = step or 0
     total_size = None

truefoundry/ml/mlfoundry_api.py

@@ -16,14 +16,16 @@ from typing import (
 )
 
 import coolname
+from truefoundry_sdk import Collaborator, MlRepoManifest, NotFoundError
+from truefoundry_sdk.core import ApiError
 
+from truefoundry import client
 from truefoundry.common.utils import ContextualDirectoryManager
 from truefoundry.ml import constants
 from truefoundry.ml._autogen.client import (  # type: ignore[attr-defined]
     ArtifactDto,
     ArtifactType,
     CreateDatasetRequestDto,
-    CreateExperimentRequestDto,
     CreateRunRequestDto,
     DatasetDto,
     ExperimentsApi,
@@ -44,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
@@ -132,24 +133,22 @@ class MlFoundry:
             str: The id of the ML Repo.
         """
         try:
-            _ml_repo_instance = self._experiments_api.get_experiment_by_name_get(
-                experiment_name=ml_repo
-            )
-            ml_repo_instance = _ml_repo_instance.experiment
-        except NotFoundException:
+            result = list(client.ml_repos.list(name=ml_repo, limit=1))
+            ml_repo_instance = result[0]
+        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 ApiException as e:
+        except ApiError as e:
             err_msg = (
                 f"Error happened in getting ML Repo based on name: "
                 f"{ml_repo}. Error details: {e}"
             )
             raise MlFoundryException(err_msg) from e
-
-        return ml_repo_instance.experiment_id
+        return ml_repo_instance.id
 
     def list_ml_repos(self) -> List[str]:
         """Returns a list of names of ML Repos accessible by the current user.
@@ -160,26 +159,14 @@ class MlFoundry:
         # TODO (chiragjn): This API should yield ML Repo Entities instead of just names
         #   Kinda useless without it
         ml_repos_names = []
-        done, page_token, max_results = False, None, 25
-        while not done:
-            try:
-                _ml_repos = self._experiments_api.list_experiments_get(
-                    view_type=ViewType.ALL.value,
-                    max_results=max_results,
-                    page_token=page_token,
-                )
-            except ApiException as e:
-                err_msg = f"Error happened in fetching ML Repos. Error details: {e}"
-                raise MlFoundryException(err_msg) from e
-            else:
-                ml_repos = _ml_repos.experiments
-                page_token = _ml_repos.next_page_token
-                for ml_repo in ml_repos:
-                    # ML Repo with experiment_id 0 represents default ML Repo which we are removing.
-                    if ml_repo.experiment_id != "0":
-                        ml_repos_names.append(ml_repo.name)
-                if not ml_repos or page_token is None:
-                    done = True
+        max_results = 25
+        try:
+            for ml_repo_instance in client.ml_repos.list(limit=max_results):
+                if ml_repo_instance.id != "0":
+                    ml_repos_names.append(ml_repo_instance.manifest.name)
+        except ApiError as e:
+            err_msg = f"Error happened in fetching ML Repos. Error details: {e}"
+            raise MlFoundryException(err_msg) from e
         return ml_repos_names
 
     def create_ml_repo(
@@ -204,56 +191,60 @@ 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)
         if description:
             _validate_ml_repo_description(description=description)
         try:
-            _ml_repo_instance = self._experiments_api.get_experiment_by_name_get(
-                experiment_name=name
-            )
-            existing_ml_repo = _ml_repo_instance.experiment
-        except NotFoundException:
+            result = list(client.ml_repos.list(name=name, limit=1))
+            existing_ml_repo = result[0]
+        except (IndexError, NotFoundError):
             existing_ml_repo = None
+        except ApiError as e:
+            err_msg = (
+                f"Error happened in getting ML Repo based on name: "
+                f"{name}. Error details: {e}"
+            )
+            raise MlFoundryException(err_msg) from e
 
         if not existing_ml_repo:
+            session = client._get_session()
+            user_info = session.user_info
+            if not user_info.email:
+                raise MlFoundryException(
+                    "Virtual accounts cannot be used to create new ML Repos"
+                )
             try:
-                self._experiments_api.create_experiment_post(
-                    create_experiment_request_dto=CreateExperimentRequestDto(
+                client.ml_repos.create_or_update(
+                    manifest=MlRepoManifest(
                         name=name,
                         description=description,
                         storage_integration_fqn=storage_integration_fqn,
-                    )
+                        collaborators=[
+                            Collaborator(
+                                subject=f"user:{user_info.email}",
+                                role_id="mlf-project-admin",
+                            )
+                        ],
+                    ),
                 )
-            except ApiException as e:
+            except ApiError as e:
                 err_msg = f"Error happened in creating ML Repo with name: {name}. Error details: {e}"
                 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(
@@ -554,16 +545,20 @@ class MlFoundry:
         """
         _validate_ml_repo_name(ml_repo_name=ml_repo)
         try:
-            _ml_repo_obj = self._experiments_api.get_experiment_by_name_get(
-                experiment_name=ml_repo
-            )
-            ml_repo_obj = _ml_repo_obj.experiment
-        except ApiException as e:
+            result = list(client.ml_repos.list(name=ml_repo, limit=1))
+            ml_repo_instance = result[0]
+        except (IndexError, NotFoundError):
+            raise MlFoundryException(
+                f"ML Repo with name {ml_repo} does not exist. "
+                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(
                 f"ML Repo with name {ml_repo} does not exist or your user does not have permission to access it: {e}"
             ) from e
 
-        ml_repo_id = ml_repo_obj.experiment_id
+        ml_repo_id = ml_repo_instance.id
 
         page_token = None
         done = False
@@ -1196,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",
@@ -1295,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/ml/mlfoundry_run.py

@@ -16,12 +16,11 @@ from typing import (
 )
 from urllib.parse import urljoin, urlsplit
 
-from truefoundry import version
+from truefoundry import client, version
 from truefoundry.ml import constants
 from truefoundry.ml._autogen.client import (  # type: ignore[attr-defined]
     ArtifactType,
     DeleteRunRequest,
-    ExperimentsApi,
     ListArtifactVersionsRequestDto,
     ListModelVersionsRequestDto,
     LogBatchRequestDto,
@@ -110,7 +109,6 @@ class MlFoundryRun:
             ACTIVE_RUNS.add_run(self)
 
         self._api_client = _get_api_client()
-        self._experiments_api = ExperimentsApi(api_client=self._api_client)
         self._runs_api = RunsApi(api_client=self._api_client)
         self._metrics_api = MetricsApi(api_client=self._api_client)
         self._mlfoundry_artifacts_api = MlfoundryArtifactsApi(
@@ -165,10 +163,8 @@ class MlFoundryRun:
     @_ensure_not_deleted
     def ml_repo(self) -> str:
         """Get ml_repo name of which the current `run` is part of"""
-        _experiment = self._experiments_api.get_experiment_get(
-            experiment_id=self._experiment_id
-        )
-        return _experiment.experiment.name
+        _experiment = client.ml_repos.get(id=self._experiment_id)
+        return _experiment.data.manifest.name
 
     @property
     @_ensure_not_deleted

{truefoundry-0.10.7rc1.dist-info → truefoundry-0.11.1rc1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: truefoundry
-Version: 0.10.7rc1
+Version: 0.11.1rc1
 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.3
+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.10.7rc1.dist-info → truefoundry-0.11.1rc1.dist-info}/RECORD

@@ -66,8 +66,8 @@ truefoundry/deploy/builder/builders/tfy_notebook_buildpack/dockerfile_template.p
 truefoundry/deploy/builder/builders/tfy_python_buildpack/__init__.py,sha256=_fjqHKn80qKi68SAMMALge7_A6e1sTsQWichw8uoGIw,2025
 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=dAmWgAHcNdBLSedpz18r5r7n9hdXKV-PB7yFPMdaoww,5990
-truefoundry/deploy/builder/builders/tfy_spark_buildpack/tfy_execute_notebook.py,sha256=PKJWNlW76dRw_XtFkRnqpzwBGuXU4w6_kIJkgoBmZ5Q,1975
+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/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,8 +127,8 @@ 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=WiIVpJRylBn8NrcAmXClQnavffqI9fhY5h-8Vwx4RKo,59872
-truefoundry/ml/mlfoundry_run.py,sha256=34yyQqgpG6EfrAJd37vkbCjrFoHkhvbOAxxSQcSWPtY,44320
+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
 truefoundry/ml/run_utils.py,sha256=0W208wSLUrbdfk2pjNcZlkUi9bNxG2JORqoe-5rVqHI,2423
@@ -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=0ng8RqSWufpPwz-uG8iyWkiKcrZ5lUAmS4_UOgmzkx0,37738
+truefoundry/ml/artifact/truefoundry_artifact_repo.py,sha256=ocX5EIcLQa8Uc_C3NxxgNorpxc-z1Yp4TLvmzSORPpw,36862
 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,16 +355,15 @@ 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
 truefoundry/ml/log_types/utils.py,sha256=xjJ21jdPScvFmw3TbVh5NCzbzJwaqiXJyiiT4xxX1EI,335
-truefoundry/ml/log_types/artifacts/artifact.py,sha256=N1M1a7Oq9fY_7s1PT7uv64s_Ek62Lj-JsPBrWdRKpdY,20050
+truefoundry/ml/log_types/artifacts/artifact.py,sha256=dZj9UkRvvU5EN-HUiuLCPsmosSS0yyHM5JNWzaaIxMA,19887
 truefoundry/ml/log_types/artifacts/constants.py,sha256=uB2JPEqwTbqevkQv2QcEMROsm_4cVAl6s0QU1MLa8SQ,1088
-truefoundry/ml/log_types/artifacts/dataset.py,sha256=UpLXoqhfONqp6YG4N8lDrDe-XhTK6ZZ9Lwg8mI0UZn4,13113
+truefoundry/ml/log_types/artifacts/dataset.py,sha256=OgWIoT59AhMw8P01FfvUKbJ3EL6HQf_Xw8X4E3Ff5Sg,13172
 truefoundry/ml/log_types/artifacts/general_artifact.py,sha256=yr-SQ2fhUR_sE1MB5zoHHYpGC8tizH_-t3lhsxCAULU,2747
-truefoundry/ml/log_types/artifacts/model.py,sha256=0_2EGwj8VYRJ3g8Ti2k75s-OiC1Tu-7Ay9U1_QYG_iw,25027
+truefoundry/ml/log_types/artifacts/model.py,sha256=kAiq7TDJ8RHG1Z4IN9mQD1dzKgmLCP1p0v_Yta9GHlM,24864
 truefoundry/ml/log_types/artifacts/utils.py,sha256=q_atcGzn3wfxItt3RABxjdris8b3njEFNuC8ihWqUSI,8088
 truefoundry/ml/log_types/image/__init__.py,sha256=fcOq8yQnNj1rkLcPeIjLXBpdA1WIeiPsXOlAAvMxx7M,76
 truefoundry/ml/log_types/image/constants.py,sha256=wLtGEOA4T5fZHSlOXPuNDLX3lpbCtwlvGKPFk_1fah0,255
@@ -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.10.7rc1.dist-info/METADATA,sha256=T0oUo-QR8p-afUbJyLte32yJXkdc5w4-iMytXtBv3iI,2508
-truefoundry-0.10.7rc1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-truefoundry-0.10.7rc1.dist-info/entry_points.txt,sha256=xVjn7RMN-MW2-9f7YU-bBdlZSvvrwzhpX1zmmRmsNPU,98
-truefoundry-0.10.7rc1.dist-info/RECORD,,
+truefoundry-0.11.1rc1.dist-info/METADATA,sha256=jIzJR7W9qc5D4a-WuEne8Y8BUotdhjnc_CEcz-Bdp78,2508
+truefoundry-0.11.1rc1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+truefoundry-0.11.1rc1.dist-info/entry_points.txt,sha256=xVjn7RMN-MW2-9f7YU-bBdlZSvvrwzhpX1zmmRmsNPU,98
+truefoundry-0.11.1rc1.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}"
-            )