Merging pull request #16877

Author: jcortes

components/lamini/actions/create-fine-tune-job/create-fine-tune-job.mjs

@@ -6,7 +6,7 @@ export default {
   key: "lamini-create-fine-tune-job",
   name: "Create Fine-Tune Job",
   description: "Create a fine-tuning job with a dataset. [See the documentation](https://docs.lamini.ai/api/).",
-  version: "0.0.1",
+  version: "0.0.2",
   type: "action",
   props: {
     app,
@@ -49,6 +49,13 @@ export default {
       description: "A human-readable name for the fine-tuned model.",
       optional: true,
     },
+    waitForCompletion: {
+      type: "boolean",
+      label: "Wait for Completion",
+      description: "If set to `true`, the action will wait and poll until the fine-tuning job is `COMPLETED`. If is set to `false`, it will return immediately after creating the job.",
+      default: false,
+      optional: true,
+    },
   },
   methods: {
     createFineTuneJob(args = {}) {
@@ -69,35 +76,87 @@ export default {
       gpuConfig,
       isPublic,
       customModelName,
+      waitForCompletion,
     } = this;
 
-    const { upload_base_path: uploadBasePath } =
-      await app.getUploadBasePath({
+    const MAX_RETRIES = 15;
+    const DELAY = 1000 * 30; // 30 seconds
+    const { run } = $.context;
+
+    // First run: Create the fine-tune job
+    if (run.runs === 1) {
+      const { upload_base_path: uploadBasePath } =
+        await app.getUploadBasePath({
+          $,
+        });
+
+      await app.getExistingDataset({
         $,
+        data: {
+          dataset_id: datasetId,
+          upload_base_path: uploadBasePath,
+        },
       });
 
-    await app.getExistingDataset({
-      $,
-      data: {
-        dataset_id: datasetId,
-        upload_base_path: uploadBasePath,
-      },
-    });
+      const response = await createFineTuneJob({
+        $,
+        data: {
+          model_name: modelName,
+          dataset_id: datasetId,
+          upload_file_path: `${uploadBasePath}/${datasetId}.jsonlines`,
+          finetune_args: utils.parseJson(fineTuneArgs),
+          gpu_config: utils.parseJson(gpuConfig),
+          is_public: isPublic,
+          custom_model_name: customModelName,
+        },
+      });
+
+      $.export("$summary", `Successfully created a fine-tune job with ID \`${response.job_id}\`.`);
+
+      // If user doesn't want to wait, return immediately
+      if (!waitForCompletion) {
+        return response;
+      }
+
+      // Store job_id for polling and start rerun
+      $.flow.rerun(DELAY, {
+        jobId: response.job_id,
+      }, MAX_RETRIES);
+      return response;
+    }
 
-    const response = await createFineTuneJob({
+    // Subsequent runs: Poll for job status
+    if (run.runs > MAX_RETRIES) {
+      throw new Error("Max retries exceeded - fine-tuning job may still be running");
+    }
+
+    const { jobId } = run.context;
+
+    // Poll for status
+    const statusResponse = await app.getJobStatus({
       $,
-      data: {
-        model_name: modelName,
-        dataset_id: datasetId,
-        upload_file_path: `${uploadBasePath}/${datasetId}.jsonlines`,
-        finetune_args: utils.parseJson(fineTuneArgs),
-        gpu_config: utils.parseJson(gpuConfig),
-        is_public: isPublic,
-        custom_model_name: customModelName,
-      },
+      jobId,
     });
 
-    $.export("$summary", `Successfully created a fine-tune job with ID \`${response.job_id}\`.`);
-    return response;
+    // If job is completed, return the final status
+    if (statusResponse.status === "COMPLETED") {
+      $.export("$summary", `Fine-tuning job \`${jobId}\` completed successfully.`);
+      return statusResponse;
+    }
+
+    // If job failed, throw error
+    if (statusResponse.status === "FAILED") {
+      throw new Error(`Fine-tuning job \`${jobId}\` failed.`);
+    }
+
+    // Otherwise, continue polling
+    $.flow.rerun(DELAY, {
+      jobId,
+    }, MAX_RETRIES);
+    return {
+      status: statusResponse.status,
+      jobId,
+      message: `Job is still running. Current status: ${statusResponse.status}`,
+    };
   },
 };

components/lamini/actions/evaluate-job/evaluate-job.mjs

@@ -4,14 +4,17 @@ export default {
   key: "lamini-evaluate-job",
   name: "Evaluate Job",
   description: "Evaluate a fine-tuning job by job ID. [See the documentation](https://docs.lamini.ai/api/).",
-  version: "0.0.1",
+  version: "0.0.2",
   type: "action",
   props: {
     app,
     jobId: {
       propDefinition: [
         app,
         "jobId",
+        () => ({
+          filter: ({ status }) => status === "COMPLETED",
+        }),
       ],
       description: "The ID of the fine-tuning job to evaluate.",
     },

components/lamini/actions/generate-completion/generate-completion.mjs

@@ -5,7 +5,7 @@ export default {
   key: "lamini-generate-completion",
   name: "Generate Completion",
   description: "Generate completions using a Lamini model. [See the documentation](https://docs.lamini.ai/api/).",
-  version: "0.0.1",
+  version: "0.0.2",
   type: "action",
   props: {
     app,

components/lamini/actions/get-batch-completions/get-batch-completions.mjs

@@ -5,7 +5,7 @@ export default {
   key: "lamini-get-batch-completions",
   name: "Get Batch Completions",
   description: "Retrieve the results of a batch completion request from Lamini. [See the documentation](https://docs.lamini.ai/api/).",
-  version: "0.0.1",
+  version: "0.0.2",
   type: "action",
   props: {
     app,

components/lamini/actions/upload-dataset/upload-dataset.mjs

@@ -7,7 +7,7 @@ export default {
   key: "lamini-upload-dataset",
   name: "Upload Dataset",
   description: "Upload a dataset to Lamini for training.",
-  version: "0.0.1",
+  version: "0.0.2",
   type: "action",
   props: {
     app,
@@ -19,12 +19,12 @@ export default {
     inputKey: {
       type: "string",
       label: "Input Key",
-      description: "Key of the JSON dictionary to use as the input. For CSV files, this should be a column header. Eg. `question`.",
+      description: "Key of the JSON dictionary to use as the input. Eg. `question`.",
     },
     outputKey: {
       type: "string",
       label: "Output Key",
-      description: "Key of the JSON dictionary to use as the output. For CSV files, this should be a column header. Eg. `answer`.",
+      description: "Key of the JSON dictionary to use as the output. Eg. `answer`.",
       optional: true,
     },
     isPublic: {

components/lamini/lamini.app.mjs

@@ -38,10 +38,10 @@ export default {
       type: "string",
       label: "Job ID",
       description: "The ID of the fine-tuning job to use.",
-      async options() {
+      async options({ filter = () => true }) {
         const jobs = await this.listTrainedJobs();
         return jobs
-          .filter(({ status }) => status === "COMPLETED")
+          .filter(filter)
           .map(({
             custom_model_name: label,
             job_id: value,

components/lamini/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@pipedream/lamini",
-  "version": "0.1.0",
+  "version": "0.1.1",
   "description": "Pipedream Lamini Components",
   "main": "lamini.app.mjs",
   "keywords": [

components/lamini/sources/get-job-status/get-job-status.mjs

@@ -8,7 +8,7 @@ export default {
   key: "lamini-get-job-status",
   name: "New Job Status",
   description: "Emit new events with the status of a training job. [See the documentation](https://docs.lamini.ai/api/).",
-  version: "0.0.1",
+  version: "0.0.2",
   type: "source",
   props: {
     app,