Skip to content

Conversation

Copy link

Copilot AI commented Jan 28, 2026

Why are the changes needed?

Two categories of unit tests were failing:

  1. test_source_code_renderer - lazy module attribute access doesn't support nested paths
  2. PyTorch tests - PyTorch 2.6 changed torch.load default from weights_only=False to True

What changes were proposed in this pull request?

1. SourceCodeRenderer pygments import

File: plugins/flytekit-deck-standard/flytekitplugins/deck/renderer.py

Replace lazy module attribute chain with direct imports:

# Before
formatter = pygments.formatters.html.HtmlFormatter(style="colorful")

# After  
from pygments import highlight
from pygments.formatters.html import HtmlFormatter
from pygments.lexers.python import PythonLexer

formatter = HtmlFormatter(style="colorful")

Aligns with pattern already used in flytekit/deck/renderer.py.

2. PyTorch backward compatibility

Files: flytekit/extras/pytorch/native.py, flytekit/extras/pytorch/checkpoint.py

Added explicit weights_only=False to torch.load calls:

return torch.load(local_path, map_location=map_location, weights_only=False)

Maintains compatibility with models/checkpoints saved before PyTorch 2.6. Includes security warnings about arbitrary code execution risk when loading untrusted pickle files.

How was this patch tested?

  • plugins/flytekit-deck-standard/tests/test_renderer.py: All 10 tests pass
  • tests/flytekit/unit/extras/pytorch/: All 13 tests pass
  • CodeQL security scan: 0 alerts

Setup process

N/A

Screenshots

N/A

Check all the applicable boxes

  • I updated the documentation accordingly.
  • All new and existing tests passed.
  • All commits are signed-off.

Related PRs

N/A

Docs link

N/A


✨ Let Copilot coding agent set things up for you — coding agent works faster and does higher quality work when set up for your repo.

The lazy module access pattern pygments.formatters.html.HtmlFormatter doesn't work
properly with nested attributes in lazy modules. Changed to use direct imports from
pygments submodules instead, matching the pattern used in flytekit/deck/renderer.py.

This fixes the AttributeError: html that was occurring in the test_source_code_renderer test.

Co-authored-by: pingsutw <37936015+pingsutw@users.noreply.github.com>
Copilot AI changed the title [WIP] Fix all unit test errors in the codebase Fix SourceCodeRenderer lazy module attribute access for pygments Jan 28, 2026
Copilot AI requested a review from pingsutw January 28, 2026 23:22
Copilot AI and others added 2 commits January 28, 2026 23:34
…lity

PyTorch 2.6 changed the default value of weights_only from False to True in torch.load.
This breaks backward compatibility when loading models and checkpoints saved before PyTorch 2.6.

Added weights_only=False to both torch.load calls in:
- flytekit/extras/pytorch/native.py (for tensors and modules)
- flytekit/extras/pytorch/checkpoint.py (for checkpoints)

This fixes the UnpicklingError that occurs when trying to load nn.Module objects with the new default.

Co-authored-by: pingsutw <37936015+pingsutw@users.noreply.github.com>
Enhanced comments to explicitly warn about the security implications of using
weights_only=False. Added warnings that this can execute arbitrary code from
pickle files and should only be used with trusted sources. Included reference
to PyTorch documentation.

Co-authored-by: pingsutw <37936015+pingsutw@users.noreply.github.com>
Copilot AI changed the title Fix SourceCodeRenderer lazy module attribute access for pygments Fix unit test failures: pygments lazy module access and PyTorch 2.6 compatibility Jan 28, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants