From 750604c7af8d983ed8e6d94b6d129efb6ffdcedc Mon Sep 17 00:00:00 2001
From: Botao Chen <markchen1015@meta.com>
Date: Thu, 2 Jan 2025 13:08:20 -0800
Subject: [PATCH] [Post Training] Fix missing import (#705)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## context
Post training apis are broken after the import * refactor
https://github.com/meta-llama/llama-stack/pull/689. This PR is adding
the missing import back

## Test
Issue a post training request from client and the training finishes
successfully

<img width="1101" alt="Screenshot 2025-01-02 at 12 18 45 PM"
src="https://github.com/user-attachments/assets/8c781459-f340-4021-85e1-fc68b1dcb8c8"
/>

<img width="782" alt="Screenshot 2025-01-02 at 12 18 52 PM"
src="https://github.com/user-attachments/assets/14b04b7d-e5c7-4662-8fa6-748446ad3511"
/>
---
 .../providers/inline/post_training/torchtune/common/utils.py    | 2 ++
 .../torchtune/recipes/lora_finetuning_single_device.py          | 1 +
 2 files changed, 3 insertions(+)

diff --git a/llama_stack/providers/inline/post_training/torchtune/common/utils.py b/llama_stack/providers/inline/post_training/torchtune/common/utils.py
index f2a2edae5..9673e0732 100644
--- a/llama_stack/providers/inline/post_training/torchtune/common/utils.py
+++ b/llama_stack/providers/inline/post_training/torchtune/common/utils.py
@@ -15,6 +15,8 @@ from typing import Any, Callable, Dict, List
 
 import torch
 from llama_models.datatypes import Model
+
+from llama_models.llama3.api.datatypes import BaseModel
 from llama_models.sku_list import resolve_model
 from llama_stack.apis.common.type_system import ParamType, StringType
 from llama_stack.apis.datasets import Datasets
diff --git a/llama_stack/providers/inline/post_training/torchtune/recipes/lora_finetuning_single_device.py b/llama_stack/providers/inline/post_training/torchtune/recipes/lora_finetuning_single_device.py
index 517be6d89..1b6c508a7 100644
--- a/llama_stack/providers/inline/post_training/torchtune/recipes/lora_finetuning_single_device.py
+++ b/llama_stack/providers/inline/post_training/torchtune/recipes/lora_finetuning_single_device.py
@@ -7,6 +7,7 @@
 import logging
 import os
 import time
+from datetime import datetime
 from functools import partial
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple