Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/c7n/resources/sagemaker.py: 58%

1# Copyright The Cloud Custodian Authors.

2# SPDX-License-Identifier: Apache-2.0

4from c7n.actions import BaseAction

5from c7n.exceptions import PolicyValidationError

6from c7n.manager import resources

7from c7n.query import QueryResourceManager, TypeInfo, DescribeSource, ConfigSource

8from c7n.utils import local_session, type_schema

9from c7n.tags import RemoveTag, Tag, TagActionFilter, TagDelayedAction

10from c7n.filters.vpc import SubnetFilter, SecurityGroupFilter

11from c7n.filters.kms import KmsRelatedFilter

12from c7n.filters.offhours import OffHour, OnHour

15@resources.register('sagemaker-notebook')

16class NotebookInstance(QueryResourceManager):

18 class resource_type(TypeInfo):

19 service = 'sagemaker'

20 enum_spec = ('list_notebook_instances', 'NotebookInstances', None)

21 detail_spec = (

22 'describe_notebook_instance', 'NotebookInstanceName',

23 'NotebookInstanceName', None)

24 arn = id = 'NotebookInstanceArn'

25 name = 'NotebookInstanceName'

26 date = 'CreationTime'

27 cfn_type = 'AWS::SageMaker::NotebookInstance'

29 permissions = ('sagemaker:ListTags',)

31 def augment(self, resources):

32 client = local_session(self.session_factory).client('sagemaker')

34 def _augment(r):

35 # List tags for the Notebook-Instance & set as attribute

36 tags = self.retry(client.list_tags,

37 ResourceArn=r['NotebookInstanceArn'])['Tags']

38 r['Tags'] = tags

39 return r

41 # Describe notebook-instance & then list tags

42 resources = super(NotebookInstance, self).augment(resources)

43 return list(map(_augment, resources))

46NotebookInstance.filter_registry.register('marked-for-op', TagActionFilter)

47NotebookInstance.filter_registry.register('offhour', OffHour)

48NotebookInstance.filter_registry.register('onhour', OnHour)

51@resources.register('sagemaker-job')

52class SagemakerJob(QueryResourceManager):

54 class resource_type(TypeInfo):

55 service = 'sagemaker'

56 enum_spec = ('list_training_jobs', 'TrainingJobSummaries', None)

57 detail_spec = (

58 'describe_training_job', 'TrainingJobName', 'TrainingJobName', None)

59 arn = id = 'TrainingJobArn'

60 name = 'TrainingJobName'

61 date = 'CreationTime'

62 permission_augment = (

63 'sagemaker:DescribeTrainingJob', 'sagemaker:ListTags')

65 def __init__(self, ctx, data):

66 super(SagemakerJob, self).__init__(ctx, data)

67 self.queries = QueryFilter.parse(

68 self.data.get('query', [

69 {'StatusEquals': 'InProgress'}]))

71 def resources(self, query=None):

72 for q in self.queries:

73 if q is None:

74 continue

75 query = query or {}

76 for k, v in q.items():

77 query[k] = v

78 return super(SagemakerJob, self).resources(query=query)

80 def augment(self, jobs):

81 client = local_session(self.session_factory).client('sagemaker')

83 def _augment(j):

84 tags = self.retry(client.list_tags,

85 ResourceArn=j['TrainingJobArn'])['Tags']

86 j['Tags'] = tags

87 return j

89 jobs = super(SagemakerJob, self).augment(jobs)

90 return list(map(_augment, jobs))

93@resources.register('sagemaker-transform-job')

94class SagemakerTransformJob(QueryResourceManager):

96 class resource_type(TypeInfo):

97 arn_type = "transform-job"

98 service = 'sagemaker'

99 enum_spec = ('list_transform_jobs', 'TransformJobSummaries', None)

100 detail_spec = (

101 'describe_transform_job', 'TransformJobName', 'TransformJobName', None)

102 arn = id = 'TransformJobArn'

103 name = 'TransformJobName'

104 date = 'CreationTime'

105 filter_name = 'NameContains'

106 filter_type = 'scalar'

107 permission_augment = ('sagemaker:DescribeTransformJob', 'sagemaker:ListTags')

108

109 def __init__(self, ctx, data):

110 super(SagemakerTransformJob, self).__init__(ctx, data)

111 self.queries = QueryFilter.parse(

112 self.data.get('query', [

113 {'StatusEquals': 'InProgress'}]))

114

115 def resources(self, query=None):

116 for q in self.queries:

117 if q is None:

118 continue

119 query = query or {}

120 for k, v in q.items():

121 query[k] = v

122 return super(SagemakerTransformJob, self).resources(query=query)

123

124 def augment(self, jobs):

125 client = local_session(self.session_factory).client('sagemaker')

126

127 def _augment(j):

128 tags = self.retry(client.list_tags,

129 ResourceArn=j['TransformJobArn'])['Tags']

130 j['Tags'] = tags

131 return j

132

133 return list(map(_augment, super(SagemakerTransformJob, self).augment(jobs)))

134

135

136class QueryFilter:

137

138 JOB_FILTERS = ('StatusEquals', 'NameContains',)

139

140 @classmethod

141 def parse(cls, data):

142 results = []

143 names = set()

144 for d in data:

145 if not isinstance(d, dict):

146 raise PolicyValidationError(

147 "Job Query Filter Invalid structure %s" % d)

148 for k, v in d.items():

149 if isinstance(v, list):

150 raise ValueError(

151 'Job query filter invalid structure %s' % v)

152 query = cls(d).validate().query()

153 if query['Name'] in names:

154 # Cannot filter multiple times on the same key

155 continue

156 names.add(query['Name'])

157 if isinstance(query['Value'], list):

158 results.append({query['Name']: query['Value'][0]})

159 continue

160 results.append({query['Name']: query['Value']})

161 if 'StatusEquals' not in names:

162 # add default StatusEquals if not included

163 results.append({'Name': 'StatusEquals', 'Value': 'InProgress'})

164 return results

165

166 def __init__(self, data):

167 self.data = data

168 self.key = None

169 self.value = None

170

171 def validate(self):

172 if not len(list(self.data.keys())) == 1:

173 raise PolicyValidationError(

174 "Job Query Filter Invalid %s" % self.data)

175 self.key = list(self.data.keys())[0]

176 self.value = list(self.data.values())[0]

177

178 if self.key not in self.JOB_FILTERS and not self.key.startswith('tag:'):

179 raise PolicyValidationError(

180 "Job Query Filter invalid filter name %s" % (

181 self.data))

182

183 if self.value is None:

184 raise PolicyValidationError(

185 "Job Query Filters must have a value, use tag-key"

186 " w/ tag name as value for tag present checks"

187 " %s" % self.data)

188 return self

189

190 def query(self):

191 value = self.value

192 if isinstance(self.value, str):

193 value = [self.value]

194 return {'Name': self.key, 'Value': value}

195

196

197@resources.register('sagemaker-endpoint')

198class SagemakerEndpoint(QueryResourceManager):

199

200 class resource_type(TypeInfo):

201 service = 'sagemaker'

202 enum_spec = ('list_endpoints', 'Endpoints', None)

203 detail_spec = (

204 'describe_endpoint', 'EndpointName',

205 'EndpointName', None)

206 arn = id = 'EndpointArn'

207 name = 'EndpointName'

208 date = 'CreationTime'

209 cfn_type = 'AWS::SageMaker::Endpoint'

210

211 permissions = ('sagemaker:ListTags',)

212

213 def augment(self, endpoints):

214 client = local_session(self.session_factory).client('sagemaker')

215

216 def _augment(e):

217 tags = self.retry(client.list_tags,

218 ResourceArn=e['EndpointArn'])['Tags']

219 e['Tags'] = tags

220 return e

221

222 # Describe endpoints & then list tags

223 endpoints = super(SagemakerEndpoint, self).augment(endpoints)

224 return list(map(_augment, endpoints))

225

226

227SagemakerEndpoint.filter_registry.register('marked-for-op', TagActionFilter)

228

229

230@resources.register('sagemaker-endpoint-config')

231class SagemakerEndpointConfig(QueryResourceManager):

232

233 class resource_type(TypeInfo):

234 service = 'sagemaker'

235 enum_spec = ('list_endpoint_configs', 'EndpointConfigs', None)

236 detail_spec = (

237 'describe_endpoint_config', 'EndpointConfigName',

238 'EndpointConfigName', None)

239 arn = id = 'EndpointConfigArn'

240 name = 'EndpointConfigName'

241 date = 'CreationTime'

242 cfn_type = 'AWS::SageMaker::EndpointConfig'

243

244 permissions = ('sagemaker:ListTags',)

245

246 def augment(self, endpoints):

247 client = local_session(self.session_factory).client('sagemaker')

248

249 def _augment(e):

250 tags = self.retry(client.list_tags,

251 ResourceArn=e['EndpointConfigArn'])['Tags']

252 e['Tags'] = tags

253 return e

254

255 endpoints = super(SagemakerEndpointConfig, self).augment(endpoints)

256 return list(map(_augment, endpoints))

257

258

259SagemakerEndpointConfig.filter_registry.register('marked-for-op', TagActionFilter)

260

261

262class DescribeModel(DescribeSource):

263

264 def augment(self, resources):

265 client = local_session(self.manager.session_factory).client('sagemaker')

266

267 def _augment(r):

268 tags = self.manager.retry(client.list_tags,

269 ResourceArn=r['ModelArn'])['Tags']

270 r.setdefault('Tags', []).extend(tags)

271 return r

272

273 resources = super(DescribeModel, self).augment(resources)

274 return list(map(_augment, resources))

275

276

277@resources.register('sagemaker-model')

278class Model(QueryResourceManager):

279

280 class resource_type(TypeInfo):

281 service = 'sagemaker'

282 enum_spec = ('list_models', 'Models', None)

283 detail_spec = (

284 'describe_model', 'ModelName',

285 'ModelName', None)

286 arn = id = 'ModelArn'

287 name = 'ModelName'

288 date = 'CreationTime'

289 cfn_type = config_type = 'AWS::SageMaker::Model'

290

291 source_mapping = {

292 'describe': DescribeModel,

293 'config': ConfigSource

294 }

295

296 permissions = ('sagemaker:ListTags',)

297

298

299Model.filter_registry.register('marked-for-op', TagActionFilter)

300

301

302@SagemakerEndpoint.action_registry.register('tag')

303@SagemakerEndpointConfig.action_registry.register('tag')

304@NotebookInstance.action_registry.register('tag')

305@SagemakerJob.action_registry.register('tag')

306@SagemakerTransformJob.action_registry.register('tag')

307@Model.action_registry.register('tag')

308class TagNotebookInstance(Tag):

309 """Action to create tag(s) on a SageMaker resource

310 (notebook-instance, endpoint, endpoint-config)

311

312 :example:

313

314 .. code-block:: yaml

315

316 policies:

317 - name: tag-sagemaker-notebook

318 resource: sagemaker-notebook

319 filters:

320 - "tag:target-tag": absent

321 actions:

322 - type: tag

323 key: target-tag

324 value: target-value

325

326 - name: tag-sagemaker-endpoint

327 resource: sagemaker-endpoint

328 filters:

329 - "tag:required-tag": absent

330 actions:

331 - type: tag

332 key: required-tag

333 value: required-value

334

335 - name: tag-sagemaker-endpoint-config

336 resource: sagemaker-endpoint-config

337 filters:

338 - "tag:required-tag": absent

339 actions:

340 - type: tag

341 key: required-tag

342 value: required-value

343

344 - name: tag-sagemaker-job

345 resource: sagemaker-job

346 filters:

347 - "tag:required-tag": absent

348 actions:

349 - type: tag

350 key: required-tag

351 value: required-value

352 """

353 permissions = ('sagemaker:AddTags',)

354

355 def process_resource_set(self, client, resources, tags):

356 mid = self.manager.resource_type.id

357 for r in resources:

358 client.add_tags(ResourceArn=r[mid], Tags=tags)

359

360

361@SagemakerEndpoint.action_registry.register('remove-tag')

362@SagemakerEndpointConfig.action_registry.register('remove-tag')

363@NotebookInstance.action_registry.register('remove-tag')

364@SagemakerJob.action_registry.register('remove-tag')

365@SagemakerTransformJob.action_registry.register('remove-tag')

366@Model.action_registry.register('remove-tag')

367class RemoveTagNotebookInstance(RemoveTag):

368 """Remove tag(s) from SageMaker resources

369 (notebook-instance, endpoint, endpoint-config)

370

371 :example:

372

373 .. code-block:: yaml

374

375 policies:

376 - name: sagemaker-notebook-remove-tag

377 resource: sagemaker-notebook

378 filters:

379 - "tag:BadTag": present

380 actions:

381 - type: remove-tag

382 tags: ["BadTag"]

383

384 - name: sagemaker-endpoint-remove-tag

385 resource: sagemaker-endpoint

386 filters:

387 - "tag:expired-tag": present

388 actions:

389 - type: remove-tag

390 tags: ["expired-tag"]

391

392 - name: sagemaker-endpoint-config-remove-tag

393 resource: sagemaker-endpoint-config

394 filters:

395 - "tag:expired-tag": present

396 actions:

397 - type: remove-tag

398 tags: ["expired-tag"]

399

400 - name: sagemaker-job-remove-tag

401 resource: sagemaker-job

402 filters:

403 - "tag:expired-tag": present

404 actions:

405 - type: remove-tag

406 tags: ["expired-tag"]

407 """

408 permissions = ('sagemaker:DeleteTags',)

409

410 def process_resource_set(self, client, resources, keys):

411 for r in resources:

412 client.delete_tags(ResourceArn=r[self.id_key], TagKeys=keys)

413

414

415@SagemakerEndpoint.action_registry.register('mark-for-op')

416@SagemakerEndpointConfig.action_registry.register('mark-for-op')

417@NotebookInstance.action_registry.register('mark-for-op')

418@Model.action_registry.register('mark-for-op')

419class MarkNotebookInstanceForOp(TagDelayedAction):

420 """Mark SageMaker resources for deferred action

421 (notebook-instance, endpoint, endpoint-config)

422

423 :example:

424

425 .. code-block:: yaml

426

427 policies:

428 - name: sagemaker-notebook-invalid-tag-stop

429 resource: sagemaker-notebook

430 filters:

431 - "tag:InvalidTag": present

432 actions:

433 - type: mark-for-op

434 op: stop

435 days: 1

436

437 - name: sagemaker-endpoint-failure-delete

438 resource: sagemaker-endpoint

439 filters:

440 - 'EndpointStatus': 'Failed'

441 actions:

442 - type: mark-for-op

443 op: delete

444 days: 1

445

446 - name: sagemaker-endpoint-config-invalid-size-delete

447 resource: sagemaker-notebook

448 filters:

449 - type: value

450 - key: ProductionVariants[].InstanceType

451 - value: 'ml.m4.10xlarge'

452 - op: contains

453 actions:

454 - type: mark-for-op

455 op: delete

456 days: 1

457 """

458

459

460@NotebookInstance.action_registry.register('start')

461class StartNotebookInstance(BaseAction):

462 """Start sagemaker-notebook(s)

463

464 :example:

465

466 .. code-block:: yaml

467

468 policies:

469 - name: start-sagemaker-notebook

470 resource: sagemaker-notebook

471 actions:

472 - start

473 """

474 schema = type_schema('start')

475 permissions = ('sagemaker:StartNotebookInstance',)

476 valid_origin_states = ('Stopped',)

477

478 def process(self, resources):

479 resources = self.filter_resources(resources, 'NotebookInstanceStatus',

480 self.valid_origin_states)

481 if not len(resources):

482 return

483

484 client = local_session(self.manager.session_factory).client('sagemaker')

485

486 for n in resources:

487 try:

488 client.start_notebook_instance(

489 NotebookInstanceName=n['NotebookInstanceName'])

490 except client.exceptions.ResourceNotFound:

491 pass

492

493

494@NotebookInstance.action_registry.register('stop')

495class StopNotebookInstance(BaseAction):

496 """Stop sagemaker-notebook(s)

497

498 :example:

499

500 .. code-block:: yaml

501

502 policies:

503 - name: stop-sagemaker-notebook

504 resource: sagemaker-notebook

505 filters:

506 - "tag:DeleteMe": present

507 actions:

508 - stop

509 """

510 schema = type_schema('stop')

511 permissions = ('sagemaker:StopNotebookInstance',)

512 valid_origin_states = ('InService',)

513

514 def process(self, resources):

515 resources = self.filter_resources(resources, 'NotebookInstanceStatus',

516 self.valid_origin_states)

517 if not len(resources):

518 return

519

520 client = local_session(self.manager.session_factory).client('sagemaker')

521

522 for n in resources:

523 try:

524 client.stop_notebook_instance(

525 NotebookInstanceName=n['NotebookInstanceName'])

526 except client.exceptions.ResourceNotFound:

527 pass

528

529

530@NotebookInstance.action_registry.register('delete')

531class DeleteNotebookInstance(BaseAction):

532 """Deletes sagemaker-notebook(s)

533

534 :example:

535

536 .. code-block:: yaml

537

538 policies:

539 - name: delete-sagemaker-notebook

540 resource: sagemaker-notebook

541 filters:

542 - "tag:DeleteMe": present

543 actions:

544 - delete

545 """

546 schema = type_schema('delete')

547 permissions = ('sagemaker:DeleteNotebookInstance',)

548 valid_origin_states = ('Stopped', 'Failed',)

549

550 def process(self, resources):

551 resources = self.filter_resources(resources, 'NotebookInstanceStatus',

552 self.valid_origin_states)

553 if not len(resources):

554 return

555

556 client = local_session(self.manager.session_factory).client('sagemaker')

557

558 for n in resources:

559 try:

560 client.delete_notebook_instance(

561 NotebookInstanceName=n['NotebookInstanceName'])

562 except client.exceptions.ResourceNotFound:

563 pass

564

565

566@NotebookInstance.filter_registry.register('security-group')

567class NotebookSecurityGroupFilter(SecurityGroupFilter):

568

569 RelatedIdsExpression = "SecurityGroups[]"

570

571

572@NotebookInstance.filter_registry.register('subnet')

573class NotebookSubnetFilter(SubnetFilter):

574

575 RelatedIdsExpression = "SubnetId"

576

577

578@NotebookInstance.filter_registry.register('kms-key')

579@SagemakerEndpointConfig.filter_registry.register('kms-key')

580class NotebookKmsFilter(KmsRelatedFilter):

581

582 RelatedIdsExpression = "KmsKeyId"

583

584

585@Model.action_registry.register('delete')

586class DeleteModel(BaseAction):

587 """Deletes sagemaker-model(s)

588

589 :example:

590

591 .. code-block:: yaml

592

593 policies:

594 - name: delete-sagemaker-model

595 resource: sagemaker-model

596 filters:

597 - "tag:DeleteMe": present

598 actions:

599 - delete

600 """

601 schema = type_schema('delete')

602 permissions = ('sagemaker:DeleteModel',)

603

604 def process(self, resources):

605 client = local_session(self.manager.session_factory).client('sagemaker')

606

607 for m in resources:

608 try:

609 client.delete_model(ModelName=m['ModelName'])

610 except client.exceptions.ResourceNotFound:

611 pass

612

613

614@SagemakerJob.action_registry.register('stop')

615class SagemakerJobStop(BaseAction):

616 """Stops a SageMaker job

617

618 :example:

619

620 .. code-block:: yaml

621

622 policies:

623 - name: stop-ml-job

624 resource: sagemaker-job

625 filters:

626 - TrainingJobName: ml-job-10

627 actions:

628 - stop

629 """

630 schema = type_schema('stop')

631 permissions = ('sagemaker:StopTrainingJob',)

632

633 def process(self, jobs):

634 client = local_session(self.manager.session_factory).client('sagemaker')

635

636 for j in jobs:

637 try:

638 client.stop_training_job(TrainingJobName=j['TrainingJobName'])

639 except client.exceptions.ResourceNotFound:

640 pass

641

642

643@SagemakerEndpoint.action_registry.register('delete')

644class SagemakerEndpointDelete(BaseAction):

645 """Delete a SageMaker endpoint

646

647 :example:

648

649 .. code-block:: yaml

650

651 policies:

652 - name: delete-sagemaker-endpoint

653 resource: sagemaker-endpoint

654 filters:

655 - EndpointName: sagemaker-ep--2018-01-01-00-00-00

656 actions:

657 - type: delete

658 """

659 permissions = (

660 'sagemaker:DeleteEndpoint',

661 'sagemaker:DeleteEndpointConfig')

662 schema = type_schema('delete')

663

664 def process(self, endpoints):

665 client = local_session(self.manager.session_factory).client('sagemaker')

666 for e in endpoints:

667 try:

668 client.delete_endpoint(EndpointName=e['EndpointName'])

669 except client.exceptions.ResourceNotFound:

670 pass

671

672

673@SagemakerEndpointConfig.action_registry.register('delete')

674class SagemakerEndpointConfigDelete(BaseAction):

675 """Delete a SageMaker endpoint

676

677 :example:

678

679 .. code-block:: yaml

680

681 policies:

682 - name: delete-sagemaker-endpoint-config

683 resource: sagemaker-endpoint-config

684 filters:

685 - EndpointConfigName: sagemaker-2018-01-01-00-00-00-T00

686 actions:

687 - delete

688 """

689 schema = type_schema('delete')

690 permissions = ('sagemaker:DeleteEndpointConfig',)

691

692 def process(self, endpoints):

693 client = local_session(self.manager.session_factory).client('sagemaker')

694 for e in endpoints:

695 try:

696 client.delete_endpoint_config(

697 EndpointConfigName=e['EndpointConfigName'])

698 except client.exceptions.ResourceNotFound:

699 pass

700

701

702@SagemakerTransformJob.action_registry.register('stop')

703class SagemakerTransformJobStop(BaseAction):

704 """Stops a SageMaker Transform job

705

706 :example:

707

708 .. code-block:: yaml

709

710 policies:

711 - name: stop-tranform-job

712 resource: sagemaker-transform-job

713 filters:

714 - TransformJobName: ml-job-10

715 actions:

716 - stop

717 """

718 schema = type_schema('stop')

719 permissions = ('sagemaker:StopTransformJob',)

720

721 def process(self, jobs):

722 client = local_session(self.manager.session_factory).client('sagemaker')

723

724 for j in jobs:

725 try:

726 client.stop_transform_job(TransformJobName=j['TransformJobName'])

727 except client.exceptions.ResourceNotFound:

728 pass