Skip to content

Commit

Permalink
fix(mlflow.api.transfer): apply prog bar flag, allow extra args, copy…
Browse files Browse the repository at this point in the history
… links by default
  • Loading branch information
lariel-fernandes committed Oct 11, 2024
1 parent 7499f11 commit 759c922
Showing 1 changed file with 12 additions and 1 deletion.
13 changes: 12 additions & 1 deletion src/mlopus/mlflow/api/common/transfer.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import inspect
from datetime import datetime
from pathlib import Path
from typing import List, Any, Iterable
Expand Down Expand Up @@ -26,6 +27,7 @@ class FileTransfer(pydantic.BaseModel):

prog_bar: bool = True
tool: Any = "rclone_python.rclone"
extra_args: dict[str, list[str]] = {"sync": ["--copy-links"]}

@pydantic.root_validator
def _find_tool(self):
Expand Down Expand Up @@ -67,4 +69,13 @@ def push_files(self, src: Path, tgt: urls.Url):
)

def _tool(self, func: str, *args, **kwargs):
return getattr(self.tool, func)(*args, **kwargs)
call = getattr(self.tool, func)

if self.prog_bar and "show_progress" in inspect.signature(call).parameters:
kwargs["show_progress"] = True

return call(
*args,
**kwargs,
args=self.extra_args.get(func) or None,
)

0 comments on commit 759c922

Please sign in to comment.