ludwig-ai--ludwig
wehub-resource-sync
593b94c120
pytest / Unit Tests (push) Has been cancelled
pytest / Integration (integration_tests_a) (push) Has been cancelled
pytest / Integration (integration_tests_b) (push) Has been cancelled
pytest / Integration (integration_tests_c) (push) Has been cancelled
pytest / Integration (integration_tests_d) (push) Has been cancelled
pytest / Integration (integration_tests_e) (push) Has been cancelled
pytest / Integration (integration_tests_f) (push) Has been cancelled
pytest / Integration (integration_tests_g) (push) Has been cancelled
pytest / Integration (integration_tests_h) (push) Has been cancelled
pytest / Integration (integration_tests_i) (push) Has been cancelled
pytest / Integration (integration_tests_j) (push) Has been cancelled
pytest / Distributed (distributed_a) (push) Has been cancelled
pytest / Distributed (distributed_b) (push) Has been cancelled
pytest / Distributed (distributed_c) (push) Has been cancelled
pytest / Distributed (distributed_d) (push) Has been cancelled
pytest / Distributed (distributed_e) (push) Has been cancelled
pytest / Distributed (distributed_f) (push) Has been cancelled
pytest / Minimal Install (push) Has been cancelled
pytest / Event File (push) Has been cancelled
pytest (slow) / py-slow (push) Has been cancelled
Publish JSON Schema / publish-schema (push) Has been cancelled
60 行
2.2 KiB
Python
60 行
2.2 KiB
Python
from typing import Any
|
|
|
|
from ludwig.api_annotations import DeveloperAPI
|
|
from ludwig.benchmarking.profiler import LudwigProfiler
|
|
from ludwig.callbacks import Callback
|
|
from ludwig.constants import EVALUATION, PREPROCESSING, TRAINING
|
|
|
|
|
|
# TODO: Change annotation to PublicAPI once Ludwig 0.7 is released
|
|
@DeveloperAPI
|
|
class LudwigProfilerCallback(Callback):
|
|
"""Class that defines the methods necessary to hook into process."""
|
|
|
|
def __init__(self, experiment: dict[str, Any]):
|
|
self.experiment_name = experiment["experiment_name"]
|
|
self.use_torch_profiler = experiment["profiler"]["use_torch_profiler"]
|
|
self.logging_interval = experiment["profiler"]["logging_interval"]
|
|
self.preprocess_profiler = None
|
|
self.train_profiler = None
|
|
self.evaluation_profiler = None
|
|
|
|
def on_preprocess_start(self, *args, **kwargs):
|
|
self.preprocess_profiler = LudwigProfiler(
|
|
tag=PREPROCESSING,
|
|
output_dir=self.experiment_name,
|
|
use_torch_profiler=self.use_torch_profiler,
|
|
logging_interval=self.logging_interval,
|
|
)
|
|
self.preprocess_profiler.__enter__()
|
|
|
|
def on_preprocess_end(self, *args, **kwargs):
|
|
self.preprocess_profiler.__exit__(None, None, None)
|
|
del self.preprocess_profiler
|
|
|
|
def on_train_start(self, *args, **kwargs):
|
|
self.train_profiler = LudwigProfiler(
|
|
tag=TRAINING,
|
|
output_dir=self.experiment_name,
|
|
use_torch_profiler=self.use_torch_profiler,
|
|
logging_interval=self.logging_interval,
|
|
)
|
|
self.train_profiler.__enter__()
|
|
|
|
def on_train_end(self, *args, **kwargs):
|
|
self.train_profiler.__exit__(None, None, None)
|
|
del self.train_profiler
|
|
|
|
def on_evaluation_start(self):
|
|
self.evaluation_profiler = LudwigProfiler(
|
|
tag=EVALUATION,
|
|
output_dir=self.experiment_name,
|
|
use_torch_profiler=self.use_torch_profiler,
|
|
logging_interval=self.logging_interval,
|
|
)
|
|
self.evaluation_profiler.__enter__()
|
|
|
|
def on_evaluation_end(self):
|
|
self.evaluation_profiler.__exit__(None, None, None)
|
|
del self.evaluation_profiler
|