[feat] Add --exclude-artifacts option to exclude artifacts based on glob expression

aimlflow/cli.py

@@ -25,18 +25,19 @@ def cli_entry_point():
                                                             writable=True))
 @click.option('--mlflow-tracking-uri',  required=False, default=None)
 @click.option('--experiment', '-e', required=False, default=None)
-def sync(aim_repo, mlflow_tracking_uri, experiment):
+@click.option('--excluded-artifacts', required=False, default=None)
+def sync(aim_repo, mlflow_tracking_uri, experiment, excluded_artifacts):
     repo_path = clean_repo_path(aim_repo) or Repo.default_repo_path()
     repo_inst = Repo.from_path(repo_path)
 
     mlflow_tracking_uri = mlflow_tracking_uri or os.environ.get('MLFLOW_TRACKING_URI')
     if not mlflow_tracking_uri:
         raise ClickException('MLFlow tracking URI must be provided either through ENV or CLI.')
 
-    watcher = MLFlowWatcher(repo_inst, mlflow_tracking_uri, experiment)
+    watcher = MLFlowWatcher(repo_inst, mlflow_tracking_uri, experiment, excluded_artifacts)
 
     click.echo('Converting existing MLflow logs.')
-    convert_existing_logs(repo_inst, mlflow_tracking_uri, experiment)
+    convert_existing_logs(repo_inst, mlflow_tracking_uri, experiment, excluded_artifacts)
 
     click.echo(f'Starting watcher on {mlflow_tracking_uri}.')
     watcher.start()

aimlflow/utils.py

@@ -1,3 +1,4 @@
+import fnmatch
 import click
 import mlflow
 import json
@@ -114,7 +115,10 @@ def collect_run_params(aim_run, mlflow_run):
     }
 
 
-def collect_artifacts(aim_run, mlflow_run, mlflow_client):
+def collect_artifacts(aim_run, mlflow_run, mlflow_client, excluded_artifacts):
+    if excluded_artifacts == '*':
+        return
+
     run_id = mlflow_run.info.run_id
 
     artifacts_cache_key = '_mlflow_artifacts_cache'
@@ -136,6 +140,10 @@ def collect_artifacts(aim_run, mlflow_run, mlflow_client):
                     continue
                 else:
                     artifacts_cache.append(file_info.path)
+
+                if fnmatch.fnmatch(file_info.path, excluded_artifacts):
+                    continue
+
                 downloaded_path = mlflow_client.download_artifacts(run_id, file_info.path, dst_path=temp_path)
                 if file_info.path.endswith(HTML_EXTENSIONS):
                     if not __html_warning_issued:
@@ -196,7 +204,7 @@ def collect_metrics(aim_run, mlflow_run, mlflow_client, timestamp=None):
             aim_run.track(m.value, step=m.step, name=m.key)
 
 
-def convert_existing_logs(repo_inst, tracking_uri, experiment=None, no_cache=False):
+def convert_existing_logs(repo_inst, tracking_uri, experiment=None, excluded_artifacts=None, no_cache=False):
     client = mlflow.tracking.client.MlflowClient(tracking_uri=tracking_uri)
 
     experiments = get_mlflow_experiments(client, experiment)
@@ -215,7 +223,7 @@ def convert_existing_logs(repo_inst, tracking_uri, experiment=None, no_cache=Fal
             collect_metrics(aim_run, run, client)
 
             # Collect artifacts
-            collect_artifacts(aim_run, run, client)
+            collect_artifacts(aim_run, run, client, excluded_artifacts)
 
     run_cache.refresh()
 

aimlflow/watcher.py

@@ -28,6 +28,7 @@ def __init__(self,
                  repo: 'Repo',
                  tracking_uri: str,
                  experiment: str = None,
+                 excluded_artifacts: str = None,
                  interval: Union[int, float] = WATCH_INTERVAL_DEFAULT,
                  ):
 
@@ -38,6 +39,7 @@ def __init__(self,
 
         self._client = MlflowClient(tracking_uri)
         self._experiments = get_mlflow_experiments(self._client, experiment)
+        self._excluded_artifacts = excluded_artifacts
         self._repo = repo
 
         self._th_collector = Thread(target=self._watch, daemon=True)
@@ -77,7 +79,7 @@ def _process_single_run(self, aim_run, mlflow_run):
         collect_metrics(aim_run, mlflow_run, self._client, timestamp=self._last_watch_time)
 
         # Collect artifacts
-        collect_artifacts(aim_run, mlflow_run, self._client)
+        collect_artifacts(aim_run, mlflow_run, self._client, self._excluded_artifacts)
 
     def _process_runs(self):
         watch_started_time = time.time()