chore: import upstream snapshot with attribution

python/ray/train/examples/experiment_tracking/lightning_exp_tracking_comet.py

@@ -0,0 +1,41 @@
+# isort: skip_file
+from lightning_exp_tracking_model_dl import DummyModel, dataloader
+
+# __lightning_experiment_tracking_comet_start__
+import os
+import ray
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+
+import lightning.pytorch as pl
+from lightning.pytorch.loggers import CometLogger
+
+
+def train_func(config):
+    logger = None
+    if ray.train.get_context().get_world_rank() == 0:
+        logger = CometLogger(api_key=os.environ["COMET_API_KEY"])
+
+    ptl_trainer = pl.Trainer(
+        max_epochs=5,
+        accelerator="cpu",
+        logger=logger,
+        log_every_n_steps=1,
+    )
+    model = DummyModel()
+    ptl_trainer.fit(model, train_dataloaders=dataloader)
+
+
+scaling_config = ScalingConfig(num_workers=2, use_gpu=False)
+
+assert (
+    "COMET_API_KEY" in os.environ
+), 'Please do COMET_API_KEY="abcde" when running this script.'
+# This makes sure that all workers have this env var set.
+ray.init(runtime_env={"env_vars": {"COMET_API_KEY": os.environ["COMET_API_KEY"]}})
+trainer = TorchTrainer(
+    train_func,
+    scaling_config=scaling_config,
+)
+
+trainer.fit()

python/ray/train/examples/experiment_tracking/lightning_exp_tracking_mlflow.py

@@ -0,0 +1,63 @@
+# ruff: noqa
+# isort: skip_file
+import os
+import tempfile
+
+tempdir = tempfile.TemporaryDirectory()
+os.environ["SHARED_STORAGE_PATH"] = tempdir.name
+
+from ray.train.examples.experiment_tracking.lightning_exp_tracking_model_dl import (
+    DummyModel,
+    dataloader,
+)
+
+
+# __lightning_experiment_tracking_mlflow_start__
+import os
+
+import ray
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+
+import lightning.pytorch as pl
+from lightning.pytorch.loggers import MLFlowLogger
+
+
+def train_func(config):
+
+    save_dir = config["save_dir"]
+    logger = None
+    if ray.train.get_context().get_world_rank() == 0:
+        logger = MLFlowLogger(
+            experiment_name="demo-project",
+            tracking_uri=f"file:{save_dir}",
+        )
+
+    ptl_trainer = pl.Trainer(
+        max_epochs=5,
+        accelerator="cpu",
+        logger=logger,
+        log_every_n_steps=1,
+    )
+    model = DummyModel()
+    ptl_trainer.fit(model, train_dataloaders=dataloader)
+
+
+scaling_config = ScalingConfig(num_workers=2, use_gpu=False)
+
+assert (
+    "SHARED_STORAGE_PATH" in os.environ
+), "Please do SHARED_STORAGE_PATH=/a/b/c when running this script."
+
+trainer = TorchTrainer(
+    train_func,
+    train_loop_config={
+        "save_dir": os.path.join(os.environ["SHARED_STORAGE_PATH"], "mlruns")
+    },
+    scaling_config=scaling_config,
+)
+
+trainer.fit()
+# __lightning_experiment_tracking_mlflow_end__
+
+tempdir.cleanup()

python/ray/train/examples/experiment_tracking/lightning_exp_tracking_model_dl.py

@@ -0,0 +1,46 @@
+# ruff: noqa
+# fmt: off
+# # isort: skip_file
+
+
+# __model_dl_start__
+import lightning.pytorch as pl
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, TensorDataset
+
+# Create dummy data
+X = torch.randn(128, 3)  # 128 samples, 3 features
+y = torch.randint(0, 2, (128,))  # 128 binary labels
+
+# Create a TensorDataset to wrap the data
+dataset = TensorDataset(X, y)
+
+# Create a DataLoader to iterate over the dataset
+batch_size = 8
+dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+
+
+# Define a dummy model
+class DummyModel(pl.LightningModule):
+    def __init__(self):
+        super().__init__()
+        self.layer = torch.nn.Linear(3, 1)
+
+    def forward(self, x):
+        return self.layer(x)
+
+    def training_step(self, batch, batch_idx):
+        x, y = batch
+        y_hat = self(x)
+        loss = F.binary_cross_entropy_with_logits(y_hat.flatten(), y.float())
+
+        # The metrics below will be reported to Loggers
+        self.log("train_loss", loss)
+        self.log_dict({
+            "metric_1": 1 / (batch_idx + 1), "metric_2": batch_idx * 100
+        })
+        return loss
+
+    def configure_optimizers(self):
+        return torch.optim.Adam(self.parameters(), lr=1e-3)

python/ray/train/examples/experiment_tracking/lightning_exp_tracking_tensorboard.py

@@ -0,0 +1,60 @@
+# ruff: noqa
+# isort: skip_file
+
+import os
+import tempfile
+
+tempdir = tempfile.TemporaryDirectory()
+os.environ["SHARED_STORAGE_PATH"] = tempdir.name
+
+from ray.train.examples.experiment_tracking.lightning_exp_tracking_model_dl import (
+    DummyModel,
+    dataloader,
+)
+
+# __lightning_experiment_tracking_tensorboard_start__
+import os
+
+import ray
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+
+import lightning.pytorch as pl
+from lightning.pytorch.loggers import TensorBoardLogger
+
+
+def train_func(config):
+
+    save_dir = config["save_dir"]
+    logger = None
+    if ray.train.get_context().get_world_rank() == 0:
+        logger = TensorBoardLogger(name="demo-run", save_dir=f"file:{save_dir}")
+
+    ptl_trainer = pl.Trainer(
+        max_epochs=5,
+        accelerator="cpu",
+        logger=logger,
+        log_every_n_steps=1,
+    )
+    model = DummyModel()
+    ptl_trainer.fit(model, train_dataloaders=dataloader)
+
+
+scaling_config = ScalingConfig(num_workers=2, use_gpu=False)
+
+assert (
+    "SHARED_STORAGE_PATH" in os.environ
+), "Please do SHARED_STORAGE_PATH=/a/b/c when running this script."
+
+trainer = TorchTrainer(
+    train_func,
+    train_loop_config={
+        "save_dir": os.path.join(os.environ["SHARED_STORAGE_PATH"], "tensorboard")
+    },
+    scaling_config=scaling_config,
+)
+
+trainer.fit()
+# __lightning_experiment_tracking_tensorboard_end__
+
+tempdir.cleanup()

python/ray/train/examples/experiment_tracking/lightning_exp_tracking_wandb.py

@@ -0,0 +1,50 @@
+# ruff: noqa
+# fmt: off
+# # isort: skip_file
+
+from lightning_exp_tracking_model_dl import DummyModel, dataloader
+
+# __lightning_experiment_tracking_wandb_start__
+import os
+import wandb
+import ray
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+
+import lightning.pytorch as pl
+from lightning.pytorch.loggers import WandbLogger
+
+
+def train_func(config):
+    logger = None
+    if ray.train.get_context().get_world_rank() == 0:
+        logger = WandbLogger(name="demo-run", project="demo-project")
+
+    ptl_trainer = pl.Trainer(
+        max_epochs=5,
+        accelerator="cpu",
+        logger=logger,
+        log_every_n_steps=1,
+    )
+    model = DummyModel()
+    ptl_trainer.fit(model, train_dataloaders=dataloader)
+    if ray.train.get_context().get_world_rank() == 0:
+        wandb.finish()
+
+
+scaling_config = ScalingConfig(num_workers=2, use_gpu=False)
+
+assert (
+    "WANDB_API_KEY" in os.environ
+), 'Please set WANDB_API_KEY="abcde" when running this script.'
+
+# This ensures that all workers have this env var set.
+ray.init(
+    runtime_env={"env_vars": {"WANDB_API_KEY": os.environ["WANDB_API_KEY"]}}
+)
+trainer = TorchTrainer(
+    train_func,
+    scaling_config=scaling_config,
+)
+
+trainer.fit()

python/ray/train/examples/experiment_tracking/torch_exp_tracking_mlflow.py

@@ -0,0 +1,85 @@
+# ruff: noqa
+# isort: skip_file
+from filelock import FileLock
+import os
+import tempfile
+
+tempdir = tempfile.TemporaryDirectory()
+os.environ["SHARED_STORAGE_PATH"] = tempdir.name
+
+# __start__
+# Run the following script with the SHARED_STORAGE_PATH env var set.
+# The MLflow offline logs are saved to SHARED_STORAGE_PATH/mlruns.
+
+import mlflow
+import os
+import ray
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+import torch
+from torchvision import datasets, transforms
+from torchvision.models import resnet18
+from torch.utils.data import DataLoader
+
+assert os.environ.get(
+    "SHARED_STORAGE_PATH", None
+), "Please set SHARED_STORAGE_PATH env var."
+
+
+# Assumes you are passing a `save_dir` in `config`
+def train_func(config):
+    save_dir = config["save_dir"]
+    if ray.train.get_context().get_world_rank() == 0:
+        mlflow.set_tracking_uri(f"file:{save_dir}")
+        mlflow.set_experiment("my_experiment")
+        mlflow.start_run()
+
+    # Model, Loss, Optimizer
+    model = resnet18(num_classes=10)
+    model.conv1 = torch.nn.Conv2d(
+        1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
+    )
+    model = ray.train.torch.prepare_model(model)
+    criterion = torch.nn.CrossEntropyLoss()
+    optimizer = torch.optim.Adam(model.module.parameters(), lr=0.001)
+
+    # Data
+    transform = transforms.Compose(
+        [transforms.ToTensor(), transforms.Normalize((0.28604,), (0.32025,))]
+    )
+    with FileLock("./data.lock"):
+        train_data = datasets.FashionMNIST(
+            root="./data", train=True, download=True, transform=transform
+        )
+    train_loader = DataLoader(train_data, batch_size=128, shuffle=True)
+    train_loader = ray.train.torch.prepare_data_loader(train_loader)
+
+    # Training
+    for epoch in range(1):
+        if ray.train.get_context().get_world_size() > 1:
+            train_loader.sampler.set_epoch(epoch)
+
+        for images, labels in train_loader:
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            if ray.train.get_context().get_world_rank() == 0:
+                mlflow.log_metrics({"loss": loss.item(), "epoch": epoch})
+
+    if ray.train.get_context().get_world_rank() == 0:
+        mlflow.end_run()
+
+
+trainer = TorchTrainer(
+    train_func,
+    train_loop_config={
+        "save_dir": os.path.join(os.environ["SHARED_STORAGE_PATH"], "mlruns")
+    },
+    scaling_config=ScalingConfig(num_workers=2),
+)
+trainer.fit()
+# __end__
+
+tempdir.cleanup()

python/ray/train/examples/experiment_tracking/torch_exp_tracking_wandb.py

@@ -0,0 +1,75 @@
+# ruff: noqa
+# fmt: off
+# isort: off
+
+# __start__
+from filelock import FileLock
+import os
+
+import torch
+import wandb
+
+from torch.utils.data import DataLoader
+from torchvision import datasets, transforms
+from torchvision.models import resnet18
+
+import ray
+from ray.train import ScalingConfig
+from ray.train.torch import TorchTrainer
+
+# Run the following script with the WANDB_API_KEY env var set.
+assert os.environ.get("WANDB_API_KEY", None), "Please set WANDB_API_KEY env var."
+
+# This makes sure that all workers have this env var set.
+ray.init(
+    runtime_env={"env_vars": {"WANDB_API_KEY": os.environ["WANDB_API_KEY"]}}
+)
+
+
+def train_func(config):
+    if ray.train.get_context().get_world_rank() == 0:
+        wandb.init()
+
+    # Model, Loss, Optimizer
+    model = resnet18(num_classes=10)
+    model.conv1 = torch.nn.Conv2d(
+        1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
+    )
+    model = ray.train.torch.prepare_model(model)
+    criterion = torch.nn.CrossEntropyLoss()
+    optimizer = torch.optim.Adam(model.module.parameters(), lr=0.001)
+
+    # Data
+    transform = transforms.Compose(
+        [transforms.ToTensor(), transforms.Normalize((0.28604,), (0.32025,))]
+    )
+    with FileLock("./data.lock"):
+        train_data = datasets.FashionMNIST(
+            root="./data", train=True, download=True, transform=transform
+        )
+    train_loader = DataLoader(train_data, batch_size=128, shuffle=True)
+    train_loader = ray.train.torch.prepare_data_loader(train_loader)
+
+    # Training
+    for epoch in range(1):
+        if ray.train.get_context().get_world_size() > 1:
+            train_loader.sampler.set_epoch(epoch)
+
+        for images, labels in train_loader:
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            if ray.train.get_context().get_world_rank() == 0:
+                wandb.log({"loss": loss, "epoch": epoch})
+
+    if ray.train.get_context().get_world_rank() == 0:
+        wandb.finish()
+
+
+trainer = TorchTrainer(
+    train_func,
+    scaling_config=ScalingConfig(num_workers=2),
+)
+trainer.fit()