collie-mlops 0.1.1b0__tar.gz → 0.1.2b0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (54) hide show
  1. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/MANIFEST.in +1 -0
  2. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/PKG-INFO +121 -22
  3. collie-mlops-0.1.1b0/collie_mlops.egg-info/PKG-INFO → collie_mlops-0.1.2b0/README.md +83 -56
  4. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/contracts/mlflow.py +10 -4
  5. collie_mlops-0.1.2b0/collie_mlops.egg-info/PKG-INFO +312 -0
  6. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie_mlops.egg-info/SOURCES.txt +0 -9
  7. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie_mlops.egg-info/requires.txt +13 -33
  8. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/pyproject.toml +55 -3
  9. collie-mlops-0.1.1b0/README.md +0 -170
  10. collie-mlops-0.1.1b0/collie/helper/__init__.py +0 -0
  11. collie-mlops-0.1.1b0/collie/helper/pytorch/__init__.py +0 -0
  12. collie-mlops-0.1.1b0/collie/helper/pytorch/callback/__init__.py +0 -0
  13. collie-mlops-0.1.1b0/collie/helper/pytorch/callback/callback.py +0 -155
  14. collie-mlops-0.1.1b0/collie/helper/pytorch/callback/earlystop.py +0 -54
  15. collie-mlops-0.1.1b0/collie/helper/pytorch/callback/model_checkpoint.py +0 -100
  16. collie-mlops-0.1.1b0/collie/helper/pytorch/model/__init__.py +0 -0
  17. collie-mlops-0.1.1b0/collie/helper/pytorch/model/loader.py +0 -55
  18. collie-mlops-0.1.1b0/collie/helper/pytorch/trainer.py +0 -304
  19. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/LICENSE +0 -0
  20. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/__init__.py +0 -0
  21. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/_common/__init__.py +0 -0
  22. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/_common/decorator.py +0 -0
  23. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/_common/exceptions.py +0 -0
  24. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/_common/mlflow_model_io/__init__.py +0 -0
  25. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/_common/mlflow_model_io/base_flavor_handler.py +0 -0
  26. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/_common/mlflow_model_io/flavor_registry.py +0 -0
  27. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/_common/mlflow_model_io/model_flavors.py +0 -0
  28. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/_common/mlflow_model_io/model_io.py +0 -0
  29. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/_common/utils.py +0 -0
  30. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/contracts/__init__.py +0 -0
  31. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/contracts/event.py +0 -0
  32. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/contracts/orchestrator.py +0 -0
  33. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/core/__init__.py +0 -0
  34. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/core/enums/__init__.py +0 -0
  35. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/core/enums/components.py +0 -0
  36. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/core/enums/ml_models.py +0 -0
  37. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/core/evaluator/__init__.py +0 -0
  38. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/core/evaluator/evaluator.py +0 -0
  39. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/core/models.py +0 -0
  40. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/core/orchestrator/__init__.py +0 -0
  41. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/core/orchestrator/orchestrator.py +0 -0
  42. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/core/pusher/__init__.py +0 -0
  43. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/core/pusher/pusher.py +0 -0
  44. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/core/trainer/__init__.py +0 -0
  45. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/core/trainer/trainer.py +0 -0
  46. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/core/transform/__init__.py +0 -0
  47. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/core/transform/transform.py +0 -0
  48. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/core/tuner/__init__.py +0 -0
  49. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie/core/tuner/tuner.py +0 -0
  50. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie_mlops.egg-info/dependency_links.txt +0 -0
  51. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie_mlops.egg-info/not-zip-safe +0 -0
  52. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/collie_mlops.egg-info/top_level.txt +0 -0
  53. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/requirements.txt +0 -0
  54. {collie-mlops-0.1.1b0 → collie_mlops-0.1.2b0}/setup.cfg +0 -0
@@ -9,6 +9,7 @@ recursive-include collie py.typed
9
9
  recursive-exclude tests *
10
10
  recursive-exclude example *
11
11
  recursive-exclude deploy *
12
+ recursive-exclude collie/helper *
12
13
  recursive-exclude __pycache__ *
13
14
  recursive-exclude *.pyc
14
15
  recursive-exclude *.pyo
@@ -1,6 +1,6 @@
1
- Metadata-Version: 2.1
1
+ Metadata-Version: 2.4
2
2
  Name: collie-mlops
3
- Version: 0.1.1b0
3
+ Version: 0.1.2b0
4
4
  Summary: A Lightweight MLOps Framework for Machine Learning Workflows
5
5
  Home-page: https://github.com/ChingHuanChiu/collie
6
6
  Author: ChingHuanChiu
@@ -29,17 +29,46 @@ Classifier: Topic :: Software Development :: Libraries :: Python Modules
29
29
  Classifier: Typing :: Typed
30
30
  Requires-Python: >=3.10
31
31
  Description-Content-Type: text/markdown
32
+ License-File: LICENSE
33
+ Requires-Dist: mlflow>=2.0.0
34
+ Requires-Dist: pydantic>=2.0.0
35
+ Requires-Dist: pandas>=1.3.0
36
+ Requires-Dist: numpy<2.0.0,>=1.20.0
32
37
  Provides-Extra: sklearn
38
+ Requires-Dist: scikit-learn>=1.0.0; extra == "sklearn"
33
39
  Provides-Extra: xgboost
34
- Provides-Extra: pytorch
40
+ Requires-Dist: xgboost>=1.5.0; extra == "xgboost"
35
41
  Provides-Extra: lightgbm
36
- Provides-Extra: transformers
37
- Provides-Extra: config
38
- Provides-Extra: dev
39
- Provides-Extra: test
40
- Provides-Extra: docs
42
+ Requires-Dist: lightgbm>=3.0.0; extra == "lightgbm"
43
+ Provides-Extra: pytorch
44
+ Requires-Dist: torch>=1.9.0; extra == "pytorch"
45
+ Requires-Dist: pytorch-lightning>=2.0.0; extra == "pytorch"
46
+ Requires-Dist: transformers>=4.0.0; extra == "pytorch"
47
+ Requires-Dist: sentence-transformers>=2.0.0; extra == "pytorch"
48
+ Provides-Extra: tabular
49
+ Requires-Dist: scikit-learn>=1.0.0; extra == "tabular"
50
+ Requires-Dist: xgboost>=1.5.0; extra == "tabular"
51
+ Requires-Dist: lightgbm>=3.0.0; extra == "tabular"
52
+ Provides-Extra: deep-learning
53
+ Requires-Dist: torch>=1.9.0; extra == "deep-learning"
54
+ Requires-Dist: pytorch-lightning>=2.0.0; extra == "deep-learning"
55
+ Requires-Dist: transformers>=4.0.0; extra == "deep-learning"
56
+ Requires-Dist: sentence-transformers>=2.0.0; extra == "deep-learning"
41
57
  Provides-Extra: all
42
- License-File: LICENSE
58
+ Requires-Dist: scikit-learn>=1.0.0; extra == "all"
59
+ Requires-Dist: xgboost>=1.5.0; extra == "all"
60
+ Requires-Dist: lightgbm>=3.0.0; extra == "all"
61
+ Requires-Dist: torch>=1.9.0; extra == "all"
62
+ Requires-Dist: pytorch-lightning>=2.0.0; extra == "all"
63
+ Requires-Dist: transformers>=4.0.0; extra == "all"
64
+ Requires-Dist: sentence-transformers>=2.0.0; extra == "all"
65
+ Provides-Extra: dev
66
+ Requires-Dist: pytest>=7.0.0; extra == "dev"
67
+ Requires-Dist: pytest-cov>=3.0.0; extra == "dev"
68
+ Requires-Dist: black>=22.0.0; extra == "dev"
69
+ Requires-Dist: mypy>=0.990; extra == "dev"
70
+ Requires-Dist: ruff>=0.0.260; extra == "dev"
71
+ Dynamic: license-file
43
72
 
44
73
  # Collie
45
74
 
@@ -79,16 +108,42 @@ Collie follows an event-driven architecture with the following core components:
79
108
 
80
109
  ### Installation
81
110
 
111
+ #### Basic Installation (Core Framework Only)
82
112
  ```bash
83
113
  pip install collie-mlops
84
114
  ```
85
115
 
86
- This will install Collie with all supported ML frameworks including:
87
- - scikit-learn
88
- - PyTorch
89
- - XGBoost
90
- - LightGBM
91
- - Transformers (with Sentence Transformers)
116
+ This installs the core MLOps orchestration framework with MLflow integration (~100MB).
117
+
118
+ #### Install with ML Frameworks
119
+
120
+ Choose the installation that fits your needs:
121
+
122
+ **For Traditional ML (Tabular Data)**
123
+ ```bash
124
+ # Individual frameworks
125
+ pip install collie-mlops[sklearn] # scikit-learn support
126
+ pip install collie-mlops[xgboost] # XGBoost support
127
+ pip install collie-mlops[lightgbm] # LightGBM support
128
+
129
+ # Or install all tabular ML frameworks (~250MB)
130
+ pip install collie-mlops[tabular]
131
+ ```
132
+
133
+ **For Deep Learning**
134
+ ```bash
135
+ # PyTorch ecosystem (includes Transformers for NLP/Vision) (~3GB)
136
+ pip install collie-mlops[pytorch]
137
+
138
+ # Or use the alias
139
+ pip install collie-mlops[deep-learning]
140
+ ```
141
+
142
+ **For Complete Installation**
143
+ ```bash
144
+ # All frameworks (~3.5GB)
145
+ pip install collie-mlops[all]
146
+ ```
92
147
 
93
148
  ### Prerequisites
94
149
 
@@ -157,6 +212,32 @@ class CustomPusher(Pusher):
157
212
  return Event(payload=PusherPayload(model_uri=model_uri))
158
213
  ```
159
214
 
215
+ ### Orchestrator
216
+ Coordinates the execution of all components in the pipeline.
217
+
218
+ ```python
219
+ from collie import Orchestrator
220
+
221
+ # Create orchestrator with all components
222
+ orchestrator = Orchestrator(
223
+ components=[
224
+ CustomTransformer(),
225
+ CustomTuner(),
226
+ CustomTrainer(),
227
+ CustomEvaluator(),
228
+ CustomPusher()
229
+ ],
230
+ tracking_uri="http://localhost:5000",
231
+ experiment_name="my_experiment",
232
+ registered_model_name="my_model",
233
+ mlflow_tags={"project": "my_project"},
234
+ description="My ML Pipeline"
235
+ )
236
+
237
+ # Run the entire pipeline
238
+ orchestrator.run()
239
+ ```
240
+
160
241
  ## Configuration
161
242
 
162
243
  ### MLflow Setup
@@ -173,13 +254,27 @@ mlflow server \
173
254
 
174
255
  ## Supported Frameworks
175
256
 
176
- Collie supports multiple ML frameworks through its model flavor system currently:
257
+ Collie supports multiple ML frameworks through its flexible optional dependency system:
258
+
259
+ ### Available Frameworks
260
+ - **scikit-learn** - Traditional ML algorithms
261
+ - **XGBoost** - Gradient boosting for tabular data
262
+ - **LightGBM** - Fast gradient boosting framework
263
+ - **PyTorch** - Deep learning framework
264
+ - **PyTorch Lightning** - High-level PyTorch wrapper
265
+ - **Transformers** - Hugging Face transformers for NLP
266
+ - **Sentence Transformers** - Sentence embeddings
267
+
268
+ ### Installation Options
177
269
 
178
- - **PyTorch**
179
- - **scikit-learn**
180
- - **XGBoost**
181
- - **LightGBM**
182
- - **Transformers**
270
+ | Use Case | Command | Size | Frameworks Included |
271
+ |----------|---------|------|---------------------|
272
+ | **Core Only** | `pip install collie-mlops` | ~100MB | MLflow orchestration only |
273
+ | **Tabular ML** | `pip install collie-mlops[tabular]` | ~250MB | sklearn, XGBoost, LightGBM |
274
+ | **Deep Learning** | `pip install collie-mlops[pytorch]` | ~3GB | PyTorch, Lightning, Transformers |
275
+ | **Complete** | `pip install collie-mlops[all]` | ~3.5GB | All frameworks |
276
+
277
+ > **Note**: Install only what you need to keep your environment lightweight!
183
278
 
184
279
 
185
280
  ## Documentation
@@ -188,9 +283,13 @@ Collie supports multiple ML frameworks through its model flavor system currently
188
283
 
189
284
  ## Roadmap
190
285
 
286
+ ### Core Features
287
+ - [ ] **Pipeline Checkpoint & Resume** - Save intermediate results and resume from failure points
288
+
289
+ ### Framework Support
191
290
  - [ ] TensorFlow/Keras support
192
291
  - [ ] Model monitoring and drift detection
193
- - [ ] Integrate an LLM training/fine-tuning framework
292
+
194
293
 
195
294
  ## License
196
295
 
@@ -1,46 +1,3 @@
1
- Metadata-Version: 2.1
2
- Name: collie-mlops
3
- Version: 0.1.1b0
4
- Summary: A Lightweight MLOps Framework for Machine Learning Workflows
5
- Home-page: https://github.com/ChingHuanChiu/collie
6
- Author: ChingHuanChiu
7
- Author-email: ChingHuanChiu <stevenchiou8@gmail.com>
8
- Maintainer-email: ChingHuanChiu <stevenchiou8@gmail.com>
9
- License: MIT
10
- Project-URL: Homepage, https://github.com/ChingHuanChiu/collie
11
- Project-URL: Documentation, https://github.com/ChingHuanChiu/collie/blob/main/README.md
12
- Project-URL: Repository, https://github.com/ChingHuanChiu/collie
13
- Project-URL: Bug Tracker, https://github.com/ChingHuanChiu/collie/issues
14
- Project-URL: Changelog, https://github.com/ChingHuanChiu/collie/blob/main/CHANGELOG.md
15
- Keywords: mlops,machine-learning,mlflow,pipeline,orchestration,deep-learning,experiment-tracking
16
- Classifier: Development Status :: 4 - Beta
17
- Classifier: Intended Audience :: Developers
18
- Classifier: Intended Audience :: Science/Research
19
- Classifier: License :: OSI Approved :: MIT License
20
- Classifier: Operating System :: OS Independent
21
- Classifier: Programming Language :: Python :: 3
22
- Classifier: Programming Language :: Python :: 3.10
23
- Classifier: Programming Language :: Python :: 3.11
24
- Classifier: Programming Language :: Python :: 3.12
25
- Classifier: Programming Language :: Python :: 3.13
26
- Classifier: Programming Language :: Python :: 3 :: Only
27
- Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
28
- Classifier: Topic :: Software Development :: Libraries :: Python Modules
29
- Classifier: Typing :: Typed
30
- Requires-Python: >=3.10
31
- Description-Content-Type: text/markdown
32
- Provides-Extra: sklearn
33
- Provides-Extra: xgboost
34
- Provides-Extra: pytorch
35
- Provides-Extra: lightgbm
36
- Provides-Extra: transformers
37
- Provides-Extra: config
38
- Provides-Extra: dev
39
- Provides-Extra: test
40
- Provides-Extra: docs
41
- Provides-Extra: all
42
- License-File: LICENSE
43
-
44
1
  # Collie
45
2
 
46
3
  [![PyPI version](https://badge.fury.io/py/collie-mlops.svg)](https://badge.fury.io/py/collie-mlops)
@@ -79,16 +36,42 @@ Collie follows an event-driven architecture with the following core components:
79
36
 
80
37
  ### Installation
81
38
 
39
+ #### Basic Installation (Core Framework Only)
82
40
  ```bash
83
41
  pip install collie-mlops
84
42
  ```
85
43
 
86
- This will install Collie with all supported ML frameworks including:
87
- - scikit-learn
88
- - PyTorch
89
- - XGBoost
90
- - LightGBM
91
- - Transformers (with Sentence Transformers)
44
+ This installs the core MLOps orchestration framework with MLflow integration (~100MB).
45
+
46
+ #### Install with ML Frameworks
47
+
48
+ Choose the installation that fits your needs:
49
+
50
+ **For Traditional ML (Tabular Data)**
51
+ ```bash
52
+ # Individual frameworks
53
+ pip install collie-mlops[sklearn] # scikit-learn support
54
+ pip install collie-mlops[xgboost] # XGBoost support
55
+ pip install collie-mlops[lightgbm] # LightGBM support
56
+
57
+ # Or install all tabular ML frameworks (~250MB)
58
+ pip install collie-mlops[tabular]
59
+ ```
60
+
61
+ **For Deep Learning**
62
+ ```bash
63
+ # PyTorch ecosystem (includes Transformers for NLP/Vision) (~3GB)
64
+ pip install collie-mlops[pytorch]
65
+
66
+ # Or use the alias
67
+ pip install collie-mlops[deep-learning]
68
+ ```
69
+
70
+ **For Complete Installation**
71
+ ```bash
72
+ # All frameworks (~3.5GB)
73
+ pip install collie-mlops[all]
74
+ ```
92
75
 
93
76
  ### Prerequisites
94
77
 
@@ -157,6 +140,32 @@ class CustomPusher(Pusher):
157
140
  return Event(payload=PusherPayload(model_uri=model_uri))
158
141
  ```
159
142
 
143
+ ### Orchestrator
144
+ Coordinates the execution of all components in the pipeline.
145
+
146
+ ```python
147
+ from collie import Orchestrator
148
+
149
+ # Create orchestrator with all components
150
+ orchestrator = Orchestrator(
151
+ components=[
152
+ CustomTransformer(),
153
+ CustomTuner(),
154
+ CustomTrainer(),
155
+ CustomEvaluator(),
156
+ CustomPusher()
157
+ ],
158
+ tracking_uri="http://localhost:5000",
159
+ experiment_name="my_experiment",
160
+ registered_model_name="my_model",
161
+ mlflow_tags={"project": "my_project"},
162
+ description="My ML Pipeline"
163
+ )
164
+
165
+ # Run the entire pipeline
166
+ orchestrator.run()
167
+ ```
168
+
160
169
  ## Configuration
161
170
 
162
171
  ### MLflow Setup
@@ -173,13 +182,27 @@ mlflow server \
173
182
 
174
183
  ## Supported Frameworks
175
184
 
176
- Collie supports multiple ML frameworks through its model flavor system currently:
185
+ Collie supports multiple ML frameworks through its flexible optional dependency system:
177
186
 
178
- - **PyTorch**
179
- - **scikit-learn**
180
- - **XGBoost**
181
- - **LightGBM**
182
- - **Transformers**
187
+ ### Available Frameworks
188
+ - **scikit-learn** - Traditional ML algorithms
189
+ - **XGBoost** - Gradient boosting for tabular data
190
+ - **LightGBM** - Fast gradient boosting framework
191
+ - **PyTorch** - Deep learning framework
192
+ - **PyTorch Lightning** - High-level PyTorch wrapper
193
+ - **Transformers** - Hugging Face transformers for NLP
194
+ - **Sentence Transformers** - Sentence embeddings
195
+
196
+ ### Installation Options
197
+
198
+ | Use Case | Command | Size | Frameworks Included |
199
+ |----------|---------|------|---------------------|
200
+ | **Core Only** | `pip install collie-mlops` | ~100MB | MLflow orchestration only |
201
+ | **Tabular ML** | `pip install collie-mlops[tabular]` | ~250MB | sklearn, XGBoost, LightGBM |
202
+ | **Deep Learning** | `pip install collie-mlops[pytorch]` | ~3GB | PyTorch, Lightning, Transformers |
203
+ | **Complete** | `pip install collie-mlops[all]` | ~3.5GB | All frameworks |
204
+
205
+ > **Note**: Install only what you need to keep your environment lightweight!
183
206
 
184
207
 
185
208
  ## Documentation
@@ -188,9 +211,13 @@ Collie supports multiple ML frameworks through its model flavor system currently
188
211
 
189
212
  ## Roadmap
190
213
 
214
+ ### Core Features
215
+ - [ ] **Pipeline Checkpoint & Resume** - Save intermediate results and resume from failure points
216
+
217
+ ### Framework Support
191
218
  - [ ] TensorFlow/Keras support
192
219
  - [ ] Model monitoring and drift detection
193
- - [ ] Integrate an LLM training/fine-tuning framework
220
+
194
221
 
195
222
  ## License
196
223
 
@@ -212,7 +212,7 @@ class _MLflowModelManager:
212
212
  def get_latest_version(
213
213
  self,
214
214
  model_name: str,
215
- stages
215
+ stages: List[str]
216
216
  ) -> Optional[str]:
217
217
 
218
218
  """
@@ -275,6 +275,8 @@ class MLFlowComponentABC(ABC):
275
275
  def model_manager(self) -> Optional[_MLflowModelManager]:
276
276
 
277
277
  if self._model_manager is None:
278
+ if self._mlflow_config is None:
279
+ raise MLflowConfigurationError("MLflow config not set")
278
280
  mlflow_client = self._mlflow_config.mlflow_client
279
281
  self._model_manager = _MLflowModelManager(mlflow_client)
280
282
  return self._model_manager
@@ -310,10 +312,11 @@ class MLFlowComponentABC(ABC):
310
312
 
311
313
  def load_model(
312
314
  self,
313
- name: Optional[str] = None
315
+ flavor: ModelFlavor,
316
+ model_uri: Optional[str] = None
314
317
  ) -> Any:
315
318
 
316
- return self.model_manager.load_model(name)
319
+ return self.model_manager.load_model(flavor, model_uri)
317
320
 
318
321
  def load_latest_model(
319
322
  self,
@@ -355,7 +358,7 @@ class MLFlowComponentABC(ABC):
355
358
  self,
356
359
  model_name: str,
357
360
  stages: List[Literal["None", "Staging", "Production", "Archived"]]
358
- ) -> int:
361
+ ) -> Optional[str]:
359
362
 
360
363
  return self.model_manager.get_latest_version(model_name, stages)
361
364
 
@@ -392,6 +395,9 @@ class MLFlowComponentABC(ABC):
392
395
  return None
393
396
  try:
394
397
  experiment = self.mlflow.get_experiment_by_name(experiment_name)
398
+ if experiment is None:
399
+ logger.warning(f"Experiment '{experiment_name}' not found")
400
+ return None
395
401
  if return_id:
396
402
  return experiment.experiment_id
397
403
  else: