Support for different OCI object store authentication mechanisms

[psarka]

🚀 Feature Request

To my knowledge, OCI has 4 different auth mechanisms (“api_key”, “resource_principal”, “instance_principal”, and "oke_principal" in ocifs lingo). Mosaicml streaming hardcodes "api_key" in the OCIUploaded and OCIDownloader classes. Please add support for the other 3, perhaps by providing an oci client dependency injection path. 🤗

Motivation

Training jobs and data preparation jobs run on the k8s clusters that use "oke_principal" or "instance_principal" depending on the setup.

[Optional] Implementation

I only have the following monkey patch that serves as a workaround for the moment:

def patch_oci_auth(client_factory: Callable) -> None:
    """Patch MDS OCIUploader and OCIDownloader to use a custom OCI client instead of oci.config.from_file()."""
    import oci

    def _patched_uploader_init(self, out, keep_local=False, progress_bar=False, retry=2, exist_ok=False):
        CloudUploader.__init__(self, out, keep_local, progress_bar, retry, exist_ok)
        self.client = client_factory()
        self.namespace = self.client.get_namespace().data
        self.upload_manager = oci.object_storage.UploadManager(self.client)
        self.check_bucket_exists(self.remote)

    def _patched_create_oci_client(self):
        self._oci_client = client_factory()

    OCIUploader.__init__ = _patched_uploader_init
    OCIDownloader._create_oci_client = _patched_create_oci_client

Additional context