forked from aws/sagemaker-python-sdk
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathestimator.py
More file actions
2740 lines (2399 loc) · 127 KB
/
estimator.py
File metadata and controls
2740 lines (2399 loc) · 127 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
# http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.
"""Placeholder docstring"""
from __future__ import absolute_import, print_function
import json
import logging
import os
import uuid
from abc import ABCMeta, abstractmethod
from six import string_types, with_metaclass
from six.moves.urllib.parse import urlparse
import sagemaker
from sagemaker import git_utils, image_uris, vpc_utils
from sagemaker.analytics import TrainingJobAnalytics
from sagemaker.debugger import ( # noqa: F401 # pylint: disable=unused-import
DEBUGGER_FLAG,
DebuggerHookConfig,
FrameworkProfile,
ProfilerConfig,
ProfilerRule,
Rule,
TensorBoardOutputConfig,
get_default_profiler_rule,
get_rule_container_image_uri,
)
from sagemaker.deprecations import removed_function, removed_kwargs, renamed_kwargs
from sagemaker.fw_utils import (
UploadedCode,
_region_supports_debugger,
_region_supports_profiler,
get_mp_parameters,
tar_and_upload_dir,
validate_source_dir,
)
from sagemaker.inputs import TrainingInput
from sagemaker.job import _Job
from sagemaker.local import LocalSession
from sagemaker.model import (
CONTAINER_LOG_LEVEL_PARAM_NAME,
DIR_PARAM_NAME,
JOB_NAME_PARAM_NAME,
NEO_ALLOWED_FRAMEWORKS,
SAGEMAKER_REGION_PARAM_NAME,
SCRIPT_PARAM_NAME,
Model,
)
from sagemaker.predictor import Predictor
from sagemaker.s3 import S3Uploader, parse_s3_url
from sagemaker.session import Session
from sagemaker.transformer import Transformer
from sagemaker.utils import (
base_from_name,
base_name_from_image,
build_dict,
get_config_value,
name_from_base,
)
from sagemaker.workflow.entities import Expression
from sagemaker.workflow.parameters import Parameter
from sagemaker.workflow.properties import Properties
logger = logging.getLogger(__name__)
class EstimatorBase(with_metaclass(ABCMeta, object)): # pylint: disable=too-many-public-methods
"""Handle end-to-end Amazon SageMaker training and deployment tasks.
For introduction to model training and deployment, see
http://docs.aws.amazon.com/sagemaker/latest/dg/how-it-works-training.html
Subclasses must define a way to determine what image to use for training,
what hyperparameters to use, and how to create an appropriate predictor
instance.
"""
def __init__(
self,
role,
instance_count=None,
instance_type=None,
volume_size=30,
volume_kms_key=None,
max_run=24 * 60 * 60,
input_mode="File",
output_path=None,
output_kms_key=None,
base_job_name=None,
sagemaker_session=None,
tags=None,
subnets=None,
security_group_ids=None,
model_uri=None,
model_channel_name="model",
metric_definitions=None,
encrypt_inter_container_traffic=False,
use_spot_instances=False,
max_wait=None,
checkpoint_s3_uri=None,
checkpoint_local_path=None,
rules=None,
debugger_hook_config=None,
tensorboard_output_config=None,
enable_sagemaker_metrics=None,
enable_network_isolation=False,
profiler_config=None,
disable_profiler=False,
environment=None,
max_retry_attempts=None,
**kwargs,
):
"""Initialize an ``EstimatorBase`` instance.
Args:
role (str): An AWS IAM role (either name or full ARN). The Amazon
SageMaker training jobs and APIs that create Amazon SageMaker
endpoints use this role to access training data and model
artifacts. After the endpoint is created, the inference code
might use the IAM role, if it needs to access an AWS resource.
instance_count (int): Number of Amazon EC2 instances to use
for training.
instance_type (str): Type of EC2 instance to use for training,
for example, 'ml.c4.xlarge'.
volume_size (int): Size in GB of the EBS volume to use for
storing input data during training (default: 30). Must be large
enough to store training data if File Mode is used (which is the
default).
volume_kms_key (str): Optional. KMS key ID for encrypting EBS
volume attached to the training instance (default: None).
max_run (int): Timeout in seconds for training (default: 24 *
60 * 60). After this amount of time Amazon SageMaker terminates
the job regardless of its current status.
input_mode (str): The input mode that the algorithm supports
(default: 'File'). Valid modes:
'File' - Amazon SageMaker copiesthe training dataset from the
S3 location to a local directory.
'Pipe' - Amazon SageMaker streams data directly from S3 to the
container via a Unix-named pipe.
'FastFile' - Amazon SageMaker streams data from S3 on demand instead of
downloading the entire dataset before training begins. This argument can
be overriden on a per-channel basis using
``sagemaker.inputs.TrainingInput.input_mode``.
output_path (str): S3 location for saving the training result (model
artifacts and output files). If not specified, results are
stored to a default bucket. If the bucket with the specific name
does not exist, the estimator creates the bucket during the
:meth:`~sagemaker.estimator.EstimatorBase.fit` method execution.
file:// urls are used for local mode. For example: 'file://model/'
will save to the model folder in the current directory.
output_kms_key (str): Optional. KMS key ID for encrypting the
training output (default: Your IAM role's KMS key for Amazon S3).
If you don't provide a KMS key ID, Amazon SageMaker uses the
default KMS key for Amazon S3 of the account linked to your
IAM role.
base_job_name (str): Prefix for training job name when the
:meth:`~sagemaker.estimator.EstimatorBase.fit` method launches.
If not specified, the estimator generates a default job name
based on the training image name and current timestamp.
sagemaker_session (sagemaker.session.Session): Session object which
manages interactions with Amazon SageMaker APIs and any other
AWS services needed. If not specified, the estimator creates one
using the default AWS configuration chain.
tags (list[dict]): List of tags for labeling a training job. For
more, see
https://docs.aws.amazon.com/sagemaker/latest/dg/API_Tag.html.
subnets (list[str]): List of subnet ids. If not specified training
job will be created without VPC config.
security_group_ids (list[str]): List of security group ids. If not
specified training job will be created without VPC config.
model_uri (str): URI where a pre-trained model is stored, either
locally or in S3 (default: None). If specified, the estimator
will create a channel pointing to the model so the training job
can download it. This model can be a 'model.tar.gz' from a
previous training job, or other artifacts coming from a
different source.
In local mode, this should point to the path in which the model
is located and not the file itself, as local Docker containers
will try to mount the URI as a volume.
More information:
https://docs.aws.amazon.com/sagemaker/latest/dg/cdf-training.html#td-deserialization
model_channel_name (str): Name of the channel where 'model_uri' will
be downloaded (default: 'model').
metric_definitions (list[dict]): A list of dictionaries that defines
the metric(s) used to evaluate the training jobs. Each
dictionary contains two keys: 'Name' for the name of the metric,
and 'Regex' for the regular expression used to extract the
metric from the logs. This should be defined only for jobs that
don't use an Amazon algorithm.
encrypt_inter_container_traffic (bool): Specifies whether traffic
between training containers is encrypted for the training job
(default: ``False``).
use_spot_instances (bool): Specifies whether to use SageMaker
Managed Spot instances for training. If enabled then the
``max_wait`` arg should also be set.
More information:
https://docs.aws.amazon.com/sagemaker/latest/dg/model-managed-spot-training.html
(default: ``False``).
max_wait (int): Timeout in seconds waiting for spot training
job (default: None). After this amount of time Amazon
SageMaker will stop waiting for managed spot training job to
complete (default: ``None``).
checkpoint_s3_uri (str): The S3 URI in which to persist checkpoints
that the algorithm persists (if any) during training. (default:
``None``).
checkpoint_local_path (str): The local path that the algorithm
writes its checkpoints to. SageMaker will persist all files
under this path to `checkpoint_s3_uri` continually during
training. On job startup the reverse happens - data from the
s3 location is downloaded to this path before the algorithm is
started. If the path is unset then SageMaker assumes the
checkpoints will be provided under `/opt/ml/checkpoints/`.
(default: ``None``).
rules (list[:class:`~sagemaker.debugger.RuleBase`]): A list of
:class:`~sagemaker.debugger.RuleBase` objects used to define
SageMaker Debugger rules for real-time analysis
(default: ``None``). For more information,
see `Continuous analyses through rules
<https://sagemaker.readthedocs.io/en/stable/amazon_sagemaker_debugger.html
#continuous-analyses-through-rules)>`_.
debugger_hook_config (:class:`~sagemaker.debugger.DebuggerHookConfig` or bool):
Configuration for how debugging information is emitted with
SageMaker Debugger. If not specified, a default one is created using
the estimator's ``output_path``, unless the region does not
support SageMaker Debugger. To disable SageMaker Debugger,
set this parameter to ``False``. For more information, see
`Capture real-time debugging data during model training in Amazon SageMaker
<https://sagemaker.readthedocs.io/en/stable/amazon_sagemaker_debugger.html#
capture-real-time-debugging-data-during-model-training-in-amazon-sagemaker>`_.
tensorboard_output_config (:class:`~sagemaker.debugger.TensorBoardOutputConfig`):
Configuration for customizing debugging visualization using TensorBoard
(default: ``None``). For more information,
see `Capture real time tensorboard data
<https://sagemaker.readthedocs.io/en/stable/amazon_sagemaker_debugger.html#
capture-real-time-tensorboard-data-from-the-debugging-hook>`_.
enable_sagemaker_metrics (bool): enable SageMaker Metrics Time
Series. For more information, see `AlgorithmSpecification API
<https://docs.aws.amazon.com/sagemaker/latest/dg/
API_AlgorithmSpecification.html#SageMaker-Type-AlgorithmSpecification-
EnableSageMakerMetricsTimeSeries>`_.
(default: ``None``).
enable_network_isolation (bool): Specifies whether container will
run in network isolation mode (default: ``False``). Network
isolation mode restricts the container access to outside networks
(such as the Internet). The container does not make any inbound or
outbound network calls. Also known as Internet-free mode.
profiler_config (:class:`~sagemaker.debugger.ProfilerConfig`):
Configuration for how SageMaker Debugger collects
monitoring and profiling information from your training job.
If not specified, a default configuration is created using
the estimator's ``output_path``, unless the region does not
support SageMaker Debugger. To disable SageMaker Debugger
monitoring and profiling, set the
``disable_profiler`` parameter to ``True``.
disable_profiler (bool): Specifies whether Debugger monitoring and profiling
will be disabled (default: ``False``).
environment (dict[str, str]) : Environment variables to be set for
use during training job (default: ``None``)
max_retry_attempts (int): The number of times to move a job to the STARTING status.
You can specify between 1 and 30 attempts.
If the value of attempts is greater than zero,
the job is retried on InternalServerFailure
the same number of attempts as the value.
You can cap the total duration for your job by setting ``max_wait`` and ``max_run``
(default: ``None``)
"""
instance_count = renamed_kwargs(
"train_instance_count", "instance_count", instance_count, kwargs
)
instance_type = renamed_kwargs(
"train_instance_type", "instance_type", instance_type, kwargs
)
max_run = renamed_kwargs("train_max_run", "max_run", max_run, kwargs)
use_spot_instances = renamed_kwargs(
"train_use_spot_instances", "use_spot_instances", use_spot_instances, kwargs
)
max_wait = renamed_kwargs("train_max_wait", "max_wait", max_wait, kwargs)
volume_size = renamed_kwargs("train_volume_size", "volume_size", volume_size, kwargs)
volume_kms_key = renamed_kwargs(
"train_volume_kms_key", "volume_kms_key", volume_kms_key, kwargs
)
if instance_count is None or instance_type is None:
raise ValueError("Both instance_count and instance_type are required.")
self.role = role
self.instance_count = instance_count
self.instance_type = instance_type
self.volume_size = volume_size
self.volume_kms_key = volume_kms_key
self.max_run = max_run
self.input_mode = input_mode
self.tags = tags
self.metric_definitions = metric_definitions
self.model_uri = model_uri
self.model_channel_name = model_channel_name
self.code_uri = None
self.code_channel_name = "code"
if self.instance_type in ("local", "local_gpu"):
if self.instance_type == "local_gpu" and self.instance_count > 1:
raise RuntimeError("Distributed Training in Local GPU is not supported")
self.sagemaker_session = sagemaker_session or LocalSession()
if not isinstance(self.sagemaker_session, sagemaker.local.LocalSession):
raise RuntimeError(
"instance_type local or local_gpu is only supported with an"
"instance of LocalSession"
)
else:
self.sagemaker_session = sagemaker_session or Session()
self.base_job_name = base_job_name
self._current_job_name = None
if (
not self.sagemaker_session.local_mode
and output_path
and output_path.startswith("file://")
):
raise RuntimeError("file:// output paths are only supported in Local Mode")
self.output_path = output_path
self.output_kms_key = output_kms_key
self.latest_training_job = None
self.jobs = []
self.deploy_instance_type = None
self._compiled_models = {}
# VPC configurations
self.subnets = subnets
self.security_group_ids = security_group_ids
self.encrypt_inter_container_traffic = encrypt_inter_container_traffic
self.use_spot_instances = use_spot_instances
self.max_wait = max_wait
self.checkpoint_s3_uri = checkpoint_s3_uri
self.checkpoint_local_path = checkpoint_local_path
self.rules = rules
self.debugger_hook_config = debugger_hook_config
self.tensorboard_output_config = tensorboard_output_config
self.debugger_rule_configs = None
self.collection_configs = None
self.enable_sagemaker_metrics = enable_sagemaker_metrics
self._enable_network_isolation = enable_network_isolation
self.profiler_config = profiler_config
self.disable_profiler = disable_profiler
self.environment = environment
self.max_retry_attempts = max_retry_attempts
if not _region_supports_profiler(self.sagemaker_session.boto_region_name):
self.disable_profiler = True
self.profiler_rule_configs = None
self.profiler_rules = None
self.debugger_rules = None
@abstractmethod
def training_image_uri(self):
"""Return the Docker image to use for training.
The :meth:`~sagemaker.estimator.EstimatorBase.fit` method, which does
the model training, calls this method to find the image to use for model
training.
Returns:
str: The URI of the Docker image.
"""
@abstractmethod
def hyperparameters(self):
"""Return the hyperparameters as a dictionary to use for training.
The :meth:`~sagemaker.estimator.EstimatorBase.fit` method, which
trains the model, calls this method to find the hyperparameters.
Returns:
dict[str, str]: The hyperparameters.
"""
def enable_network_isolation(self):
"""Return True if this Estimator will need network isolation to run.
Returns:
bool: Whether this Estimator needs network isolation or not.
"""
return self._enable_network_isolation
def prepare_workflow_for_training(self, job_name=None):
"""Calls _prepare_for_training. Used when setting up a workflow.
Args:
job_name (str): Name of the training job to be created. If not
specified, one is generated, using the base name given to the
constructor if applicable.
"""
self._prepare_for_training(job_name=job_name)
def _ensure_base_job_name(self):
"""Set ``self.base_job_name`` if it is not set already."""
# honor supplied base_job_name or generate it
if self.base_job_name is None:
self.base_job_name = base_name_from_image(self.training_image_uri())
def _get_or_create_name(self, name=None):
"""Generate a name based on the base job name or training image if needed.
Args:
name (str): User-supplied name. If not specified, a name is generated from
the base job name or training image.
Returns:
str: Either the user-supplied name or a generated name.
"""
if name:
return name
self._ensure_base_job_name()
return name_from_base(self.base_job_name)
def _prepare_for_training(self, job_name=None):
"""Set any values in the estimator that need to be set before training.
Args:
job_name (str): Name of the training job to be created. If not
specified, one is generated, using the base name given to the
constructor if applicable.
"""
self._current_job_name = self._get_or_create_name(job_name)
# if output_path was specified we use it otherwise initialize here.
# For Local Mode with local_code=True we don't need an explicit output_path
if self.output_path is None:
local_code = get_config_value("local.local_code", self.sagemaker_session.config)
if self.sagemaker_session.local_mode and local_code:
self.output_path = ""
else:
self.output_path = "s3://{}/".format(self.sagemaker_session.default_bucket())
self._prepare_rules()
self._prepare_debugger_for_training()
self._prepare_profiler_for_training()
def _prepare_rules(self):
"""Rules list includes both debugger and profiler rules.
Customer can explicitly disable any rule by setting rules to an empty list.
"""
self.debugger_rules = []
self.profiler_rules = []
if self.rules is not None:
for rule in self.rules:
if isinstance(rule, Rule):
self.debugger_rules.append(rule)
elif isinstance(rule, ProfilerRule):
self.profiler_rules.append(rule)
else:
raise RuntimeError(
"Rules list can only contain sagemaker.debugger.Rule "
+ "and sagemaker.debugger.ProfilerRule"
)
def _prepare_debugger_for_training(self):
"""Prepare debugger rules and debugger configs for training."""
if self.debugger_rules and self.debugger_hook_config is None:
self.debugger_hook_config = DebuggerHookConfig(s3_output_path=self.output_path)
# If debugger_hook_config was provided without an S3 URI, default it for the customer.
if self.debugger_hook_config and not self.debugger_hook_config.s3_output_path:
self.debugger_hook_config.s3_output_path = self.output_path
self.debugger_rule_configs = self._prepare_debugger_rules()
self._prepare_collection_configs()
def _prepare_debugger_rules(self):
"""Set any necessary values in debugger rules, if they are provided."""
debugger_rule_configs = []
if self.debugger_rules:
for rule in self.debugger_rules:
self._set_default_rule_config(rule)
self._set_source_s3_uri(rule)
rule.prepare_actions(self._current_job_name)
debugger_rule_configs.append(rule.to_debugger_rule_config_dict())
return debugger_rule_configs
def _prepare_collection_configs(self):
"""De-duplicate configurations and save them in the debugger hook configuration."""
# Create a set to de-duplicate CollectionConfigs.
self.collection_configs = set()
# Iterate through the debugger rules and add their respective CollectionConfigs to the set.
if self.debugger_rules:
for rule in self.debugger_rules:
self.collection_configs.update(rule.collection_configs)
# Add the CollectionConfigs from DebuggerHookConfig to the set.
if self.debugger_hook_config:
self.collection_configs.update(self.debugger_hook_config.collection_configs or [])
def _prepare_profiler_for_training(self):
"""Set necessary values and do basic validations in profiler config and profiler rules.
When user explicitly set rules to an empty list, default profiler rule won't be enabled.
Default profiler rule will be enabled in supported regions when either:
1. user doesn't specify any rules, i.e., rules=None; or
2. user only specify debugger rules, i.e., rules=[Rule.sagemaker(...)]
"""
if self.disable_profiler:
if self.profiler_config:
raise RuntimeError("profiler_config cannot be set when disable_profiler is True.")
if self.profiler_rules:
raise RuntimeError("ProfilerRule cannot be set when disable_profiler is True.")
elif _region_supports_profiler(self.sagemaker_session.boto_region_name):
if self.profiler_config is None:
self.profiler_config = ProfilerConfig(s3_output_path=self.output_path)
if self.rules is None or (self.rules and not self.profiler_rules):
self.profiler_rules = [get_default_profiler_rule()]
if self.profiler_config and not self.profiler_config.s3_output_path:
self.profiler_config.s3_output_path = self.output_path
self.profiler_rule_configs = self._prepare_profiler_rules()
def _prepare_profiler_rules(self):
"""Set any necessary values in profiler rules, if they are provided."""
profiler_rule_configs = []
if self.profiler_rules:
for rule in self.profiler_rules:
self._set_default_rule_config(rule)
self._set_source_s3_uri(rule)
profiler_rule_configs.append(rule.to_profiler_rule_config_dict())
return profiler_rule_configs
def _set_default_rule_config(self, rule):
"""Set default rule configurations.
Args:
rule (:class:`~sagemaker.debugger.RuleBase`): Any rule object that derives from RuleBase
"""
if rule.image_uri == "DEFAULT_RULE_EVALUATOR_IMAGE":
rule.image_uri = get_rule_container_image_uri(self.sagemaker_session.boto_region_name)
rule.instance_type = None
rule.volume_size_in_gb = None
def _set_source_s3_uri(self, rule):
"""Set updated source S3 uri when specified.
Args:
rule (:class:`~sagemaker.debugger.RuleBase`): Any rule object that derives from RuleBase
"""
if "source_s3_uri" in (rule.rule_parameters or {}):
parse_result = urlparse(rule.rule_parameters["source_s3_uri"])
if parse_result.scheme != "s3":
desired_s3_uri = os.path.join(
"s3://",
self.sagemaker_session.default_bucket(),
rule.name,
str(uuid.uuid4()),
)
s3_uri = S3Uploader.upload(
local_path=rule.rule_parameters["source_s3_uri"],
desired_s3_uri=desired_s3_uri,
sagemaker_session=self.sagemaker_session,
)
rule.rule_parameters["source_s3_uri"] = s3_uri
def latest_job_debugger_artifacts_path(self):
"""Gets the path to the DebuggerHookConfig output artifacts.
Returns:
str: An S3 path to the output artifacts.
"""
self._ensure_latest_training_job(
error_message="""Cannot get the Debugger artifacts path.
The Estimator is not associated with a training job."""
)
if self.debugger_hook_config is not None:
return os.path.join(
self.debugger_hook_config.s3_output_path,
self.latest_training_job.name,
"debug-output",
)
return None
def latest_job_tensorboard_artifacts_path(self):
"""Gets the path to the TensorBoardOutputConfig output artifacts.
Returns:
str: An S3 path to the output artifacts.
"""
self._ensure_latest_training_job(
error_message="""Cannot get the TensorBoard artifacts path.
The Estimator is not associated with a training job."""
)
if self.debugger_hook_config is not None:
return os.path.join(
self.tensorboard_output_config.s3_output_path,
self.latest_training_job.name,
"tensorboard-output",
)
return None
def latest_job_profiler_artifacts_path(self):
"""Gets the path to the profiling output artifacts.
Returns:
str: An S3 path to the output artifacts.
"""
self._ensure_latest_training_job(
error_message="""Cannot get the profiling output artifacts path.
The Estimator is not associated with a training job."""
)
if self.profiler_config is not None:
return os.path.join(
self.profiler_config.s3_output_path,
self.latest_training_job.name,
"profiler-output",
)
return None
def fit(self, inputs=None, wait=True, logs="All", job_name=None, experiment_config=None):
"""Train a model using the input training dataset.
The API calls the Amazon SageMaker CreateTrainingJob API to start
model training. The API uses configuration you provided to create the
estimator and the specified input training data to send the
CreatingTrainingJob request to Amazon SageMaker.
This is a synchronous operation. After the model training
successfully completes, you can call the ``deploy()`` method to host the
model using the Amazon SageMaker hosting services.
Args:
inputs (str or dict or sagemaker.inputs.TrainingInput or
sagemaker.inputs.FileSystemInput): Information about the training data.
This can be one of four types:
* (str) the S3 location where training data is saved, or a file:// path in
local mode.
* (dict[str, str] or dict[str, sagemaker.inputs.TrainingInput] or
dict[str, sagemaker.inputs.FileSystemInput]) If using multiple channels for
training data, you can specify a dict mapping channel names to strings or
:func:`~sagemaker.inputs.TrainingInput` objects or
:func:`~sagemaker.inputs.FileSystemInput` objects.
* (sagemaker.inputs.TrainingInput) - channel configuration for S3 data sources
that can provide additional information as well as the path to the training
dataset.
See :func:`sagemaker.inputs.TrainingInput` for full details.
* (sagemaker.inputs.FileSystemInput) - channel configuration for
a file system data source that can provide additional information as well as
the path to the training dataset.
wait (bool): Whether the call should wait until the job completes (default: True).
logs ([str]): A list of strings specifying which logs to print. Acceptable
strings are "All", "None", "Training", or "Rules". To maintain backwards
compatibility, boolean values are also accepted and converted to strings.
Only meaningful when wait is True.
job_name (str): Training job name. If not specified, the estimator generates
a default job name based on the training image name and current timestamp.
experiment_config (dict[str, str]): Experiment management configuration.
Optionally, the dict can contain three keys:
'ExperimentName', 'TrialName', and 'TrialComponentDisplayName'.
The behavior of setting these keys is as follows:
* If `ExperimentName` is supplied but `TrialName` is not a Trial will be
automatically created and the job's Trial Component associated with the Trial.
* If `TrialName` is supplied and the Trial already exists the job's Trial Component
will be associated with the Trial.
* If both `ExperimentName` and `TrialName` are not supplied the trial component
will be unassociated.
* `TrialComponentDisplayName` is used for display in Studio.
"""
self._prepare_for_training(job_name=job_name)
self.latest_training_job = _TrainingJob.start_new(self, inputs, experiment_config)
self.jobs.append(self.latest_training_job)
if wait:
self.latest_training_job.wait(logs=logs)
def _compilation_job_name(self):
"""Placeholder docstring"""
base_name = self.base_job_name or base_name_from_image(self.training_image_uri())
return name_from_base("compilation-" + base_name)
def compile_model(
self,
target_instance_family,
input_shape,
output_path,
framework=None,
framework_version=None,
compile_max_run=15 * 60,
tags=None,
target_platform_os=None,
target_platform_arch=None,
target_platform_accelerator=None,
compiler_options=None,
**kwargs,
):
"""Compile a Neo model using the input model.
Args:
target_instance_family (str): Identifies the device that you want to
run your model after compilation, for example: ml_c5. For allowed
strings see
https://docs.aws.amazon.com/sagemaker/latest/dg/API_OutputConfig.html.
input_shape (dict): Specifies the name and shape of the expected
inputs for your trained model in json dictionary form, for
example: {'data':[1,3,1024,1024]}, or {'var1': [1,1,28,28],
'var2':[1,1,28,28]}
output_path (str): Specifies where to store the compiled model
framework (str): The framework that is used to train the original
model. Allowed values: 'mxnet', 'tensorflow', 'keras', 'pytorch',
'onnx', 'xgboost'
framework_version (str): The version of the framework
compile_max_run (int): Timeout in seconds for compilation (default:
15 * 60). After this amount of time Amazon SageMaker Neo
terminates the compilation job regardless of its current status.
tags (list[dict]): List of tags for labeling a compilation job. For
more, see
https://docs.aws.amazon.com/sagemaker/latest/dg/API_Tag.html.
target_platform_os (str): Target Platform OS, for example: 'LINUX'.
For allowed strings see
https://docs.aws.amazon.com/sagemaker/latest/dg/API_OutputConfig.html.
It can be used instead of target_instance_family.
target_platform_arch (str): Target Platform Architecture, for example: 'X86_64'.
For allowed strings see
https://docs.aws.amazon.com/sagemaker/latest/dg/API_OutputConfig.html.
It can be used instead of target_instance_family.
target_platform_accelerator (str, optional): Target Platform Accelerator,
for example: 'NVIDIA'. For allowed strings see
https://docs.aws.amazon.com/sagemaker/latest/dg/API_OutputConfig.html.
It can be used instead of target_instance_family.
compiler_options (dict, optional): Additional parameters for compiler.
Compiler Options are TargetPlatform / target_instance_family specific. See
https://docs.aws.amazon.com/sagemaker/latest/dg/API_OutputConfig.html for details.
**kwargs: Passed to invocation of ``create_model()``.
Implementations may customize ``create_model()`` to accept
``**kwargs`` to customize model creation during deploy. For
more, see the implementation docs.
Returns:
sagemaker.model.Model: A SageMaker ``Model`` object. See
:func:`~sagemaker.model.Model` for full details.
"""
if framework and framework not in NEO_ALLOWED_FRAMEWORKS:
raise ValueError(
"Please use valid framework, allowed values: {}".format(NEO_ALLOWED_FRAMEWORKS)
)
if (framework is None) != (framework_version is None):
raise ValueError("You should provide framework and framework_version at the same time.")
model = self.create_model(**kwargs)
self._compiled_models[target_instance_family] = model.compile(
target_instance_family,
input_shape,
output_path,
self.role,
tags,
self._compilation_job_name(),
compile_max_run,
framework=framework,
framework_version=framework_version,
target_platform_os=target_platform_os,
target_platform_arch=target_platform_arch,
target_platform_accelerator=target_platform_accelerator,
compiler_options=compiler_options,
)
return self._compiled_models[target_instance_family]
@classmethod
def attach(cls, training_job_name, sagemaker_session=None, model_channel_name="model"):
"""Attach to an existing training job.
Create an Estimator bound to an existing training job, each subclass
is responsible to implement
``_prepare_init_params_from_job_description()`` as this method delegates
the actual conversion of a training job description to the arguments
that the class constructor expects. After attaching, if the training job
has a Complete status, it can be ``deploy()`` ed to create a SageMaker
Endpoint and return a ``Predictor``.
If the training job is in progress, attach will block until the training job
completes, but logs of the training job will not display. To see the logs
content, please call ``logs()``
Examples:
>>> my_estimator.fit(wait=False)
>>> training_job_name = my_estimator.latest_training_job.name
Later on:
>>> attached_estimator = Estimator.attach(training_job_name)
>>> attached_estimator.logs()
>>> attached_estimator.deploy()
Args:
training_job_name (str): The name of the training job to attach to.
sagemaker_session (sagemaker.session.Session): Session object which
manages interactions with Amazon SageMaker APIs and any other
AWS services needed. If not specified, the estimator creates one
using the default AWS configuration chain.
model_channel_name (str): Name of the channel where pre-trained
model data will be downloaded (default: 'model'). If no channel
with the same name exists in the training job, this option will
be ignored.
Returns:
Instance of the calling ``Estimator`` Class with the attached
training job.
"""
sagemaker_session = sagemaker_session or Session()
job_details = sagemaker_session.sagemaker_client.describe_training_job(
TrainingJobName=training_job_name
)
init_params = cls._prepare_init_params_from_job_description(job_details, model_channel_name)
tags = sagemaker_session.sagemaker_client.list_tags(
ResourceArn=job_details["TrainingJobArn"]
)["Tags"]
init_params.update(tags=tags)
estimator = cls(sagemaker_session=sagemaker_session, **init_params)
estimator.latest_training_job = _TrainingJob(
sagemaker_session=sagemaker_session, job_name=training_job_name
)
estimator._current_job_name = estimator.latest_training_job.name
estimator.latest_training_job.wait(logs="None")
return estimator
def logs(self):
"""Display the logs for Estimator's training job.
If the output is a tty or a Jupyter cell, it will be color-coded based
on which instance the log entry is from.
"""
self.sagemaker_session.logs_for_job(self.latest_training_job.name, wait=True)
def deploy(
self,
initial_instance_count,
instance_type,
serializer=None,
deserializer=None,
accelerator_type=None,
endpoint_name=None,
use_compiled_model=False,
wait=True,
model_name=None,
kms_key=None,
data_capture_config=None,
tags=None,
**kwargs,
):
"""Deploy the trained model to an Amazon SageMaker endpoint.
And then return ``sagemaker.Predictor`` object.
More information:
http://docs.aws.amazon.com/sagemaker/latest/dg/how-it-works-training.html
Args:
initial_instance_count (int): Minimum number of EC2 instances to
deploy to an endpoint for prediction.
instance_type (str): Type of EC2 instance to deploy to an endpoint
for prediction, for example, 'ml.c4.xlarge'.
serializer (:class:`~sagemaker.serializers.BaseSerializer`): A
serializer object, used to encode data for an inference endpoint
(default: None). If ``serializer`` is not None, then
``serializer`` will override the default serializer. The
default serializer is set by the ``predictor_cls``.
deserializer (:class:`~sagemaker.deserializers.BaseDeserializer`): A
deserializer object, used to decode data from an inference
endpoint (default: None). If ``deserializer`` is not None, then
``deserializer`` will override the default deserializer. The
default deserializer is set by the ``predictor_cls``.
accelerator_type (str): Type of Elastic Inference accelerator to
attach to an endpoint for model loading and inference, for
example, 'ml.eia1.medium'. If not specified, no Elastic
Inference accelerator will be attached to the endpoint. For more
information:
https://docs.aws.amazon.com/sagemaker/latest/dg/ei.html
endpoint_name (str): Name to use for creating an Amazon SageMaker
endpoint. If not specified, the name of the training job is
used.
use_compiled_model (bool): Flag to select whether to use compiled
(optimized) model. Default: False.
wait (bool): Whether the call should wait until the deployment of
model completes (default: True).
model_name (str): Name to use for creating an Amazon SageMaker
model. If not specified, the estimator generates a default job name
based on the training image name and current timestamp.
kms_key (str): The ARN of the KMS key that is used to encrypt the
data on the storage volume attached to the instance hosting the
endpoint.
data_capture_config (sagemaker.model_monitor.DataCaptureConfig): Specifies
configuration related to Endpoint data capture for use with
Amazon SageMaker Model Monitoring. Default: None.
tags(List[dict[str, str]]): Optional. The list of tags to attach to this specific
endpoint. Example:
>>> tags = [{'Key': 'tagname', 'Value': 'tagvalue'}]
For more information about tags, see
https://boto3.amazonaws.com/v1/documentation\
/api/latest/reference/services/sagemaker.html#SageMaker.Client.add_tags
**kwargs: Passed to invocation of ``create_model()``.
Implementations may customize ``create_model()`` to accept
``**kwargs`` to customize model creation during deploy.
For more, see the implementation docs.
Returns:
sagemaker.predictor.Predictor: A predictor that provides a ``predict()`` method,
which can be used to send requests to the Amazon SageMaker
endpoint and obtain inferences.
"""
removed_kwargs("update_endpoint", kwargs)
self._ensure_latest_training_job()
self._ensure_base_job_name()
default_name = name_from_base(self.base_job_name)
endpoint_name = endpoint_name or default_name
model_name = model_name or default_name
self.deploy_instance_type = instance_type
if use_compiled_model:
family = "_".join(instance_type.split(".")[:-1])
if family not in self._compiled_models:
raise ValueError(
"No compiled model for {}. "
"Please compile one with compile_model before deploying.".format(family)
)
model = self._compiled_models[family]
else:
kwargs["model_kms_key"] = self.output_kms_key
model = self.create_model(**kwargs)
model.name = model_name
return model.deploy(
instance_type=instance_type,
initial_instance_count=initial_instance_count,
serializer=serializer,
deserializer=deserializer,
accelerator_type=accelerator_type,
endpoint_name=endpoint_name,
tags=tags or self.tags,
wait=wait,
kms_key=kms_key,
data_capture_config=data_capture_config,
)
def register(
self,
content_types,
response_types,
inference_instances,
transform_instances,
image_uri=None,
model_package_name=None,
model_package_group_name=None,
model_metrics=None,
metadata_properties=None,
marketplace_cert=False,
approval_status=None,
description=None,
compile_model_family=None,
model_name=None,
drift_check_baselines=None,
**kwargs,
):
"""Creates a model package for creating SageMaker models or listing on Marketplace.
Args:
content_types (list): The supported MIME types for the input data.
response_types (list): The supported MIME types for the output data.
inference_instances (list): A list of the instance types that are used to
generate inferences in real-time.
transform_instances (list): A list of the instance types on which a transformation
job can be run or on which an endpoint can be deployed.
image_uri (str): The container image uri for Model Package, if not specified,
Estimator's training container image will be used (default: None).
model_package_name (str): Model Package name, exclusive to `model_package_group_name`,
using `model_package_name` makes the Model Package un-versioned (default: None).
model_package_group_name (str): Model Package Group name, exclusive to
`model_package_name`, using `model_package_group_name` makes the Model Package
versioned (default: None).
model_metrics (ModelMetrics): ModelMetrics object (default: None).
metadata_properties (MetadataProperties): MetadataProperties (default: None).