chore: import upstream snapshot with attribution

python/ray/train/tests/test_train_usage.py

@@ -0,0 +1,148 @@
+import pytest
+import torch
+
+import ray
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+
+
+@pytest.fixture
+def shutdown_only():
+    yield None
+    ray.shutdown()
+
+
+def run_torch():
+    from torch.utils.data import DataLoader, TensorDataset
+
+    from ray.train.torch import (
+        get_device,
+        get_devices,
+        prepare_data_loader,
+        prepare_model,
+    )
+
+    def train_func():
+        # Create dummy model and data loader
+        model = torch.nn.Linear(10, 10)
+        inputs, targets = torch.randn(128, 10), torch.randn(128, 1)
+        dataloader = DataLoader(TensorDataset(inputs, targets), batch_size=32)
+
+        # Test Torch Utilities
+        prepare_data_loader(dataloader)
+        prepare_model(model)
+        get_device()
+        get_devices()
+
+    trainer = TorchTrainer(
+        train_func, scaling_config=ScalingConfig(num_workers=2, use_gpu=False)
+    )
+    trainer.fit()
+
+
+def run_lightning():
+    import lightning.pytorch as pl
+
+    from ray.train.lightning import (
+        RayDDPStrategy,
+        RayDeepSpeedStrategy,
+        RayFSDPStrategy,
+        RayLightningEnvironment,
+        RayTrainReportCallback,
+        prepare_trainer,
+    )
+
+    def train_func():
+        # Test Lighting utilites
+        strategy = RayFSDPStrategy()
+        strategy = RayDeepSpeedStrategy()
+        strategy = RayDDPStrategy()
+        ray_environment = RayLightningEnvironment()
+        report_callback = RayTrainReportCallback()
+
+        trainer = pl.Trainer(
+            devices="auto",
+            accelerator="auto",
+            strategy=strategy,
+            plugins=[ray_environment],
+            callbacks=[report_callback],
+        )
+        trainer = prepare_trainer(trainer)
+
+    trainer = TorchTrainer(
+        train_func, scaling_config=ScalingConfig(num_workers=2, use_gpu=False)
+    )
+
+    trainer.fit()
+
+
+def run_transformers():
+    from datasets import Dataset
+    from transformers import Trainer, TrainingArguments
+
+    from ray.train.huggingface.transformers import (
+        RayTrainReportCallback,
+        prepare_trainer,
+    )
+
+    def train_func():
+        # Create dummy model and datasets
+        dataset = Dataset.from_dict({"text": ["text1", "text2"], "label": [0, 1]})
+        model = torch.nn.Linear(10, 10)
+
+        # Test Transformers utilites
+        training_args = TrainingArguments(output_dir="./results", use_cpu=True)
+        trainer = Trainer(model=model, args=training_args, train_dataset=dataset)
+
+        trainer.add_callback(RayTrainReportCallback())
+        trainer = prepare_trainer(trainer)
+
+    trainer = TorchTrainer(
+        train_func, scaling_config=ScalingConfig(num_workers=2, use_gpu=False)
+    )
+
+    trainer.fit()
+
+
+@pytest.mark.parametrize("framework", ["torch", "lightning", "transformers"])
+def test_torch_utility_usage_tags(shutdown_only, framework):
+    from ray._common.usage.usage_lib import TagKey, get_extra_usage_tags_to_report
+
+    ctx = ray.init()
+    gcs_client = ray._raylet.GcsClient(address=ctx.address_info["gcs_address"])
+
+    if framework == "torch":
+        run_torch()
+        expected_tags = [
+            TagKey.TRAIN_TORCH_GET_DEVICE,
+            TagKey.TRAIN_TORCH_GET_DEVICES,
+            TagKey.TRAIN_TORCH_PREPARE_MODEL,
+            TagKey.TRAIN_TORCH_PREPARE_DATALOADER,
+        ]
+    elif framework == "lightning":
+        run_lightning()
+        expected_tags = [
+            TagKey.TRAIN_LIGHTNING_PREPARE_TRAINER,
+            TagKey.TRAIN_LIGHTNING_RAYTRAINREPORTCALLBACK,
+            TagKey.TRAIN_LIGHTNING_RAYDDPSTRATEGY,
+            TagKey.TRAIN_LIGHTNING_RAYFSDPSTRATEGY,
+            TagKey.TRAIN_LIGHTNING_RAYDEEPSPEEDSTRATEGY,
+            TagKey.TRAIN_LIGHTNING_RAYLIGHTNINGENVIRONMENT,
+        ]
+    elif framework == "transformers":
+        run_transformers()
+        expected_tags = [
+            TagKey.TRAIN_TRANSFORMERS_PREPARE_TRAINER,
+            TagKey.TRAIN_TRANSFORMERS_RAYTRAINREPORTCALLBACK,
+        ]
+
+    result = get_extra_usage_tags_to_report(gcs_client)
+    assert set(result.keys()).issuperset(
+        {TagKey.Name(tag).lower() for tag in expected_tags}
+    )
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))