Fix misuse of private method in mlflow.data module (mlflow#11369)

Signed-off-by: chenmoneygithub <chen.qian@databricks.com>

mlflow/data/artifact_dataset_sources.py

@@ -140,7 +140,7 @@ def _can_resolve(raw_source: Any):
         def _resolve(cls, raw_source: Any) -> DatasetForArtifactRepoSourceType:
             return cls(str(raw_source))
 
-        def _to_dict(self) -> Dict[Any, Any]:
+        def to_dict(self) -> Dict[Any, Any]:
             """
             Returns:
                 A JSON-compatible dictionary representation of the {dataset_source_name}.
@@ -150,7 +150,7 @@ def _to_dict(self) -> Dict[Any, Any]:
             }
 
         @classmethod
-        def _from_dict(cls, source_dict: Dict[Any, Any]) -> DatasetForArtifactRepoSourceType:
+        def from_dict(cls, source_dict: Dict[Any, Any]) -> DatasetForArtifactRepoSourceType:
             uri = source_dict.get("uri")
             if uri is None:
                 raise MlflowException(
@@ -163,7 +163,7 @@ def _from_dict(cls, source_dict: Dict[Any, Any]) -> DatasetForArtifactRepoSource
     ArtifactRepoSource.__name__ = dataset_source_name
     ArtifactRepoSource.__qualname__ = dataset_source_name
     ArtifactRepoSource.__doc__ = class_docstring
-    ArtifactRepoSource._to_dict.__doc__ = ArtifactRepoSource._to_dict.__doc__.format(
+    ArtifactRepoSource.to_dict.__doc__ = ArtifactRepoSource.to_dict.__doc__.format(
         dataset_source_name=dataset_source_name
     )
     ArtifactRepoSource.uri.__doc__ = ArtifactRepoSource.uri.__doc__.format(scheme=scheme)

mlflow/data/code_dataset_source.py

@@ -28,11 +28,11 @@ def _can_resolve(raw_source: Any):
     def _resolve(cls, raw_source: str) -> "CodeDatasetSource":
         raise NotImplementedError
 
-    def _to_dict(self) -> Dict[Any, Any]:
+    def to_dict(self) -> Dict[Any, Any]:
         return {"tags": self._tags}
 
     @classmethod
-    def _from_dict(cls, source_dict: Dict[Any, Any]) -> "CodeDatasetSource":
+    def from_dict(cls, source_dict: Dict[Any, Any]) -> "CodeDatasetSource":
         return cls(
             tags=source_dict.get("tags"),
         )

mlflow/data/dataset_source.py

@@ -64,7 +64,7 @@ def _resolve(cls, raw_source: Any) -> "DatasetSource":
         """
 
     @abstractmethod
-    def _to_dict(self) -> Dict[str, Any]:
+    def to_dict(self) -> Dict[str, Any]:
         """Obtains a JSON-compatible dictionary representation of the DatasetSource.
 
         Returns:
@@ -81,11 +81,11 @@ def to_json(self) -> str:
             A JSON string representation of the
             :py:class:`DatasetSource <mlflow.data.dataset_source.DatasetSource>`.
         """
-        return json.dumps(self._to_dict())
+        return json.dumps(self.to_dict())
 
     @classmethod
     @abstractmethod
-    def _from_dict(cls, source_dict: Dict[Any, Any]) -> "DatasetSource":
+    def from_dict(cls, source_dict: Dict[Any, Any]) -> "DatasetSource":
         """Constructs an instance of the DatasetSource from a dictionary representation.
 
         Args:
@@ -107,4 +107,4 @@ def from_json(cls, source_json: str) -> "DatasetSource":
             A DatasetSource instance.
 
         """
-        return cls._from_dict(json.loads(source_json))
+        return cls.from_dict(json.loads(source_json))

mlflow/data/delta_dataset_source.py

@@ -137,7 +137,7 @@ def _lookup_table_id(self, table_name):
         except Exception:
             return None
 
-    def _to_dict(self) -> Dict[Any, Any]:
+    def to_dict(self) -> Dict[Any, Any]:
         info = {}
         if self._path:
             info["path"] = self._path
@@ -154,7 +154,7 @@ def _to_dict(self) -> Dict[Any, Any]:
         return info
 
     @classmethod
-    def _from_dict(cls, source_dict: Dict[Any, Any]) -> "DeltaDatasetSource":
+    def from_dict(cls, source_dict: Dict[Any, Any]) -> "DeltaDatasetSource":
         return cls(
             path=source_dict.get("path"),
             delta_table_name=source_dict.get("delta_table_name"),

mlflow/data/filesystem_dataset_source.py

@@ -66,15 +66,15 @@ def _resolve(cls, raw_source: Any) -> "FileSystemDatasetSource":
         """
 
     @abstractmethod
-    def _to_dict(self) -> Dict[Any, Any]:
+    def to_dict(self) -> Dict[Any, Any]:
         """
         Returns:
             A JSON-compatible dictionary representation of the FileSystemDatasetSource.
         """
 
     @classmethod
     @abstractmethod
-    def _from_dict(cls, source_dict: Dict[Any, Any]) -> "FileSystemDatasetSource":
+    def from_dict(cls, source_dict: Dict[Any, Any]) -> "FileSystemDatasetSource":
         """
         Args:
             source_dict: A dictionary representation of the FileSystemDatasetSource.

mlflow/data/http_dataset_source.py

@@ -120,7 +120,7 @@ def _resolve(cls, raw_source: Any) -> "HTTPDatasetSource":
         """
         return HTTPDatasetSource(raw_source)
 
-    def _to_dict(self) -> Dict[Any, Any]:
+    def to_dict(self) -> Dict[Any, Any]:
         """
         Returns:
             A JSON-compatible dictionary representation of the HTTPDatasetSource.
@@ -130,7 +130,7 @@ def _to_dict(self) -> Dict[Any, Any]:
         }
 
     @classmethod
-    def _from_dict(cls, source_dict: Dict[Any, Any]) -> "HTTPDatasetSource":
+    def from_dict(cls, source_dict: Dict[Any, Any]) -> "HTTPDatasetSource":
         """
         Args:
             source_dict: A dictionary representation of the HTTPDatasetSource.

mlflow/data/huggingface_dataset_source.py

@@ -87,7 +87,7 @@ def _can_resolve(raw_source: Any):
     def _resolve(cls, raw_source: str) -> "HuggingFaceDatasetSource":
         raise NotImplementedError
 
-    def _to_dict(self) -> Dict[Any, Any]:
+    def to_dict(self) -> Dict[Any, Any]:
         return {
             "path": self.path,
             "config_name": self.config_name,
@@ -98,7 +98,7 @@ def _to_dict(self) -> Dict[Any, Any]:
         }
 
     @classmethod
-    def _from_dict(cls, source_dict: Dict[Any, Any]) -> "HuggingFaceDatasetSource":
+    def from_dict(cls, source_dict: Dict[Any, Any]) -> "HuggingFaceDatasetSource":
         return cls(
             path=source_dict.get("path"),
             config_name=source_dict.get("config_name"),

mlflow/data/spark_dataset_source.py

@@ -55,7 +55,7 @@ def _can_resolve(raw_source: Any):
     def _resolve(cls, raw_source: str) -> "SparkDatasetSource":
         raise NotImplementedError
 
-    def _to_dict(self) -> Dict[Any, Any]:
+    def to_dict(self) -> Dict[Any, Any]:
         info = {}
         if self._path is not None:
             info["path"] = self._path
@@ -66,7 +66,7 @@ def _to_dict(self) -> Dict[Any, Any]:
         return info
 
     @classmethod
-    def _from_dict(cls, source_dict: Dict[Any, Any]) -> "SparkDatasetSource":
+    def from_dict(cls, source_dict: Dict[Any, Any]) -> "SparkDatasetSource":
         return cls(
             path=source_dict.get("path"),
             table_name=source_dict.get("table_name"),

mlflow/data/uc_volume_dataset_source.py

@@ -65,9 +65,9 @@ def _can_resolve(raw_source: Any):
     def _resolve(cls, raw_source: str):
         raise NotImplementedError
 
-    def _to_dict(self) -> Dict[Any, Any]:
+    def to_dict(self) -> Dict[Any, Any]:
         return {"path": self.path}
 
     @classmethod
-    def _from_dict(cls, source_dict: Dict[Any, Any]) -> "UCVolumeDatasetSource":
+    def from_dict(cls, source_dict: Dict[Any, Any]) -> "UCVolumeDatasetSource":
         return cls(**source_dict)

tests/data/test_code_dataset_source.py

@@ -7,7 +7,7 @@ def test_code_dataset_source_from_path():
         "mlflow_source_name": "some_random_notebook_path",
     }
     code_datasource = CodeDatasetSource(tags)
-    assert code_datasource._to_dict() == {
+    assert code_datasource.to_dict() == {
         "tags": tags,
     }
 

tests/resources/data/dataset_source.py

@@ -39,11 +39,11 @@ def _can_resolve(raw_source: Any) -> bool:
     def _resolve(cls, raw_source: Any) -> DatasetSource:
         return cls(raw_source)
 
-    def _to_dict(self) -> Dict[Any, Any]:
+    def to_dict(self) -> Dict[Any, Any]:
         return {"uri": self.uri}
 
     @classmethod
-    def _from_dict(cls, source_dict: Dict[Any, Any]) -> DatasetSource:
+    def from_dict(cls, source_dict: Dict[Any, Any]) -> DatasetSource:
         uri = source_dict.get("uri")
         if uri is None:
             raise MlflowException(