Merge pull request #210 from comet-ml/update-ray-train-new-api

Update Ray Train examples for new public API

Author: Lothiraldan

.github/workflows/test-examples.yml

@@ -92,10 +92,9 @@ jobs:
           COMET_INTERNAL_SENTRY_DSN: ${{ secrets.COMET_INTERNAL_SENTRY_DSN }}
           COMET_WORKSPACE: cometexamples-tests
       - name: debugging-save-logs
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         if:  runner.debug == '1' && failure()
         with:
-          name: debug-logs
           path: ${{ env.COMET_LOG_DIR }}
 
   test-scripts:

integrations/model-training/ray-train/notebooks/Comet_with_ray_train_huggingface_transformers.ipynb

@@ -41,7 +41,7 @@
    },
    "outputs": [],
    "source": [
-    "%pip install \"comet_ml>=3.31.5\" \"ray[air]>=2.1.0\" \"transformers>=4.43.0\" \"accelerate>=0.12.0\" \"datasets\" \"sentencepiece\" scipy \"scikit-learn\" protobuf \"torch>=1.3\" evaluate"
+    "%pip install \"comet_ml>=3.49.0\" \"ray[air]>=2.1.0\" \"transformers>=4.43.0\" \"accelerate>=0.12.0\" \"datasets\" \"sentencepiece\" scipy \"scikit-learn\" protobuf \"torch>=1.3\" evaluate"
    ]
   },
   {
@@ -62,7 +62,6 @@
    "outputs": [],
    "source": [
     "import comet_ml\n",
-    "import comet_ml.integration.ray\n",
     "\n",
     "comet_ml.init()"
    ]
@@ -101,7 +100,9 @@
     "\n",
     "import ray.train.huggingface.transformers\n",
     "from ray.train import ScalingConfig, RunConfig\n",
-    "from ray.train.torch import TorchTrainer"
+    "from ray.train.torch import TorchTrainer\n",
+    "import comet_ml.integration.ray\n",
+    "from comet_ml.integration.ray import comet_worker"
    ]
   },
   {
@@ -164,63 +165,62 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "@comet_worker\n",
     "def train_func(config):\n",
     "    from comet_ml import get_running_experiment\n",
-    "    from comet_ml.integration.ray import comet_worker_logger\n",
-    "\n",
-    "    with comet_worker_logger(config) as experiment:\n",
-    "        small_train_dataset, small_eval_dataset = get_dataset()\n",
-    "\n",
-    "        # Model\n",
-    "        model = AutoModelForSequenceClassification.from_pretrained(\n",
-    "            \"google-bert/bert-base-cased\", num_labels=5\n",
-    "        )\n",
-    "\n",
-    "        # Evaluation Metrics\n",
-    "        metric = evaluate.load(\"accuracy\")\n",
-    "\n",
-    "        def compute_metrics(eval_pred):\n",
-    "            logits, labels = eval_pred\n",
-    "            predictions = np.argmax(logits, axis=-1)\n",
-    "\n",
-    "            experiment = comet_ml.get_running_experiment()\n",
-    "            if experiment:\n",
-    "                experiment.log_confusion_matrix(predictions, labels)\n",
-    "\n",
-    "            return metric.compute(predictions=predictions, references=labels)\n",
-    "\n",
-    "        # Hugging Face Trainer\n",
-    "        training_args = TrainingArguments(\n",
-    "            do_eval=True,\n",
-    "            do_train=True,\n",
-    "            eval_strategy=\"epoch\",\n",
-    "            num_train_epochs=config[\"epochs\"],\n",
-    "            output_dir=\"./results\",\n",
-    "            overwrite_output_dir=True,\n",
-    "            per_device_eval_batch_size=4,\n",
-    "            per_device_train_batch_size=4,\n",
-    "            report_to=[\"comet_ml\"],\n",
-    "            seed=SEED,\n",
-    "        )\n",
-    "        trainer = Trainer(\n",
-    "            model=model,\n",
-    "            args=training_args,\n",
-    "            train_dataset=small_train_dataset,\n",
-    "            eval_dataset=small_eval_dataset,\n",
-    "            compute_metrics=compute_metrics,\n",
-    "        )\n",
-    "\n",
-    "        # Report Metrics and Checkpoints to Ray Train\n",
-    "        callback = ray.train.huggingface.transformers.RayTrainReportCallback()\n",
-    "        trainer.add_callback(callback)\n",
-    "\n",
-    "        # Prepare Transformers Trainer\n",
-    "        trainer = ray.train.huggingface.transformers.prepare_trainer(trainer)\n",
-    "\n",
-    "        # Start Training\n",
-    "        trainer.train()\n",
-    "\n",
-    "    comet_ml.get_running_experiment().end()"
+    "\n",
+    "    small_train_dataset, small_eval_dataset = get_dataset()\n",
+    "\n",
+    "    # Model\n",
+    "    model = AutoModelForSequenceClassification.from_pretrained(\n",
+    "        \"google-bert/bert-base-cased\", num_labels=5\n",
+    "    )\n",
+    "\n",
+    "    # Evaluation Metrics\n",
+    "    metric = evaluate.load(\"accuracy\")\n",
+    "\n",
+    "    def compute_metrics(eval_pred):\n",
+    "        logits, labels = eval_pred\n",
+    "        predictions = np.argmax(logits, axis=-1)\n",
+    "\n",
+    "        experiment = comet_ml.get_running_experiment()\n",
+    "        if experiment:\n",
+    "            experiment.log_confusion_matrix(predictions, labels)\n",
+    "\n",
+    "        return metric.compute(predictions=predictions, references=labels)\n",
+    "\n",
+    "    # Hugging Face Trainer\n",
+    "    training_args = TrainingArguments(\n",
+    "        do_eval=True,\n",
+    "        do_train=True,\n",
+    "        eval_strategy=\"epoch\",\n",
+    "        num_train_epochs=config[\"epochs\"],\n",
+    "        output_dir=\"./results\",\n",
+    "        overwrite_output_dir=True,\n",
+    "        per_device_eval_batch_size=4,\n",
+    "        per_device_train_batch_size=4,\n",
+    "        report_to=[\"comet_ml\"],\n",
+    "        seed=SEED,\n",
+    "    )\n",
+    "    trainer = Trainer(\n",
+    "        model=model,\n",
+    "        args=training_args,\n",
+    "        train_dataset=small_train_dataset,\n",
+    "        eval_dataset=small_eval_dataset,\n",
+    "        compute_metrics=compute_metrics,\n",
+    "    )\n",
+    "\n",
+    "    # Report Metrics and Checkpoints to Ray Train\n",
+    "    callback = ray.train.huggingface.transformers.RayTrainReportCallback()\n",
+    "    trainer.add_callback(callback)\n",
+    "\n",
+    "    # Prepare Transformers Trainer\n",
+    "    trainer = ray.train.huggingface.transformers.prepare_trainer(trainer)\n",
+    "\n",
+    "    # Start Training\n",
+    "    trainer.train()\n",
+    "\n",
+    "    comet_ml.end()"
    ]
   },
   {
@@ -240,16 +240,15 @@
     "    scaling_config = ScalingConfig(num_workers=num_workers, use_gpu=use_gpu)\n",
     "    config = {\"use_gpu\": use_gpu, \"epochs\": 2}\n",
     "\n",
-    "    callback = comet_ml.integration.ray.CometTrainLoggerCallback(\n",
-    "        config, project_name=\"comet-example-ray-train-hugginface-transformers\"\n",
-    "    )\n",
-    "\n",
     "    ray_trainer = TorchTrainer(\n",
     "        train_func,\n",
     "        scaling_config=scaling_config,\n",
     "        train_loop_config=config,\n",
-    "        run_config=RunConfig(callbacks=[callback]),\n",
     "    )\n",
+    "    comet_ml.integration.ray.comet_ray_train_logger(\n",
+    "        ray_trainer, project_name=\"comet-example-ray-train-hugginface-transformers\"\n",
+    "    )\n",
+    "\n",
     "    result = ray_trainer.fit()"
    ]
   },
@@ -278,13 +277,6 @@
     "\n",
     "train(num_workers, use_gpu=False, epochs=5)"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {

integrations/model-training/ray-train/notebooks/Comet_with_ray_train_keras.ipynb

@@ -41,7 +41,7 @@
    },
    "outputs": [],
    "source": [
-    "%pip install -U \"comet_ml>=3.44.0\" \"ray[air]>=2.1.0\" \"keras<3\" \"tensorflow<2.16.0\""
+    "%pip install -U \"comet_ml>=3.49.0\" \"ray[air]>=2.1.0\" \"keras<3\" \"tensorflow<2.16.0\""
    ]
   },
   {
@@ -88,6 +88,7 @@
     "import os\n",
     "\n",
     "import comet_ml.integration.ray\n",
+    "from comet_ml.integration.ray import comet_worker\n",
     "\n",
     "import numpy as np\n",
     "import ray\n",
@@ -172,45 +173,43 @@
    },
    "outputs": [],
    "source": [
+    "@comet_worker\n",
     "def train_func(config: dict):\n",
-    "    from comet_ml.integration.ray import comet_worker_logger\n",
     "    from ray.air import session\n",
     "\n",
     "    per_worker_batch_size = config.get(\"batch_size\", 64)\n",
     "    epochs = config.get(\"epochs\", 3)\n",
     "    steps_per_epoch = config.get(\"steps_per_epoch\", 70)\n",
     "\n",
-    "    with comet_worker_logger(config) as experiment:\n",
+    "    tf_config = json.loads(os.environ[\"TF_CONFIG\"])\n",
+    "    num_workers = len(tf_config[\"cluster\"][\"worker\"])\n",
     "\n",
-    "        tf_config = json.loads(os.environ[\"TF_CONFIG\"])\n",
-    "        num_workers = len(tf_config[\"cluster\"][\"worker\"])\n",
+    "    strategy = tf.distribute.MultiWorkerMirroredStrategy()\n",
     "\n",
-    "        strategy = tf.distribute.MultiWorkerMirroredStrategy()\n",
+    "    global_batch_size = per_worker_batch_size * num_workers\n",
+    "    multi_worker_dataset = mnist_dataset(global_batch_size)\n",
     "\n",
-    "        global_batch_size = per_worker_batch_size * num_workers\n",
-    "        multi_worker_dataset = mnist_dataset(global_batch_size)\n",
-    "\n",
-    "        with strategy.scope():\n",
-    "            # Model building/compiling need to be within `strategy.scope()`.\n",
-    "            multi_worker_model = build_cnn_model()\n",
-    "            learning_rate = config.get(\"lr\", 0.001)\n",
-    "            multi_worker_model.compile(\n",
-    "                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
-    "                optimizer=tf.keras.optimizers.SGD(learning_rate=learning_rate),\n",
-    "                metrics=[\"accuracy\"],\n",
-    "            )\n",
+    "    with strategy.scope():\n",
+    "        # Model building/compiling need to be within `strategy.scope()`.\n",
+    "        multi_worker_model = build_cnn_model()\n",
+    "        learning_rate = config.get(\"lr\", 0.001)\n",
+    "        multi_worker_model.compile(\n",
+    "            loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
+    "            optimizer=tf.keras.optimizers.SGD(learning_rate=learning_rate),\n",
+    "            metrics=[\"accuracy\"],\n",
+    "        )\n",
     "\n",
-    "        callbacks = []\n",
-    "        if session.get_world_rank() == 0:\n",
-    "            callbacks.append(experiment.get_callback(\"tf-keras\"))\n",
+    "    callbacks = []\n",
+    "    if session.get_world_rank() == 0:\n",
+    "        callbacks.append(comet_ml.get_running_experiment().get_callback(\"tf-keras\"))\n",
     "\n",
-    "        history = multi_worker_model.fit(\n",
-    "            multi_worker_dataset,\n",
-    "            epochs=epochs,\n",
-    "            steps_per_epoch=steps_per_epoch,\n",
-    "            callbacks=callbacks,\n",
-    "        )\n",
-    "        results = history.history\n",
+    "    history = multi_worker_model.fit(\n",
+    "        multi_worker_dataset,\n",
+    "        epochs=epochs,\n",
+    "        steps_per_epoch=steps_per_epoch,\n",
+    "        callbacks=callbacks,\n",
+    "    )\n",
+    "    results = history.history\n",
     "\n",
     "    return results"
    ]
@@ -233,14 +232,15 @@
     ") -> Result:\n",
     "    config = {\"lr\": 1e-3, \"batch_size\": 64, \"epochs\": epochs}\n",
     "\n",
-    "    callback = comet_ml.integration.ray.CometTrainLoggerCallback(config)\n",
-    "\n",
     "    trainer = TensorflowTrainer(\n",
     "        train_loop_per_worker=train_func,\n",
     "        train_loop_config=config,\n",
     "        scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),\n",
-    "        run_config=RunConfig(callbacks=[callback]),\n",
     "    )\n",
+    "    comet_ml.integration.ray.comet_ray_train_logger(\n",
+    "        trainer, project_name=\"comet-example-ray-train-keras\"\n",
+    "    )\n",
+    "\n",
     "    results = trainer.fit()\n",
     "    return results"
    ]
@@ -270,6 +270,15 @@
     "\n",
     "train_tensorflow_mnist(num_workers, use_gpu=False, epochs=10)"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "comet_ml.end()"
+   ]
   }
  ],
  "metadata": {
@@ -291,7 +300,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.12"
+   "version": "3.9.1"
   }
  },
  "nbformat": 4,