Repository: mlflow/mlflow

Issue:
Enable `consider-using-set-comprehension` for pylint
Enable the `consider-using-set-comprehension` rule (https://pylint.pycqa.org/en/latest/user_guide/messages/refactor/consider-using-set-comprehension.html) to improve readability and avoid transient list creation.

Instructions:

1. Add `consider-using-set-comprehension` in `pylintrc`:

```diff
diff --git a/pylintrc b/pylintrc
index 24ff2174e..a86adc713 100644
--- a/pylintrc
+++ b/pylintrc
@@ -179,7 +179,8 @@ enable=c-extension-no-member,
        consider-using-from-import,
        consider-merging-isinstance,
        consider-using-dict-items,
-       unnecessary-comprehension
+       unnecessary-comprehension,
+       consider-using-set-comprehension
```

2. Run `pylint $(git ls-files | grep '\.py$')`. The output should look like:

```
************* Module mlflow.pipelines.utils.metrics
mlflow/pipelines/utils/metrics.py:58:26: R1718: Consider using a set comprehension (consider-using-set-comprehension)
mlflow/pipelines/utils/metrics.py:59:27: R1718: Consider using a set comprehension (consider-using-set-comprehension)
************* Module mlflow.store.model_registry.sqlalchemy_store
mlflow/store/model_registry/sqlalchemy_store.py:553:34: R1718: Consider using a set comprehension (consider-using-set-comprehension)
mlflow/store/model_registry/sqlalchemy_store.py:555:34: R1718: Consider using a set comprehension (consider-using-set-comprehension)
...
```

3. Fix the lines listed in the previous step.
4. File a PR.


Discussion:

Resolving patch:
diff --git a/mlflow/pipelines/utils/metrics.py b/mlflow/pipelines/utils/metrics.py
--- a/mlflow/pipelines/utils/metrics.py
+++ b/mlflow/pipelines/utils/metrics.py
@@ -55,8 +55,8 @@ def _get_custom_metrics(step_config: Dict) -> List[Dict]:
     custom_metrics = [
         PipelineMetric.from_custom_metric_dict(metric_dict) for metric_dict in custom_metric_dicts
     ]
-    custom_metric_names = set([metric.name for metric in custom_metrics])
-    builtin_metric_names = set([metric.name for metric in BUILTIN_PIPELINE_METRICS])
+    custom_metric_names = {metric.name for metric in custom_metrics}
+    builtin_metric_names = {metric.name for metric in BUILTIN_PIPELINE_METRICS}
     overridden_builtin_metrics = custom_metric_names.intersection(builtin_metric_names)
     if overridden_builtin_metrics:
         _logger.warning(
diff --git a/mlflow/store/model_registry/sqlalchemy_store.py b/mlflow/store/model_registry/sqlalchemy_store.py
--- a/mlflow/store/model_registry/sqlalchemy_store.py
+++ b/mlflow/store/model_registry/sqlalchemy_store.py
@@ -550,9 +550,9 @@ def get_latest_versions(self, name, stages=None):
             # Convert to RegisteredModel entity first and then extract latest_versions
             latest_versions = sql_registered_model.to_mlflow_entity().latest_versions
             if stages is None or len(stages) == 0:
-                expected_stages = set([get_canonical_stage(stage) for stage in ALL_STAGES])
+                expected_stages = {get_canonical_stage(stage) for stage in ALL_STAGES}
             else:
-                expected_stages = set([get_canonical_stage(stage) for stage in stages])
+                expected_stages = {get_canonical_stage(stage) for stage in stages}
             return [mv for mv in latest_versions if mv.current_stage in expected_stages]
 
     @classmethod


Test patch:
diff --git a/tests/entities/model_registry/test_model_version.py b/tests/entities/model_registry/test_model_version.py
--- a/tests/entities/model_registry/test_model_version.py
+++ b/tests/entities/model_registry/test_model_version.py
@@ -101,8 +101,8 @@ def test_creation_and_hydration(self):
         self.assertEqual(proto.version, "5")
         self.assertEqual(proto.status, ModelVersionStatus.from_string("READY"))
         self.assertEqual(proto.status_message, "Model version #5 is ready to use.")
-        self.assertEqual(set([tag.key for tag in proto.tags]), {"key", "randomKey"})
-        self.assertEqual(set([tag.value for tag in proto.tags]), {"value", "not a random value"})
+        self.assertEqual({tag.key for tag in proto.tags}, {"key", "randomKey"})
+        self.assertEqual({tag.value for tag in proto.tags}, {"value", "not a random value"})
         mvd_2 = ModelVersion.from_proto(proto)
         self._check(
             mvd_2,
diff --git a/tests/entities/model_registry/test_registered_model.py b/tests/entities/model_registry/test_registered_model.py
--- a/tests/entities/model_registry/test_registered_model.py
+++ b/tests/entities/model_registry/test_registered_model.py
@@ -96,17 +96,20 @@ def test_with_latest_model_versions(self):
         proto = rmd_1.to_proto()
         self.assertEqual(proto.creation_timestamp, 1)
         self.assertEqual(proto.last_updated_timestamp, 4000)
-        self.assertEqual(set([mvd.version for mvd in proto.latest_versions]), {"1", "4"})
-        self.assertEqual(set([mvd.name for mvd in proto.latest_versions]), {name})
+        self.assertEqual({mvd.version for mvd in proto.latest_versions}, {"1", "4"})
+        self.assertEqual({mvd.name for mvd in proto.latest_versions}, {name})
         self.assertEqual(
-            set([mvd.current_stage for mvd in proto.latest_versions]),
+            {mvd.current_stage for mvd in proto.latest_versions},
             {"Production", "Staging"},
         )
         self.assertEqual(
-            set([mvd.last_updated_timestamp for mvd in proto.latest_versions]), {2000, 2002}
+            {mvd.last_updated_timestamp for mvd in proto.latest_versions},
+            {2000, 2002},
         )
+
         self.assertEqual(
-            set([mvd.creation_timestamp for mvd in proto.latest_versions]), {1300, 1000}
+            {mvd.creation_timestamp for mvd in proto.latest_versions},
+            {1300, 1000},
         )
 
     def test_with_tags(self):
@@ -128,8 +131,14 @@ def test_with_tags(self):
         proto = rmd_1.to_proto()
         self.assertEqual(proto.creation_timestamp, 1)
         self.assertEqual(proto.last_updated_timestamp, 4000)
-        self.assertEqual(set([tag.key for tag in proto.tags]), {"key", "randomKey"})
-        self.assertEqual(set([tag.value for tag in proto.tags]), {"value", "not a random value"})
+        self.assertEqual(
+            {tag.key for tag in proto.tags},
+            {"key", "randomKey"},
+        )
+        self.assertEqual(
+            {tag.value for tag in proto.tags},
+            {"value", "not a random value"},
+        )
 
     def test_string_repr(self):
         rmd = RegisteredModel(
diff --git a/tests/entities/test_run_data.py b/tests/entities/test_run_data.py
--- a/tests/entities/test_run_data.py
+++ b/tests/entities/test_run_data.py
@@ -7,15 +7,21 @@
 
 class TestRunData(unittest.TestCase):
     def _check_metrics(self, metric_objs, metrics_dict, expected_metrics):
-        self.assertEqual(set([m.key for m in metric_objs]), set([m.key for m in expected_metrics]))
         self.assertEqual(
-            set([m.value for m in metric_objs]), set([m.value for m in expected_metrics])
+            {m.key for m in metric_objs},
+            {m.key for m in expected_metrics},
         )
         self.assertEqual(
-            set([m.timestamp for m in metric_objs]), set([m.timestamp for m in expected_metrics])
+            {m.value for m in metric_objs},
+            {m.value for m in expected_metrics},
         )
         self.assertEqual(
-            set([m.step for m in metric_objs]), set([m.step for m in expected_metrics])
+            {m.timestamp for m in metric_objs},
+            {m.timestamp for m in expected_metrics},
+        )
+        self.assertEqual(
+            {m.step for m in metric_objs},
+            {m.step for m in expected_metrics},
         )
         assert len(metrics_dict) == len(expected_metrics)
         assert metrics_dict == {m.key: m.value for m in expected_metrics}
diff --git a/tests/pipelines/test_pipeline.py b/tests/pipelines/test_pipeline.py
--- a/tests/pipelines/test_pipeline.py
+++ b/tests/pipelines/test_pipeline.py
@@ -146,7 +146,7 @@ def test_pipelines_log_to_expected_mlflow_backend_with_expected_run_tags_once_on
     artifacts = MlflowClient(tracking_uri).list_artifacts(
         run_id=logged_run.info.run_id, path="train"
     )
-    assert set([artifact.path for artifact in artifacts]) == {
+    assert {artifact.path for artifact in artifacts} == {
         "train/card.html",
         "train/estimator",
         "train/model",
diff --git a/tests/store/artifact/test_databricks_artifact_repo.py b/tests/store/artifact/test_databricks_artifact_repo.py
--- a/tests/store/artifact/test_databricks_artifact_repo.py
+++ b/tests/store/artifact/test_databricks_artifact_repo.py
@@ -548,9 +548,9 @@ def test_list_artifacts_handles_pagination(self, databricks_artifact_repo):
             call_endpoint_mock.side_effect = list_artifact_paginated_response_protos
             message_mock.return_value = None
             artifacts = databricks_artifact_repo.list_artifacts()
-            assert set(["a.txt", "b", "c.txt", "d", "e.txt", "f"]) == set(
-                [file.path for file in artifacts]
-            )
+            assert set(["a.txt", "b", "c.txt", "d", "e.txt", "f"]) == {
+                file.path for file in artifacts
+            }
             calls = [
                 mock.call(ListArtifacts(run_id=MOCK_RUN_ID, path="")),
                 mock.call(ListArtifacts(run_id=MOCK_RUN_ID, path="", page_token="2")),
diff --git a/tests/store/tracking/__init__.py b/tests/store/tracking/__init__.py
--- a/tests/store/tracking/__init__.py
+++ b/tests/store/tracking/__init__.py
@@ -61,8 +61,8 @@ def _verify_logged(store, run_id, metrics, params, tags):
         all_metrics = sum([store.get_metric_history(run_id, key) for key in run.data.metrics], [])
         assert len(all_metrics) == len(metrics)
         logged_metrics = [(m.key, m.value, m.timestamp, m.step) for m in all_metrics]
-        assert set(logged_metrics) == set([(m.key, m.value, m.timestamp, m.step) for m in metrics])
+        assert set(logged_metrics) == {(m.key, m.value, m.timestamp, m.step) for m in metrics}
         logged_tags = set(run.data.tags.items())
-        assert set([(tag.key, tag.value) for tag in tags]) <= logged_tags
+        assert {(tag.key, tag.value) for tag in tags} <= logged_tags
         assert len(run.data.params) == len(params)
-        assert set(run.data.params.items()) == set([(param.key, param.value) for param in params])
+        assert set(run.data.params.items()) == {(param.key, param.value) for param in params}
diff --git a/tests/store/tracking/test_sqlalchemy_store.py b/tests/store/tracking/test_sqlalchemy_store.py
--- a/tests/store/tracking/test_sqlalchemy_store.py
+++ b/tests/store/tracking/test_sqlalchemy_store.py
@@ -337,7 +337,7 @@ def test_list_experiments_paginated_last_page(self):
             returned_experiments.extend(result)
         self.assertEqual(result.token, None)
         # make sure that at least all the experiments created in this test are found
-        returned_exp_id_set = set([exp.experiment_id for exp in returned_experiments])
+        returned_exp_id_set = {exp.experiment_id for exp in returned_experiments}
         self.assertEqual(set(experiments) - returned_exp_id_set, set())
 
     def test_list_experiments_paginated_returns_in_correct_order(self):
diff --git a/tests/tensorflow/test_tensorflow2_autolog.py b/tests/tensorflow/test_tensorflow2_autolog.py
--- a/tests/tensorflow/test_tensorflow2_autolog.py
+++ b/tests/tensorflow/test_tensorflow2_autolog.py
@@ -408,7 +408,7 @@ def test_tf_keras_autolog_records_metrics_for_last_epoch(random_train_data, rand
     run_metrics = client.get_run(run.info.run_id).data.metrics
     assert "accuracy" in run_metrics
     all_epoch_acc = client.get_metric_history(run.info.run_id, "accuracy")
-    assert set([metric.step for metric in all_epoch_acc]) == set([0, 5, 10, 15])
+    assert {metric.step for metric in all_epoch_acc} == {0, 5, 10, 15}
 
 
 def test_tf_keras_autolog_logs_metrics_for_single_epoch_training(
@@ -1115,7 +1115,7 @@ def test_tf_estimator_autolog_logs_metrics(tmpdir, export, use_v1_estimator):
     assert "loss" in run.data.metrics
     assert "steps" in run.data.params
     metrics = client.get_metric_history(run_id, "loss")
-    assert set([metric.step for metric in metrics]) == set([1, 6, 11, 16])
+    assert {metric.step for metric in metrics} == {1, 6, 11, 16}
 
 
 @pytest.mark.parametrize("export", [True])
diff --git a/tests/tracking/test_model_registry.py b/tests/tracking/test_model_registry.py
--- a/tests/tracking/test_model_registry.py
+++ b/tests/tracking/test_model_registry.py
@@ -376,8 +376,8 @@ def test_create_and_query_model_version_flow(mlflow_client, backend_store_uri):
         rm.latest_versions for rm in mlflow_client.list_registered_models() if rm.name == name
     ]
     model_versions_by_name = mlflow_client.search_model_versions("name = '%s'" % name)
-    assert {"1", "2"} == set([mv.version for mv in model_versions_by_name])
-    assert {name} == set([mv.name for mv in model_versions_by_name])
+    assert {"1", "2"} == {mv.version for mv in model_versions_by_name}
+    assert {name} == {mv.name for mv in model_versions_by_name}
 
     mv3 = mlflow_client.create_model_version(name, "another_path/to/model", "run_id_2")
     assert mv3.version == "3"
@@ -537,16 +537,16 @@ def test_delete_model_version_flow(mlflow_client, backend_store_uri):
     ]
     assert 1 == len(model_versions_detailed)
     assert "3" == model_versions_detailed[0][0].version
-    assert {"1", "2", "3"} == set(
-        [mv.version for mv in mlflow_client.search_model_versions("name = '%s'" % name)]
-    )
+    assert {"1", "2", "3"} == {
+        mv.version for mv in mlflow_client.search_model_versions("name = '%s'" % name)
+    }
 
     start_time_2 = now()
     mlflow_client.delete_model_version(name, "1")
     end_time_2 = now()
-    assert {"2", "3"} == set(
-        [mv.version for mv in mlflow_client.search_model_versions("name = '%s'" % name)]
-    )
+    assert {"2", "3"} == {
+        mv.version for mv in mlflow_client.search_model_versions("name = '%s'" % name)
+    }
     rmd3 = mlflow_client.get_registered_model(name)
     # deleting model versions changes last_updated_timestamp for registered model
     assert_is_between(start_time_0, end_time_0, rmd3.creation_timestamp)
@@ -563,17 +563,15 @@ def test_delete_model_version_flow(mlflow_client, backend_store_uri):
         mlflow_client.transition_model_version_stage(name=name, version=1, stage="Staging")
 
     mlflow_client.delete_model_version(name, 3)
-    assert {"2"} == set(
-        [mv.version for mv in mlflow_client.search_model_versions("name = '%s'" % name)]
-    )
+    assert {"2"} == {mv.version for mv in mlflow_client.search_model_versions("name = '%s'" % name)}
 
     # new model versions will not reuse existing version numbers
     mv4 = mlflow_client.create_model_version(name, "a/b/c", "run_id_2")
     assert mv4.version == "4"
     assert mv4.name == name
-    assert {"2", "4"} == set(
-        [mv.version for mv in mlflow_client.search_model_versions("name = '%s'" % name)]
-    )
+    assert {"2", "4"} == {
+        mv.version for mv in mlflow_client.search_model_versions("name = '%s'" % name)
+    }
 
 
 def test_set_delete_model_version_tag_flow(mlflow_client, backend_store_uri):
diff --git a/tests/tracking/test_rest_tracking.py b/tests/tracking/test_rest_tracking.py
--- a/tests/tracking/test_rest_tracking.py
+++ b/tests/tracking/test_rest_tracking.py
@@ -88,29 +88,27 @@ def test_create_get_list_experiment(mlflow_client):
     assert exp.tags["key2"] == "val2"
 
     experiments = mlflow_client.list_experiments()
-    assert set([e.name for e in experiments]) == {"My Experiment", "Default"}
+    assert {e.name for e in experiments} == {"My Experiment", "Default"}
     mlflow_client.delete_experiment(experiment_id)
-    assert set([e.name for e in mlflow_client.list_experiments()]) == {"Default"}
-    assert set([e.name for e in mlflow_client.list_experiments(ViewType.ACTIVE_ONLY)]) == {
-        "Default"
-    }
-    assert set([e.name for e in mlflow_client.list_experiments(ViewType.DELETED_ONLY)]) == {
+    assert {e.name for e in mlflow_client.list_experiments()} == {"Default"}
+    assert {e.name for e in mlflow_client.list_experiments(ViewType.ACTIVE_ONLY)} == {"Default"}
+    assert {e.name for e in mlflow_client.list_experiments(ViewType.DELETED_ONLY)} == {
         "My Experiment"
     }
-    assert set([e.name for e in mlflow_client.list_experiments(ViewType.ALL)]) == {
+    assert {e.name for e in mlflow_client.list_experiments(ViewType.ALL)} == {
         "My Experiment",
         "Default",
     }
     active_exps_paginated = mlflow_client.list_experiments(max_results=1)
-    assert set([e.name for e in active_exps_paginated]) == {"Default"}
+    assert {e.name for e in active_exps_paginated} == {"Default"}
     assert active_exps_paginated.token is None
 
     all_exps_paginated = mlflow_client.list_experiments(max_results=1, view_type=ViewType.ALL)
-    first_page_names = set([e.name for e in all_exps_paginated])
+    first_page_names = {e.name for e in all_exps_paginated}
     all_exps_second_page = mlflow_client.list_experiments(
         max_results=1, view_type=ViewType.ALL, page_token=all_exps_paginated.token
     )
-    second_page_names = set([e.name for e in all_exps_second_page])
+    second_page_names = {e.name for e in all_exps_second_page}
     assert len(first_page_names) == 1
     assert len(second_page_names) == 1
     assert first_page_names.union(second_page_names) == {"Default", "My Experiment"}
@@ -672,10 +670,10 @@ def test_artifacts(mlflow_client, tmp_path):
     mlflow_client.log_artifacts(run_id, src_dir, "dir")
 
     root_artifacts_list = mlflow_client.list_artifacts(run_id)
-    assert set([a.path for a in root_artifacts_list]) == {"my.file", "dir"}
+    assert {a.path for a in root_artifacts_list} == {"my.file", "dir"}
 
     dir_artifacts_list = mlflow_client.list_artifacts(run_id, "dir")
-    assert set([a.path for a in dir_artifacts_list]) == {"dir/my.file"}
+    assert {a.path for a in dir_artifacts_list} == {"dir/my.file"}
 
     all_artifacts = mlflow_client.download_artifacts(run_id, ".")
     assert open("%s/my.file" % all_artifacts, "r").read() == "Hello, World!"
diff --git a/tests/tracking/test_tracking.py b/tests/tracking/test_tracking.py
--- a/tests/tracking/test_tracking.py
+++ b/tests/tracking/test_tracking.py
@@ -179,7 +179,7 @@ def test_list_experiments_paginated():
         assert len(result) <= max_results
         returned_experiments.extend(result)
     assert result.token is None
-    returned_exp_id_set = set([exp.experiment_id for exp in returned_experiments])
+    returned_exp_id_set = {exp.experiment_id for exp in returned_experiments}
     assert set(experiments) - returned_exp_id_set == set()
 
 
@@ -842,7 +842,7 @@ def test_search_runs():
         mlflow.set_tag("t2", "second-tag-val")
 
     def verify_runs(runs, expected_set):
-        assert set([r.info.run_id for r in runs]) == set([logged_runs[r] for r in expected_set])
+        assert {r.info.run_id for r in runs} == {logged_runs[r] for r in expected_set}
 
     experiment_id = MlflowClient().get_experiment_by_name("exp-for-search").experiment_id
 
diff --git a/tests/utils/test_search_utils.py b/tests/utils/test_search_utils.py
--- a/tests/utils/test_search_utils.py
+++ b/tests/utils/test_search_utils.py
@@ -298,7 +298,7 @@ def test_correct_filtering(filter_string, matching_runs):
         ),
     ]
     filtered_runs = SearchUtils.filter(runs, filter_string)
-    assert set(filtered_runs) == set([runs[i] for i in matching_runs])
+    assert set(filtered_runs) == {runs[i] for i in matching_runs}
 
 
 def test_filter_runs_by_start_time():