microsoft--rd-agent
wehub-resource-sync
105 行
3.6 KiB
Python
105 行
3.6 KiB
Python
"""
|
|
FT-specific Workspace implementation with minimal checkpoint strategy.
|
|
|
|
This module provides FTWorkspace, which configures checkpoint to only save
|
|
configuration files (train.yaml), excluding all training outputs.
|
|
|
|
Design Philosophy:
|
|
- Checkpoint is for code version control during CoSTEER evolution
|
|
- Model persistence is handled separately by Runner's save_model()
|
|
- This separation keeps concerns clear and checkpoints lightweight
|
|
"""
|
|
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from rdagent.components.coder.finetune.conf import FT_YAML_FILE_NAME
|
|
from rdagent.core.conf import RD_AGENT_SETTINGS
|
|
from rdagent.core.experiment import FBWorkspace
|
|
from rdagent.log import rdagent_logger as logger
|
|
from rdagent.utils.env import CacheKeyFunc, DockerEnv, LocalEnv
|
|
|
|
if TYPE_CHECKING:
|
|
from rdagent.utils.env import Env
|
|
|
|
from rdagent.utils.env import EnvResult
|
|
|
|
|
|
class FTWorkspace(FBWorkspace):
|
|
"""
|
|
Fine-tuning workspace with minimal checkpoint strategy and unified Docker logging.
|
|
|
|
Checkpoint Strategy:
|
|
- Only saves configuration files (train.yaml) for version control
|
|
- Training outputs (models, checkpoints) are excluded by design
|
|
- Final model persistence is Runner's responsibility, not checkpoint's
|
|
"""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
|
|
# Configure checkpoint to save essential files for training
|
|
# Training outputs (models, checkpoints) are managed separately by save_final_model()
|
|
RD_AGENT_SETTINGS.workspace_ckp_white_list_names = [
|
|
FT_YAML_FILE_NAME, # train.yaml - training config
|
|
"dataset_info.json", # LlamaFactory dataset config
|
|
]
|
|
RD_AGENT_SETTINGS.workspace_ckp_size_limit = 100 * 1024
|
|
|
|
def run(
|
|
self,
|
|
env: "Env",
|
|
entry: str,
|
|
env_vars: dict | None = None,
|
|
cache_key_extra_func: CacheKeyFunc | None = None,
|
|
cache_files_to_extract: list[str] | None = None,
|
|
) -> "EnvResult":
|
|
"""Execute the code in the environment with unified Docker logging.
|
|
|
|
Args:
|
|
env: The environment to run in (DockerEnv, LocalEnv, etc.)
|
|
entry: The command to execute
|
|
env_vars: Optional additional environment variables (e.g., LLM API keys)
|
|
Will be merged with default {"PYTHONPATH": "./"}
|
|
cache_key_extra_func: Optional extra function for cache key calculation
|
|
cache_files_to_extract: Optional list of files to extract from cache
|
|
|
|
Returns:
|
|
EnvResult with stdout, exit_code, running_time
|
|
"""
|
|
self.prepare()
|
|
self.inject_files(**self.file_dict)
|
|
|
|
# Merge default env with custom env_vars
|
|
run_env = {"PYTHONPATH": "./"}
|
|
if env_vars:
|
|
run_env.update(env_vars)
|
|
|
|
result = env.run(
|
|
entry,
|
|
str(self.workspace_path),
|
|
env=run_env,
|
|
cache_key_extra_func=cache_key_extra_func,
|
|
cache_files_to_extract=cache_files_to_extract,
|
|
)
|
|
|
|
# Unified execution logging for FT scenario (supports both Docker and Conda)
|
|
if isinstance(env, DockerEnv):
|
|
tag_prefix = "docker_run"
|
|
elif isinstance(env, LocalEnv):
|
|
tag_prefix = "conda_run"
|
|
else:
|
|
tag_prefix = "env_run"
|
|
|
|
logger.log_object(
|
|
{
|
|
"exit_code": result.exit_code,
|
|
"stdout": result.stdout or "",
|
|
"running_time": result.running_time,
|
|
"entry": entry,
|
|
"workspace_path": str(self.workspace_path),
|
|
},
|
|
tag=f"{tag_prefix}.FTWorkspace",
|
|
)
|
|
|
|
return result
|